Merge pull request #5198 from ethereum-optimism/feat/migration-parallel-check

op-chain-ops: Check and mutate DB in parallel during OVM_ETH migration

op-chain-ops/ether/migrate.go

@@ -3,6 +3,10 @@ package ether
 import (
 	"fmt"
 	"math/big"
+	"sync"
+
+	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/trie"
 
 	"github.com/ethereum-optimism/optimism/op-chain-ops/crossdomain"
 	"github.com/ethereum-optimism/optimism/op-chain-ops/util"
@@ -13,6 +17,20 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 )
 
+const (
+	// checkJobs is the number of parallel workers to spawn
+	// when iterating the storage trie.
+	checkJobs = 64
+
+	// BalanceSlot is an ordinal used to represent slots corresponding to OVM_ETH
+	// balances in the state.
+	BalanceSlot = 1
+
+	// AllowanceSlot is an ordinal used to represent slots corresponding to OVM_ETH
+	// allowances in the state.
+	AllowanceSlot = 2
+)
+
 var (
 	// OVMETHAddress is the address of the OVM ETH predeploy.
 	OVMETHAddress = common.HexToAddress("0xDeadDeAddeAddEAddeadDEaDDEAdDeaDDeAD0000")
@@ -27,73 +45,348 @@ var (
 		common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000005"): true,
 		common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000006"): true,
 	}
+
+	// maxSlot is the maximum possible storage slot.
+	maxSlot = common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
+
+	// sequencerEntrypointAddr is the address of the OVM sequencer entrypoint contract.
+	sequencerEntrypointAddr = common.HexToAddress("0x4200000000000000000000000000000000000005")
 )
 
-type FilteredOVMETHAddresses []common.Address
+// accountData is a wrapper struct that contains the balance and address of an account.
+// It gets passed via channel to the collector process.
+type accountData struct {
+	balance    *big.Int
+	legacySlot common.Hash
+	address    common.Address
+}
+
+type DBFactory func() (*state.StateDB, error)
 
