From 7e8777d574c35fb8860844021c155b6ca5b7a1ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20M=C3=BCller?= <bastian@turbolent.com>
Date: Mon, 26 Feb 2024 16:53:54 -0800
Subject: [PATCH] make staged contract code update migration aware of updated
 system contracts

---
 cmd/util/ledger/migrations/cadence.go         | 25 ++++++++++++++-----
 .../change_contract_code_migration.go         | 11 --------
 .../migrations/staged_contracts_migration.go  |  8 ++++--
 3 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/cmd/util/ledger/migrations/cadence.go b/cmd/util/ledger/migrations/cadence.go
index 1b6418e41eb..9e3b90db5d5 100644
--- a/cmd/util/ledger/migrations/cadence.go
+++ b/cmd/util/ledger/migrations/cadence.go
@@ -209,9 +209,27 @@ func NewCadence1ContractsMigrations(
 	stagedContracts []StagedContract,
 ) []ledger.Migration {
 
+	systemContractsMigration := NewChangeContractCodeMigration(chainID)
 	stagedContractsMigration := NewStagedContractsMigration(chainID).
 		WithContractUpdateValidation()
 
+	systemContractChanges := SystemContractChanges(
+		chainID,
+		SystemContractChangesOptions{
+			EVM: evmContractChange,
+		},
+	)
+
+	for _, change := range systemContractChanges {
+		systemContractsMigration.RegisterContractChange(change)
+
+		// Make sure the staged contracts migration is aware of the new system contract code
+		stagedContractsMigration.RegisterContract(
+			change.Address,
+			change.Contract,
+		)
+	}
+
 	stagedContractsMigration.RegisterContractUpdates(stagedContracts)
 
 	return []ledger.Migration{
@@ -219,12 +237,7 @@ func NewCadence1ContractsMigrations(
 			log,
 			nWorker,
 			[]AccountBasedMigration{
-				NewSystemContactsMigration(
-					chainID,
-					SystemContractChangesOptions{
-						EVM: evmContractChange,
-					},
-				),
+				systemContractsMigration,
 			},
 		),
 		NewBurnerDeploymentMigration(chainID, log),
diff --git a/cmd/util/ledger/migrations/change_contract_code_migration.go b/cmd/util/ledger/migrations/change_contract_code_migration.go
index eadfa65113d..f977aa340c6 100644
--- a/cmd/util/ledger/migrations/change_contract_code_migration.go
+++ b/cmd/util/ledger/migrations/change_contract_code_migration.go
@@ -318,14 +318,3 @@ func mustHexAddress(hexAddress string) common.Address {
 	}
 	return address
 }
-
-func NewSystemContactsMigration(
-	chainID flow.ChainID,
-	options SystemContractChangesOptions,
-) *ChangeContractCodeMigration {
-	migration := NewChangeContractCodeMigration(chainID)
-	for _, change := range SystemContractChanges(chainID, options) {
-		migration.RegisterContractChange(change)
-	}
-	return migration
-}
diff --git a/cmd/util/ledger/migrations/staged_contracts_migration.go b/cmd/util/ledger/migrations/staged_contracts_migration.go
index 3c3e1734546..f81f1e3cb18 100644
--- a/cmd/util/ledger/migrations/staged_contracts_migration.go
+++ b/cmd/util/ledger/migrations/staged_contracts_migration.go
@@ -142,11 +142,15 @@ func (m *StagedContractsMigration) RegisterContractChange(change StagedContract)
 
 	m.stagedContracts[address][registerID] = change.Contract
 
+	m.RegisterContract(address, change.Contract)
+}
+
+func (m *StagedContractsMigration) RegisterContract(address common.Address, contract Contract) {
 	location := common.AddressLocation{
-		Name:    name,
+		Name:    contract.Name,
 		Address: address,
 	}
-	m.contractsByLocation[location] = change.Code
+	m.contractsByLocation[location] = contract.Code
 }
 
 func (m *StagedContractsMigration) contractUpdatesForAccount(