From e73a0a115dfb0850b80ce62c018064482e66bc5e Mon Sep 17 00:00:00 2001
From: Giovanni Sanchez <108043524+sisyphusSmiling@users.noreply.github.com>
Date: Mon, 15 Apr 2024 16:59:33 -0500
Subject: [PATCH] fix updated EVM dependency aliasing

---
 .../change_contract_code_migration.go         |  2 ++
 fvm/bootstrap.go                              |  6 +++---
 fvm/evm/stdlib/contract.go                    | 19 +++++++++++++++----
 fvm/evm/stdlib/contract_test.go               |  2 +-
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/cmd/util/ledger/migrations/change_contract_code_migration.go b/cmd/util/ledger/migrations/change_contract_code_migration.go
index ff4271aa1ff..c327a677e1f 100644
--- a/cmd/util/ledger/migrations/change_contract_code_migration.go
+++ b/cmd/util/ledger/migrations/change_contract_code_migration.go
@@ -223,6 +223,8 @@ func SystemContractChanges(chainID flow.ChainID, options SystemContractsMigratio
 			NewSystemContractChange(
 				systemContracts.EVMContract,
 				evm.ContractCode(
+					flow.HexToAddress(env.NonFungibleTokenAddress),
+					flow.HexToAddress(env.FungibleTokenAddress),
 					flow.HexToAddress(env.FlowTokenAddress),
 				),
 			),
diff --git a/fvm/bootstrap.go b/fvm/bootstrap.go
index 85beee49d50..5c7e09d31ed 100644
--- a/fvm/bootstrap.go
+++ b/fvm/bootstrap.go
@@ -434,7 +434,7 @@ func (b *bootstrapExecutor) Execute() error {
 	b.deployStakingCollection(service, &env)
 
 	// sets up the EVM environment
-	b.setupEVM(service, fungibleToken, flowToken)
+	b.setupEVM(service, nonFungibleToken, fungibleToken, flowToken)
 
 	err = expectAccounts(systemcontracts.EVMStorageAccountIndex)
 	if err != nil {
@@ -895,7 +895,7 @@ func (b *bootstrapExecutor) setStakingAllowlist(
 	panicOnMetaInvokeErrf("failed to set staking allow-list: %s", txError, err)
 }
 
-func (b *bootstrapExecutor) setupEVM(serviceAddress, fungibleTokenAddress, flowTokenAddress flow.Address) {
+func (b *bootstrapExecutor) setupEVM(serviceAddress, nonFungibleTokenAddress, fungibleTokenAddress, flowTokenAddress flow.Address) {
 	if b.setupEVMEnabled {
 		// account for storage
 		// we dont need to deploy anything to this account, but it needs to exist
@@ -906,7 +906,7 @@ func (b *bootstrapExecutor) setupEVM(serviceAddress, fungibleTokenAddress, flowT
 		// deploy the EVM contract to the service account
 		tx := blueprints.DeployContractTransaction(
 			serviceAddress,
-			stdlib.ContractCode(flowTokenAddress),
+			stdlib.ContractCode(nonFungibleTokenAddress, fungibleTokenAddress, flowTokenAddress),
 			stdlib.ContractName,
 		)
 		// WithEVMEnabled should only be used after we create an account for storage
diff --git a/fvm/evm/stdlib/contract.go b/fvm/evm/stdlib/contract.go
index 532a327956e..8a06af45871 100644
--- a/fvm/evm/stdlib/contract.go
+++ b/fvm/evm/stdlib/contract.go
@@ -27,13 +27,24 @@ import (
 //go:embed contract.cdc
 var contractCode string
 
-var flowTokenImportPattern = regexp.MustCompile(`(?m)^import "FlowToken"\n`)
+var nftImportPattern = regexp.MustCompile(`(?m)^import "NonFungibleToken"`)
+var fungibleTokenImportPattern = regexp.MustCompile(`(?m)^import "FungibleToken"`)
+var flowTokenImportPattern = regexp.MustCompile(`(?m)^import "FlowToken"`)
 
-func ContractCode(flowTokenAddress flow.Address) []byte {
-	return []byte(flowTokenImportPattern.ReplaceAllString(
+func ContractCode(nonFungibleTokenAddress, fungibleTokenAddress, flowTokenAddress flow.Address) []byte {
+	contractCode = nftImportPattern.ReplaceAllString(
+		contractCode,
+		fmt.Sprintf("import NonFungibleToken from %s", nonFungibleTokenAddress.HexWithPrefix()),
+	)
+	contractCode = fungibleTokenImportPattern.ReplaceAllString(
+		contractCode,
+		fmt.Sprintf("import FungibleToken from %s", fungibleTokenAddress.HexWithPrefix()),
+	)
+	contractCode = flowTokenImportPattern.ReplaceAllString(
 		contractCode,
 		fmt.Sprintf("import FlowToken from %s", flowTokenAddress.HexWithPrefix()),
-	))
+	)
+	return []byte(contractCode)
 }
 
 const ContractName = "EVM"
diff --git a/fvm/evm/stdlib/contract_test.go b/fvm/evm/stdlib/contract_test.go
index 5e3c542288e..2006eeb1d75 100644
--- a/fvm/evm/stdlib/contract_test.go
+++ b/fvm/evm/stdlib/contract_test.go
@@ -238,7 +238,7 @@ func deployContracts(
 		},
 		{
 			name: stdlib.ContractName,
-			code: stdlib.ContractCode(contractsAddress),
+			code: stdlib.ContractCode(contractsAddress, contractsAddress, contractsAddress),
 		},
 	}