-func MigrateLegacyETH(db *state.StateDB, addresses FilteredOVMETHAddresses, chainID int, noCheck bool) error {
+// MigrateBalances migrates all balances in the LegacyERC20ETH contract into state. It performs checks
+// in parallel with mutations in order to reduce overall migration time.
+func MigrateBalances(mutableDB *state.StateDB, dbFactory DBFactory, addresses []common.Address, allowances []*crossdomain.Allowance, chainID int, noCheck bool) error {
 	// Chain params to use for integrity checking.
 	params := crossdomain.ParamsByChainID[chainID]
 	if params == nil {
 		return fmt.Errorf("no chain params for %d", chainID)
 	}
 
-	// Log the chain params for debugging purposes.
-	log.Info("Chain params", "chain-id", chainID, "supply-delta", params.ExpectedSupplyDelta)
+	return doMigration(mutableDB, dbFactory, addresses, allowances, params.ExpectedSupplyDelta, noCheck)
+}
+
+func doMigration(mutableDB *state.StateDB, dbFactory DBFactory, addresses []common.Address, allowances []*crossdomain.Allowance, expDiff *big.Int, noCheck bool) error {
+	// We'll need to maintain a list of all addresses that we've seen along with all of the storage
+	// slots based on the witness data.
+	slotsAddrs := make(map[common.Hash]common.Address)
+	slotsInp := make(map[common.Hash]int)
 
-	// Migrate the legacy ETH to ETH.
-	log.Info("Migrating legacy ETH to ETH", "num-accounts", len(addresses))
-	totalMigrated := new(big.Int)
-	logAccountProgress := util.ProgressLogger(1000, "imported accounts")
+	// For each known address, compute its balance key and add it to the list of addresses.
+	// Mint events are instrumented as regular ETH events in the witness data, so we no longer
+	// need to iterate over mint events during the migration.
 	for _, addr := range addresses {
-		// Balances are pre-checked not have any balances in state.
+		sk := CalcOVMETHStorageKey(addr)
+		slotsAddrs[sk] = addr
+		slotsInp[sk] = BalanceSlot
+	}
+
+	// For each known allowance, compute its storage key and add it to the list of addresses.
+	for _, allowance := range allowances {
+		sk := CalcAllowanceStorageKey(allowance.From, allowance.To)
+		slotsAddrs[sk] = allowance.From
+		slotsInp[sk] = AllowanceSlot
+	}
 
-		// Pull out the OVM ETH balance.
-		ovmBalance := GetOVMETHBalance(db, addr)
+	// Add the old SequencerEntrypoint because someone sent it ETH a long time ago and it has a
+	// balance but none of our instrumentation could easily find it. Special case.
+	entrySK := CalcOVMETHStorageKey(sequencerEntrypointAddr)
+	slotsAddrs[entrySK] = sequencerEntrypointAddr
+	slotsInp[entrySK] = BalanceSlot
 
-		// Actually perform the migration by setting the appropriate values in state.
-		db.SetBalance(addr, ovmBalance)
-		db.SetState(predeploys.LegacyERC20ETHAddr, CalcOVMETHStorageKey(addr), common.Hash{})
+	// WaitGroup to wait on each iteration job to finish.
+	var wg sync.WaitGroup
+	// Channel to receive storage slot keys and values from each iteration job.
+	outCh := make(chan accountData)
+	// Channel to receive errors from each iteration job.
+	errCh := make(chan error, checkJobs)
+	// Channel to cancel all iteration jobs as well as the collector.
+	cancelCh := make(chan struct{})
 
-		// Bump the total OVM balance.
-		totalMigrated = totalMigrated.Add(totalMigrated, ovmBalance)
+	// Define a worker function to iterate over each partition.
+	worker := func(start, end common.Hash) {
+		// Decrement the WaitGroup when the function returns.
+		defer wg.Done()
 
-		// Log progress.
-		logAccountProgress()
+		db, err := dbFactory()
+		if err != nil {
+			log.Crit("cannot get database", "err", err)
 		}
 
-	// Make sure that the total supply delta matches the expected delta. This is equivalent to
-	// checking that the total migrated is equal to the total found, since we already performed the
-	// same check against the total found (a = b, b = c => a = c).
-	totalSupply := getOVMETHTotalSupply(db)
-	delta := new(big.Int).Sub(totalSupply, totalMigrated)
-	if delta.Cmp(params.ExpectedSupplyDelta) != 0 {
+		// Create a new storage trie. Each trie returned by db.StorageTrie
+		// is a copy, so this is safe for concurrent use.
+		st, err := db.StorageTrie(predeploys.LegacyERC20ETHAddr)
+		if err != nil {
+			// Should never happen, so explode if it does.
+			log.Crit("cannot get storage trie for LegacyERC20ETHAddr", "err", err)
+		}
+		if st == nil {
+			// Should never happen, so explode if it does.
+			log.Crit("nil storage trie for LegacyERC20ETHAddr")
+		}
+
+		it := trie.NewIterator(st.NodeIterator(start.Bytes()))
+
+		// Below code is largely based on db.ForEachStorage. We can't use that
+		// because it doesn't allow us to specify a start and end key.
+		for it.Next() {
+			select {
+			case <-cancelCh:
+				// If one of the workers encounters an error, cancel all of them.
+				return
+			default:
+				break
+			}
+
+			// Use the raw (i.e., secure hashed) key to check if we've reached
+			// the end of the partition. Use > rather than >= here to account for
+			// the fact that the values returned by PartitionKeys are inclusive.
+			// Duplicate addresses that may be returned by this iteration are
+			// filtered out in the collector.
+			if new(big.Int).SetBytes(it.Key).Cmp(end.Big()) > 0 {
+				return
+			}
+
+			// Skip if the value is empty.
+			rawValue := it.Value
+			if len(rawValue) == 0 {
+				continue
+			}
+
+			// Get the preimage.
+			rawKey := st.GetKey(it.Key)
+			if rawKey == nil {
+				// Should never happen, so explode if it does.
+				log.Crit("cannot get preimage for storage key", "key", it.Key)
+			}
+			key := common.BytesToHash(rawKey)
+
+			// Parse the raw value.
+			_, content, _, err := rlp.Split(rawValue)
+			if err != nil {
+				// Should never happen, so explode if it does.
+				log.Crit("mal-formed data in state: %v", err)
+			}
+
+			// We can safely ignore specific slots (totalSupply, name, symbol).
+			if ignoredSlots[key] {
+				continue
+			}
+
+			slotType, ok := slotsInp[key]
+			if !ok {
 				if noCheck {
+					log.Error("ignoring unknown storage slot in state", "slot", key.String())
+				} else {
+					errCh <- fmt.Errorf("unknown storage slot in state: %s", key.String())
+					return
+				}
+			}
+
+			// No accounts should have a balance in state. If they do, bail.
+			addr, ok := slotsAddrs[key]
+			if !ok {
+				log.Crit("could not find address in map - should never happen")
+			}
+			bal := db.GetBalance(addr)
+			if bal.Sign() != 0 {
 				log.Error(
-				"supply mismatch",
-				"migrated", totalMigrated.String(),
-				"supply", totalSupply.String(),
-				"delta", delta.String(),
-				"exp_delta", params.ExpectedSupplyDelta.String(),
+					"account has non-zero balance in state - should never happen",
+					"addr", addr,
+					"balance", bal.String(),
 				)
+				if !noCheck {
+					errCh <- fmt.Errorf("account has non-zero balance in state - should never happen: %s", addr.String())
+					return
+				}
+			}
+
+			// Add balances to the total found.
+			switch slotType {
+			case BalanceSlot:
+				// Convert the value to a common.Hash, then send to the channel.
+				value := common.BytesToHash(content)
+				outCh <- accountData{
+					balance:    value.Big(),
+					legacySlot: key,
+					address:    addr,
+				}
+			case AllowanceSlot:
+				// Allowance slot.
+				continue
+			default:
+				// Should never happen.
+				if noCheck {
+					log.Error("unknown slot type", "slot", key, "type", slotType)
 				} else {
-			log.Crit(
+					log.Crit("unknown slot type %d, should never happen", slotType)
+				}
+			}
+		}
+	}
+
+	for i := 0; i < checkJobs; i++ {
+		wg.Add(1)
+
+		// Partition the keyspace per worker.
+		start, end := PartitionKeyspace(i, checkJobs)
+
+		// Kick off our worker.
+		go worker(start, end)
+	}
+
+	// Make a channel to make sure that the collector process completes.
+	collectorCloseCh := make(chan struct{})
+
+	// Keep track of the last error seen.
+	var lastErr error
+
+	// There are multiple ways that the cancel channel can be closed:
+	// - if we receive an error from the errCh
+	// - if the collector process completes
+	// To prevent panics, we wrap the close in a sync.Once.
+	var cancelOnce sync.Once
+
+	// Create a map of accounts we've seen so that we can filter out duplicates.
+	seenAccounts := make(map[common.Address]bool)
+
+	// Keep track of the total migrated supply.
+	totalFound := new(big.Int)
+
+	// Kick off another background process to collect
+	// values from the channel and add them to the map.
+	var count int
+	progress := util.ProgressLogger(1000, "Migrated OVM_ETH storage slot")
+	go func() {
+		defer func() {
+			collectorCloseCh <- struct{}{}
+		}()
+		for {
+			select {
+			case account := <-outCh:
+				progress()
+
+				// Filter out duplicate accounts. See the below note about keyspace iteration for
+				// why we may have to filter out duplicates.
+				if seenAccounts[account.address] {
+					log.Info("skipping duplicate account during iteration", "addr", account.address)
+					continue
+				}
+
+				// Accumulate addresses and total supply.
+				totalFound = new(big.Int).Add(totalFound, account.balance)
+
+				mutableDB.SetBalance(account.address, account.balance)
+				mutableDB.SetState(predeploys.LegacyERC20ETHAddr, account.legacySlot, common.Hash{})
+				count++
+				seenAccounts[account.address] = true
+			case err := <-errCh:
+				cancelOnce.Do(func() {
+					lastErr = err
+					close(cancelCh)
+				})
+			case <-cancelCh:
+				return
+			}
+		}
+	}()
+
+	// Wait for the workers to finish.
+	wg.Wait()
+	// Close the cancel channel to signal the collector process to stop.
+	cancelOnce.Do(func() {
+		close(cancelCh)
+	})
+
+	// Wait for the collector process to finish.
+	<-collectorCloseCh
+
+	// If we saw an error, return it.
+	if lastErr != nil {
+		return lastErr
+	}
+
+	// Log how many slots were iterated over.
+	log.Info("Iterated legacy balances", "count", count)
+
+	// Verify the supply delta. Recorded total supply in the LegacyERC20ETH contract may be higher
+	// than the actual migrated amount because self-destructs will remove ETH supply in a way that
+	// cannot be reflected in the contract. This is fine because self-destructs just mean the L2 is
+	// actually *overcollateralized* by some tiny amount.
+	db, err := dbFactory()
+	if err != nil {
+		log.Crit("cannot get database", "err", err)
+	}
+
+	totalSupply := getOVMETHTotalSupply(db)
+	delta := new(big.Int).Sub(totalSupply, totalFound)
+	if delta.Cmp(expDiff) != 0 {
+		log.Error(
 			"supply mismatch",
-				"migrated", totalMigrated.String(),
+			"migrated", totalFound.String(),
 			"supply", totalSupply.String(),
 			"delta", delta.String(),
-				"exp_delta", params.ExpectedSupplyDelta.String(),
+			"exp_delta", expDiff.String(),
 		)
+		if !noCheck {
+			return fmt.Errorf("supply mismatch: %s", delta.String())
 		}
 	}
 
+	// Supply is verified.
+	log.Info(
+		"supply verified OK",
+		"migrated", totalFound.String(),
+		"supply", totalSupply.String(),
+		"delta", delta.String(),
+		"exp_delta", expDiff.String(),
+	)
+
 	// Set the total supply to 0. We do this because the total supply is necessarily going to be
 	// different than the sum of all balances since we no longer track balances inside the contract
 	// itself. The total supply is going to be weird no matter what, might as well set it to zero
 	// so it's explicitly weird instead of implicitly weird.
-	db.SetState(predeploys.LegacyERC20ETHAddr, getOVMETHTotalSupplySlot(), common.Hash{})
+	mutableDB.SetState(predeploys.LegacyERC20ETHAddr, getOVMETHTotalSupplySlot(), common.Hash{})
 	log.Info("Set the totalSupply to 0")
 
-	// Fin.
 	return nil
 }
+
+// PartitionKeyspace divides the key space into partitions by dividing the maximum keyspace
+// by count then multiplying by i. This will leave some slots left over, which we handle below. It
+// returns the start and end keys for the partition as a common.Hash. Note that the returned range
+// of keys is inclusive, i.e., [start, end] NOT [start, end).
+func PartitionKeyspace(i int, count int) (common.Hash, common.Hash) {
+	if i < 0 || count < 0 {
+		panic("i and count must be greater than 0")
+	}
+
+	if i > count-1 {
+		panic("i must be less than count - 1")
+	}
+
+	// Divide the key space into partitions by dividing the key space by the number
+	// of jobs. This will leave some slots left over, which we handle below.
+	partSize := new(big.Int).Div(maxSlot.Big(), big.NewInt(int64(count)))
+
+	start := common.BigToHash(new(big.Int).Mul(big.NewInt(int64(i)), partSize))
+	var end common.Hash
+	if i < count-1 {
+		// If this is not the last partition, use the next partition's start key as the end.
+		end = common.BigToHash(new(big.Int).Mul(big.NewInt(int64(i+1)), partSize))
+	} else {
+		// If this is the last partition, use the max slot as the end.
+		end = maxSlot
+	}
+
+	return start, end
+}

op-chain-ops/ether/precheck_test.go → op-chain-ops/ether/migrate_test.go

 package ether
 
 import (
-	"bytes"
+	"fmt"
 	"math/big"
 	"math/rand"
-	"os"
-	"sort"
 	"testing"
 
-	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum-optimism/optimism/op-bindings/predeploys"
 
 	"github.com/ethereum-optimism/optimism/op-chain-ops/crossdomain"
 	"github.com/ethereum/go-ethereum/common"
@@ -18,9 +16,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestPreCheckBalances(t *testing.T) {
-	log.Root().SetHandler(log.StreamHandler(os.Stderr, log.TerminalFormat(true)))
-
+func TestMigrateBalances(t *testing.T) {
 	tests := []struct {
 		name            string
 		totalSupply     *big.Int
@@ -29,7 +25,7 @@ func TestPreCheckBalances(t *testing.T) {
 		stateAllowances map[common.Address]common.Address
 		inputAddresses  []common.Address
 		inputAllowances []*crossdomain.Allowance
-		check           func(t *testing.T, addrs FilteredOVMETHAddresses, err error)
+		check           func(t *testing.T, db *state.StateDB, err error)
 	}{
 		{
 			name:        "everything matches",
@@ -52,12 +48,11 @@ func TestPreCheckBalances(t *testing.T) {
 					To:   common.HexToAddress("0x456"),
 				},
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.NoError(t, err)
-				require.EqualValues(t, FilteredOVMETHAddresses{
-					common.HexToAddress("0x123"),
-					common.HexToAddress("0x456"),
-				}, addrs)
+				require.EqualValues(t, common.Big1, db.GetBalance(common.HexToAddress("0x123")))
+				require.EqualValues(t, common.Big2, db.GetBalance(common.HexToAddress("0x456")))
+				require.EqualValues(t, common.Hash{}, db.GetState(predeploys.LegacyERC20ETHAddr, GetOVMETHTotalSupplySlot()))
 			},
 		},
 		{
@@ -71,11 +66,11 @@ func TestPreCheckBalances(t *testing.T) {
 				common.HexToAddress("0x123"),
 				common.HexToAddress("0x456"),
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.NoError(t, err)
-				require.EqualValues(t, FilteredOVMETHAddresses{
-					common.HexToAddress("0x123"),
-				}, addrs)
+				require.EqualValues(t, common.Big1, db.GetBalance(common.HexToAddress("0x123")))
+				require.EqualValues(t, common.Big0, db.GetBalance(common.HexToAddress("0x456")))
+				require.EqualValues(t, common.Hash{}, db.GetState(predeploys.LegacyERC20ETHAddr, GetOVMETHTotalSupplySlot()))
 			},
 		},
 		{
@@ -102,11 +97,11 @@ func TestPreCheckBalances(t *testing.T) {
 					To:   common.HexToAddress("0x789"),
 				},
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.NoError(t, err)
-				require.EqualValues(t, FilteredOVMETHAddresses{
-					common.HexToAddress("0x123"),
-				}, addrs)
+				require.EqualValues(t, common.Big1, db.GetBalance(common.HexToAddress("0x123")))
+				require.EqualValues(t, common.Big0, db.GetBalance(common.HexToAddress("0x456")))
+				require.EqualValues(t, common.Hash{}, db.GetState(predeploys.LegacyERC20ETHAddr, GetOVMETHTotalSupplySlot()))
 			},
 		},
 		{
@@ -120,7 +115,7 @@ func TestPreCheckBalances(t *testing.T) {
 			inputAddresses: []common.Address{
 				common.HexToAddress("0x123"),
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.Error(t, err)
 				require.ErrorContains(t, err, "unknown storage slot")
 			},
@@ -145,7 +140,7 @@ func TestPreCheckBalances(t *testing.T) {
 					To:   common.HexToAddress("0x456"),
 				},
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.Error(t, err)
 				require.ErrorContains(t, err, "unknown storage slot")
 			},
@@ -162,7 +157,7 @@ func TestPreCheckBalances(t *testing.T) {
 				common.HexToAddress("0x123"),
 				common.HexToAddress("0x456"),
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.Error(t, err)
 				require.ErrorContains(t, err, "supply mismatch")
 			},
@@ -179,35 +174,25 @@ func TestPreCheckBalances(t *testing.T) {
 				common.HexToAddress("0x123"),
 				common.HexToAddress("0x456"),
 			},
-			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+			check: func(t *testing.T, db *state.StateDB, err error) {
 				require.NoError(t, err)
-				require.EqualValues(t, FilteredOVMETHAddresses{
-					common.HexToAddress("0x123"),
-					common.HexToAddress("0x456"),
-				}, addrs)
+				require.EqualValues(t, common.Big1, db.GetBalance(common.HexToAddress("0x123")))
+				require.EqualValues(t, common.Big2, db.GetBalance(common.HexToAddress("0x456")))
 			},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			db := makeLegacyETH(t, tt.totalSupply, tt.stateBalances, tt.stateAllowances)
-			factory := func() (*state.StateDB, error) {
-				return db, nil
-			}
-			addrs, err := doMigration(factory, tt.inputAddresses, tt.inputAllowances, tt.expDiff, false)
-
-			// Sort the addresses since they come in in a random order.
-			sort.Slice(addrs, func(i, j int) bool {
-				return bytes.Compare(addrs[i][:], addrs[j][:]) < 0
-			})
-
-			tt.check(t, addrs, err)
+			db, factory := makeLegacyETH(t, tt.totalSupply, tt.stateBalances, tt.stateAllowances)
+			err := doMigration(db, factory, tt.inputAddresses, tt.inputAllowances, tt.expDiff, false)
+			tt.check(t, db, err)
 		})
 	}
 }
 
-func makeLegacyETH(t *testing.T, totalSupply *big.Int, balances map[common.Address]*big.Int, allowances map[common.Address]common.Address) *state.StateDB {
-	db, err := state.New(common.Hash{}, state.NewDatabaseWithConfig(rawdb.NewMemoryDatabase(), &trie.Config{
+func makeLegacyETH(t *testing.T, totalSupply *big.Int, balances map[common.Address]*big.Int, allowances map[common.Address]common.Address) (*state.StateDB, DBFactory) {
+	memDB := rawdb.NewMemoryDatabase()
+	db, err := state.New(common.Hash{}, state.NewDatabaseWithConfig(memDB, &trie.Config{
 		Preimages: true,
 		Cache:     1024,
 	}), nil)
@@ -235,13 +220,19 @@ func makeLegacyETH(t *testing.T, totalSupply *big.Int, balances map[common.Addre
 	err = db.Database().TrieDB().Commit(root, true)
 	require.NoError(t, err)
 
-	return db
+	return db, func() (*state.StateDB, error) {
+		return state.New(root, state.NewDatabaseWithConfig(memDB, &trie.Config{
+			Preimages: true,
+			Cache:     1024,
+		}), nil)
+	}
 }
 
-// TestPreCheckBalancesRandom tests that the pre-check balances function works
+// TestMigrateBalancesRandom tests that the pre-check balances function works
 // with random addresses. This test makes sure that the partition logic doesn't
 // miss anything.
-func TestPreCheckBalancesRandom(t *testing.T) {
+func TestMigrateBalancesRandom(t *testing.T) {
+	for i := 0; i < 100; i++ {
 		addresses := make([]common.Address, 0)
 		stateBalances := make(map[common.Address]*big.Int)
 
@@ -250,19 +241,14 @@ func TestPreCheckBalancesRandom(t *testing.T) {
 
 		totalSupply := big.NewInt(0)
 
-	for i := 0; i < 100; i++ {
-		for i := 0; i < rand.Intn(1000); i++ {
+		for j := 0; j < rand.Intn(10000); j++ {
 			addr := randAddr(t)
 			addresses = append(addresses, addr)
 			stateBalances[addr] = big.NewInt(int64(rand.Intn(1_000_000)))
 			totalSupply = new(big.Int).Add(totalSupply, stateBalances[addr])
 		}
 
-		sort.Slice(addresses, func(i, j int) bool {
-			return bytes.Compare(addresses[i][:], addresses[j][:]) < 0
-		})
-
-		for i := 0; i < rand.Intn(1000); i++ {
+		for j := 0; j < rand.Intn(1000); j++ {
 			addr := randAddr(t)
 			to := randAddr(t)
 			allowances = append(allowances, &crossdomain.Allowance{
@@ -272,19 +258,94 @@ func TestPreCheckBalancesRandom(t *testing.T) {
 			stateAllowances[addr] = to
 		}
 
-		db := makeLegacyETH(t, totalSupply, stateBalances, stateAllowances)
-		factory := func() (*state.StateDB, error) {
-			return db, nil
-		}
-
-		outAddrs, err := doMigration(factory, addresses, allowances, big.NewInt(0), false)
+		db, factory := makeLegacyETH(t, totalSupply, stateBalances, stateAllowances)
+		err := doMigration(db, factory, addresses, allowances, big.NewInt(0), false)
 		require.NoError(t, err)
 
-		sort.Slice(outAddrs, func(i, j int) bool {
-			return bytes.Compare(outAddrs[i][:], outAddrs[j][:]) < 0
+		for addr, expBal := range stateBalances {
+			actBal := db.GetBalance(addr)
+			require.EqualValues(t, expBal, actBal)
+		}
+	}
+}
+
+func TestPartitionKeyspace(t *testing.T) {
+	tests := []struct {
+		i        int
+		count    int
+		expected [2]common.Hash
+	}{
+		{
+			i:     0,
+			count: 1,
+			expected: [2]common.Hash{
+				common.HexToHash("0x00"),
+				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+		{
+			i:     0,
+			count: 2,
+			expected: [2]common.Hash{
+				common.HexToHash("0x00"),
+				common.HexToHash("0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+		{
+			i:     1,
+			count: 2,
+			expected: [2]common.Hash{
+				common.HexToHash("0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+		{
+			i:     0,
+			count: 3,
+			expected: [2]common.Hash{
+				common.HexToHash("0x00"),
+				common.HexToHash("0x5555555555555555555555555555555555555555555555555555555555555555"),
+			},
+		},
+		{
+			i:     1,
+			count: 3,
+			expected: [2]common.Hash{
+				common.HexToHash("0x5555555555555555555555555555555555555555555555555555555555555555"),
+				common.HexToHash("0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
+			},
+		},
+		{
+			i:     2,
+			count: 3,
+			expected: [2]common.Hash{
+				common.HexToHash("0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
+				common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("i %d, count %d", tt.i, tt.count), func(t *testing.T) {
+			start, end := PartitionKeyspace(tt.i, tt.count)
+			require.Equal(t, tt.expected[0], start)
+			require.Equal(t, tt.expected[1], end)
 		})
-		require.EqualValues(t, addresses, outAddrs)
 	}
+
+	t.Run("panics on invalid i or count", func(t *testing.T) {
+		require.Panics(t, func() {
+			PartitionKeyspace(1, 1)
+		})
+		require.Panics(t, func() {
+			PartitionKeyspace(-1, 1)
+		})
+		require.Panics(t, func() {
+			PartitionKeyspace(0, -1)
+		})
+		require.Panics(t, func() {
+			PartitionKeyspace(-1, -1)
+		})
+	})
 }
 
 func randAddr(t *testing.T) common.Address {

op-chain-ops/ether/precheck.go

-package ether
-
-import (
-	"fmt"
-	"math/big"
-	"sync"
-
-	"github.com/ethereum/go-ethereum/rlp"
-	"github.com/ethereum/go-ethereum/trie"
-
-	"github.com/ethereum-optimism/optimism/op-chain-ops/crossdomain"
-	"github.com/ethereum-optimism/optimism/op-chain-ops/util"
-
-	"github.com/ethereum-optimism/optimism/op-bindings/predeploys"
-	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/core/state"
-	"github.com/ethereum/go-ethereum/log"
-)
-
-const (
-	// checkJobs is the number of parallel workers to spawn
-	// when iterating the storage trie.
-	checkJobs = 64
-)
-
-// maxSlot is the maximum possible storage slot.
-var maxSlot = common.HexToHash("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff")
-
-// accountData is a wrapper struct that contains the balance and address of an account.
-// It gets passed via channel to the collector process.
-type accountData struct {
-	balance *big.Int
-	address common.Address
-}
-
-type DBFactory func() (*state.StateDB, error)
-
-// PreCheckBalances checks that the given list of addresses and allowances represents all storage
-// slots in the LegacyERC20ETH contract. We don't have to filter out extra addresses like we do for
-// withdrawals because we'll simply carry the balance of a given address to the new system, if the
-// account is extra then it won't have any balance and nothing will happen.
-func PreCheckBalances(dbFactory DBFactory, addresses []common.Address, allowances []*crossdomain.Allowance, chainID int, noCheck bool) (FilteredOVMETHAddresses, error) {
-	// Chain params to use for integrity checking.
-	params := crossdomain.ParamsByChainID[chainID]
-	if params == nil {
-		return nil, fmt.Errorf("no chain params for %d", chainID)
-	}
-
-	return doMigration(dbFactory, addresses, allowances, params.ExpectedSupplyDelta, noCheck)
-}
-
-func doMigration(dbFactory DBFactory, addresses []common.Address, allowances []*crossdomain.Allowance, expDiff *big.Int, noCheck bool) (FilteredOVMETHAddresses, error) {
-	// We'll need to maintain a list of all addresses that we've seen along with all of the storage
-	// slots based on the witness data.
-	addrs := make([]common.Address, 0)
-	slotsAddrs := make(map[common.Hash]common.Address)
-	slotsInp := make(map[common.Hash]int)
-
-	// For each known address, compute its balance key and add it to the list of addresses.
-	// Mint events are instrumented as regular ETH events in the witness data, so we no longer
-	// need to iterate over mint events during the migration.
-	for _, addr := range addresses {
-		sk := CalcOVMETHStorageKey(addr)
-		slotsAddrs[sk] = addr
-		slotsInp[sk] = 1
-	}
-
-	// For each known allowance, compute its storage key and add it to the list of addresses.
-	for _, allowance := range allowances {
-		sk := CalcAllowanceStorageKey(allowance.From, allowance.To)
-		slotsAddrs[sk] = allowance.From
-		slotsInp[sk] = 2
-	}
-
-	// Add the old SequencerEntrypoint because someone sent it ETH a long time ago and it has a
-	// balance but none of our instrumentation could easily find it. Special case.
-	sequencerEntrypointAddr := common.HexToAddress("0x4200000000000000000000000000000000000005")
-	entrySK := CalcOVMETHStorageKey(sequencerEntrypointAddr)
-	slotsAddrs[entrySK] = sequencerEntrypointAddr
-	slotsInp[entrySK] = 1
-
-	// WaitGroup to wait on each iteration job to finish.
-	var wg sync.WaitGroup
-	// Channel to receive storage slot keys and values from each iteration job.
-	outCh := make(chan accountData)
-	// Channel to receive errors from each iteration job.
-	errCh := make(chan error, checkJobs)
-	// Channel to cancel all iteration jobs as well as the collector.
-	cancelCh := make(chan struct{})
-
-	// Keep track of the total migrated supply.
-	totalFound := new(big.Int)
-
-	// Divide the key space into partitions by dividing the key space by the number
-	// of jobs. This will leave some slots left over, which we handle below.
-	partSize := new(big.Int).Div(maxSlot.Big(), big.NewInt(checkJobs))
-
-	// Define a worker function to iterate over each partition.
-	worker := func(start, end common.Hash) {
-		// Decrement the WaitGroup when the function returns.
-		defer wg.Done()
-
-		db, err := dbFactory()
-		if err != nil {
-			log.Crit("cannot get database", "err", err)
-		}
-
-		// Create a new storage trie. Each trie returned by db.StorageTrie
-		// is a copy, so this is safe for concurrent use.
-		st, err := db.StorageTrie(predeploys.LegacyERC20ETHAddr)
-		if err != nil {
-			// Should never happen, so explode if it does.
-			log.Crit("cannot get storage trie for LegacyERC20ETHAddr", "err", err)
-		}
-		if st == nil {
-			// Should never happen, so explode if it does.
-			log.Crit("nil storage trie for LegacyERC20ETHAddr")
-		}
-
-		it := trie.NewIterator(st.NodeIterator(start.Bytes()))
-
-		// Below code is largely based on db.ForEachStorage. We can't use that
-		// because it doesn't allow us to specify a start and end key.
-		for it.Next() {
-			select {
-			case <-cancelCh:
-				// If one of the workers encounters an error, cancel all of them.
-				return
-			default:
-				break
-			}
-
-			// Use the raw (i.e., secure hashed) key to check if we've reached
-			// the end of the partition.
-			if new(big.Int).SetBytes(it.Key).Cmp(end.Big()) >= 0 {
-				return
-			}
-
-			// Skip if the value is empty.
-			rawValue := it.Value
-			if len(rawValue) == 0 {
-				continue
-			}
-
-			// Get the preimage.
-			key := common.BytesToHash(st.GetKey(it.Key))
-
-			// Parse the raw value.
-			_, content, _, err := rlp.Split(rawValue)
-			if err != nil {
-				// Should never happen, so explode if it does.
-				log.Crit("mal-formed data in state: %v", err)
-			}
-
-			// We can safely ignore specific slots (totalSupply, name, symbol).
-			if ignoredSlots[key] {
-				continue
-			}
-
-			slotType, ok := slotsInp[key]
-			if !ok {
-				if noCheck {
-					log.Error("ignoring unknown storage slot in state", "slot", key.String())
-				} else {
-					errCh <- fmt.Errorf("unknown storage slot in state: %s", key.String())
-					return
-				}
-			}
-
-			// No accounts should have a balance in state. If they do, bail.
-			addr, ok := slotsAddrs[key]
-			if !ok {
-				log.Crit("could not find address in map - should never happen")
-			}
-			bal := db.GetBalance(addr)
-			if bal.Sign() != 0 {
-				log.Error(
-					"account has non-zero balance in state - should never happen",
-					"addr", addr,
-					"balance", bal.String(),
-				)
-				if !noCheck {
-					errCh <- fmt.Errorf("account has non-zero balance in state - should never happen: %s", addr.String())
-					return
-				}
-			}
-
-			// Add balances to the total found.
-			switch slotType {
-			case 1:
-				// Convert the value to a common.Hash, then send to the channel.
-				value := common.BytesToHash(content)
-				outCh <- accountData{
-					balance: value.Big(),
-					address: addr,
-				}
-			case 2:
-				// Allowance slot.
-				continue
-			default:
-				// Should never happen.
-				if noCheck {
-					log.Error("unknown slot type", "slot", key, "type", slotType)
-				} else {
-					log.Crit("unknown slot type %d, should never happen", slotType)
-				}
-			}
-		}
-	}
-
-	for i := 0; i < checkJobs; i++ {
-		wg.Add(1)
-
-		// Compute the start and end keys for this partition.
-		start := common.BigToHash(new(big.Int).Mul(big.NewInt(int64(i)), partSize))
-		var end common.Hash
-		if i < checkJobs-1 {
-			// If this is not the last partition, use the next partition's start key as the end.
-			end = common.BigToHash(new(big.Int).Mul(big.NewInt(int64(i+1)), partSize))
-		} else {
-			// If this is the last partition, use the max slot as the end.
-			end = maxSlot
-		}
-
-		// Kick off our worker.
-		go worker(start, end)
-	}
-
-	// Make a channel to make sure that the collector process completes.
-	collectorCloseCh := make(chan struct{})
-
-	// Keep track of the last error seen.
-	var lastErr error
-
-	// There are multiple ways that the cancel channel can be closed:
-	// - if we receive an error from the errCh
-	// - if the collector process completes
-	// To prevent panics, we wrap the close in a sync.Once.
-	var cancelOnce sync.Once
-
-	// Kick off another background process to collect
-	// values from the channel and add them to the map.
-	var count int
-	progress := util.ProgressLogger(1000, "Collected OVM_ETH storage slot")
-	go func() {
-		defer func() {
-			collectorCloseCh <- struct{}{}
-		}()
-		for {
-			select {
-			case account := <-outCh:
-				progress()
-				// Accumulate addresses and total supply.
-				addrs = append(addrs, account.address)
-				totalFound = new(big.Int).Add(totalFound, account.balance)
-			case err := <-errCh:
-				lastErr = err
-				cancelOnce.Do(func() {
-					close(cancelCh)
-				})
-			case <-cancelCh:
-				return
-			}
-		}
-	}()
-
-	// Wait for the workers to finish.
-	wg.Wait()
-	// Close the cancel channel to signal the collector process to stop.
-	cancelOnce.Do(func() {
-		close(cancelCh)
-	})
-
-	// Wait for the collector process to finish.
-	<-collectorCloseCh
-
-	// If we saw an error, return it.
-	if lastErr != nil {
-		return nil, lastErr
-	}
-
-	// Log how many slots were iterated over.
-	log.Info("Iterated legacy balances", "count", count)
-
-	// Verify the supply delta. Recorded total supply in the LegacyERC20ETH contract may be higher
-	// than the actual migrated amount because self-destructs will remove ETH supply in a way that
-	// cannot be reflected in the contract. This is fine because self-destructs just mean the L2 is
-	// actually *overcollateralized* by some tiny amount.
-	db, err := dbFactory()
-	if err != nil {
-		log.Crit("cannot get database", "err", err)
-	}
-
-	totalSupply := getOVMETHTotalSupply(db)
-	delta := new(big.Int).Sub(totalSupply, totalFound)
-	if delta.Cmp(expDiff) != 0 {
-		log.Error(
-			"supply mismatch",
-			"migrated", totalFound.String(),
-			"supply", totalSupply.String(),
-			"delta", delta.String(),
-			"exp_delta", expDiff.String(),
-		)
-		if !noCheck {
-			return nil, fmt.Errorf("supply mismatch: %s", delta.String())
-		}
-	}
-
-	// Supply is verified.
-	log.Info(
-		"supply verified OK",
-		"migrated", totalFound.String(),
-		"supply", totalSupply.String(),
-		"delta", delta.String(),
-		"exp_delta", expDiff.String(),
-	)
-
-	// We know we have at least a superset of all addresses here since we know that we have every
-	// storage slot. It's fine to have extras because they won't have any balance.
-	return addrs, nil
-}

op-chain-ops/genesis/check.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"math/big"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/util"
+
 	"github.com/ethereum-optimism/optimism/op-chain-ops/ether"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -508,7 +510,9 @@ func CheckWithdrawalsAfter(db vm.StateDB, data crossdomain.MigrationData, l1Cros
 	// Now, iterate over each legacy withdrawal and check if there is a corresponding
 	// migrated withdrawal.
 	var innerErr error
+	progress := util.ProgressLogger(1000, "checking withdrawals")
 	err = db.ForEachStorage(predeploys.LegacyMessagePasserAddr, func(key, value common.Hash) bool {
+		progress()
 		// The legacy message passer becomes a proxy during the migration,
 		// so we need to ignore the implementation/admin slots.
 		if key == ImplementationSlot || key == AdminSlot {

op-chain-ops/genesis/db_migration.go

@@ -143,16 +143,6 @@ func MigrateDB(ldb ethdb.Database, config *DeployConfig, l1Block *types.Block, m
 		filteredWithdrawals = crossdomain.SafeFilteredWithdrawals(unfilteredWithdrawals)
 	}
 
-	// We also need to verify that we have all of the storage slots for the LegacyERC20ETH contract
-	// that we expect to have. An error will be thrown if there are any missing storage slots.
-	// Unlike with withdrawals, we do not need to filter out extra addresses because their balances
-	// would necessarily be zero and therefore not affect the migration.
-	log.Info("Checking addresses...", "no-check", noCheck)
-	addrs, err := ether.PreCheckBalances(dbFactory, migrationData.Addresses(), migrationData.OvmAllowances, int(config.L1ChainID), noCheck)
-	if err != nil {
-		return nil, fmt.Errorf("addresses mismatch: %w", err)
-	}
-
 	// At this point we've fully verified the witness data for the migration, so we can begin the
 	// actual migration process. This involves modifying parts of the legacy database and inserting
 	// a transition block.
@@ -202,11 +192,12 @@ func MigrateDB(ldb ethdb.Database, config *DeployConfig, l1Block *types.Block, m
 	}
 
 	// Finally we migrate the balances held inside the LegacyERC20ETH contract into the state trie.
-	// We also delete the balances from the LegacyERC20ETH contract.
+	// We also delete the balances from the LegacyERC20ETH contract. Unlike the steps above, this step
+	// combines the check and mutation steps into one in order to reduce migration time.
 	log.Info("Starting to migrate ERC20 ETH")
-	err = ether.MigrateLegacyETH(db, addrs, int(config.L1ChainID), noCheck)
+	err = ether.MigrateBalances(db, dbFactory, migrationData.Addresses(), migrationData.OvmAllowances, int(config.L1ChainID), noCheck)
 	if err != nil {
-		return nil, fmt.Errorf("cannot migrate legacy eth: %w", err)
+		return nil, fmt.Errorf("failed to migrate OVM_ETH: %w", err)
 	}
 
 	// We're done messing around with the database, so we can now commit the changes to the DB.