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

op-chain-ops: Add parallel migration

op-chain-ops/cmd/check-migration/main.go

+package main
+
+import (
+	"context"
+	"fmt"
+	"math/big"
+	"os"
+	"strings"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/crossdomain"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/db"
+	"github.com/mattn/go-isatty"
+
+	"github.com/ethereum-optimism/optimism/op-node/eth"
+	"github.com/ethereum-optimism/optimism/op-node/rollup/derive"
+	"github.com/ethereum/go-ethereum/common"
+
+	"github.com/ethereum/go-ethereum/core/types"
+
+	"github.com/ethereum/go-ethereum/log"
+
+	"github.com/ethereum-optimism/optimism/op-bindings/hardhat"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/genesis"
+	"github.com/ethereum/go-ethereum/ethclient"
+
+	"github.com/urfave/cli"
+)
+
+func main() {
+	log.Root().SetHandler(log.StreamHandler(os.Stderr, log.TerminalFormat(isatty.IsTerminal(os.Stderr.Fd()))))
+
+	app := &cli.App{
+		Name:  "check-migration",
+		Usage: "Run sanity checks on a migrated database",
+		Flags: []cli.Flag{
+			&cli.StringFlag{
+				Name:     "l1-rpc-url",
+				Value:    "http://127.0.0.1:8545",
+				Usage:    "RPC URL for an L1 Node",
+				Required: true,
+			},
+			&cli.StringFlag{
+				Name:     "ovm-addresses",
+				Usage:    "Path to ovm-addresses.json",
+				Required: true,
+			},
+			&cli.StringFlag{
+				Name:     "ovm-allowances",
+				Usage:    "Path to ovm-allowances.json",
+				Required: true,
+			},
+			&cli.StringFlag{
+				Name:     "ovm-messages",
+				Usage:    "Path to ovm-messages.json",
+				Required: true,
+			},
+			&cli.StringFlag{
+				Name:     "witness-file",
+				Usage:    "Path to witness file",
+				Required: true,
+			},
+			&cli.StringFlag{
+				Name:     "db-path",
+				Usage:    "Path to database",
+				Required: true,
+			},
+			cli.StringFlag{
+				Name:     "deploy-config",
+				Usage:    "Path to hardhat deploy config file",
+				Required: true,
+			},
+			cli.StringFlag{
+				Name:     "network",
+				Usage:    "Name of hardhat deploy network",
+				Required: true,
+			},
+			cli.StringFlag{
+				Name:     "hardhat-deployments",
+				Usage:    "Comma separated list of hardhat deployment directories",
+				Required: true,
+			},
+			cli.IntFlag{
+				Name:  "db-cache",
+				Usage: "LevelDB cache size in mb",
+				Value: 1024,
+			},
+			cli.IntFlag{
+				Name:  "db-handles",
+				Usage: "LevelDB number of handles",
+				Value: 60,
+			},
+		},
+		Action: func(ctx *cli.Context) error {
+			deployConfig := ctx.String("deploy-config")
+			config, err := genesis.NewDeployConfig(deployConfig)
+			if err != nil {
+				return err
+			}
+
+			ovmAddresses, err := crossdomain.NewAddresses(ctx.String("ovm-addresses"))
+			if err != nil {
+				return err
+			}
+			ovmAllowances, err := crossdomain.NewAllowances(ctx.String("ovm-allowances"))
+			if err != nil {
+				return err
+			}
+			ovmMessages, err := crossdomain.NewSentMessageFromJSON(ctx.String("ovm-messages"))
+			if err != nil {
+				return err
+			}
+			evmMessages, evmAddresses, err := crossdomain.ReadWitnessData(ctx.String("witness-file"))
+			if err != nil {
+				return err
+			}
+
+			log.Info(
+				"Loaded witness data",
+				"ovmAddresses", len(ovmAddresses),
+				"evmAddresses", len(evmAddresses),
+				"ovmAllowances", len(ovmAllowances),
+				"ovmMessages", len(ovmMessages),
+				"evmMessages", len(evmMessages),
+			)
+
+			migrationData := crossdomain.MigrationData{
+				OvmAddresses:  ovmAddresses,
+				EvmAddresses:  evmAddresses,
+				OvmAllowances: ovmAllowances,
+				OvmMessages:   ovmMessages,
+				EvmMessages:   evmMessages,
+			}
+
+			network := ctx.String("network")
+			deployments := strings.Split(ctx.String("hardhat-deployments"), ",")
+			hh, err := hardhat.New(network, []string{}, deployments)
+			if err != nil {
+				return err
+			}
+
+			l1RpcURL := ctx.String("l1-rpc-url")
+			l1Client, err := ethclient.Dial(l1RpcURL)
+			if err != nil {
+				return err
+			}
+
+			var block *types.Block
+			tag := config.L1StartingBlockTag
+			if tag.BlockNumber != nil {
+				block, err = l1Client.BlockByNumber(context.Background(), big.NewInt(tag.BlockNumber.Int64()))
+			} else if tag.BlockHash != nil {
+				block, err = l1Client.BlockByHash(context.Background(), *tag.BlockHash)
+			} else {
+				return fmt.Errorf("invalid l1StartingBlockTag in deploy config: %v", tag)
+			}
+			if err != nil {
+				return err
+			}
+
+			dbCache := ctx.Int("db-cache")
+			dbHandles := ctx.Int("db-handles")
+
+			// Read the required deployment addresses from disk if required
+			if err := config.GetDeployedAddresses(hh); err != nil {
+				return err
+			}
+
+			if err := config.Check(); err != nil {
+				return err
+			}
+
+			postLDB, err := db.Open(ctx.String("db-path"), dbCache, dbHandles)
+			if err != nil {
+				return err
+			}
+
+			if err := genesis.PostCheckMigratedDB(
+				postLDB,
+				migrationData,
+				&config.L1CrossDomainMessengerProxy,
+				config.L1ChainID,
+				config.FinalSystemOwner,
+				config.ProxyAdminOwner,
+				&derive.L1BlockInfo{
+					Number:        block.NumberU64(),
+					Time:          block.Time(),
+					BaseFee:       block.BaseFee(),
+					BlockHash:     block.Hash(),
+					BatcherAddr:   config.BatchSenderAddress,
+					L1FeeOverhead: eth.Bytes32(common.BigToHash(new(big.Int).SetUint64(config.GasPriceOracleOverhead))),
+					L1FeeScalar:   eth.Bytes32(common.BigToHash(new(big.Int).SetUint64(config.GasPriceOracleScalar))),
+				},
+			); err != nil {
+				return err
+			}
+
+			if err := postLDB.Close(); err != nil {
+				return err
+			}
+
+			return nil
+		},
+	}
+
+	if err := app.Run(os.Args); err != nil {
+		log.Crit("error in migration", "err", err)
+	}
+}

op-chain-ops/cmd/op-migrate/main.go

@@ -106,6 +106,11 @@ func main() {
 				Value:    "rollup.json",
 				Required: true,
 			},
+			cli.BoolFlag{
+				Name:     "post-check-only",
+				Usage:    "Only perform sanity checks",
+				Required: false,
+			},
 		},
 		Action: func(ctx *cli.Context) error {
 			deployConfig := ctx.String("deploy-config")

op-chain-ops/ether/cli.go

@@ -25,8 +25,8 @@ func GetOVMETHTotalSupplySlot() common.Hash {
 	return getOVMETHTotalSupplySlot()
 }
 
-// getOVMETHBalance gets a user's OVM ETH balance from state by querying the
+// GetOVMETHBalance gets a user's OVM ETH balance from state by querying the
 // appropriate storage slot directly.
-func getOVMETHBalance(db *state.StateDB, addr common.Address) *big.Int {
+func GetOVMETHBalance(db *state.StateDB, addr common.Address) *big.Int {
 	return db.GetState(OVMETHAddress, CalcOVMETHStorageKey(addr)).Big()
 }

op-chain-ops/ether/migrate.go

@@ -29,7 +29,9 @@ var (
 	}
 )
 
-func MigrateLegacyETH(db *state.StateDB, addresses []common.Address, chainID int, noCheck bool) error {
+type FilteredOVMETHAddresses []common.Address
+
+func MigrateLegacyETH(db *state.StateDB, addresses FilteredOVMETHAddresses, chainID int, noCheck bool) error {
 	// Chain params to use for integrity checking.
 	params := crossdomain.ParamsByChainID[chainID]
 	if params == nil {
@@ -39,28 +41,15 @@ func MigrateLegacyETH(db *state.StateDB, addresses []common.Address, chainID int
 	// Log the chain params for debugging purposes.
 	log.Info("Chain params", "chain-id", chainID, "supply-delta", params.ExpectedSupplyDelta)
 
-	// Deduplicate the list of addresses by converting to a map.
-	deduped := make(map[common.Address]bool)
-	for _, addr := range addresses {
-		deduped[addr] = true
-	}
-
 	// 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 addr := range deduped {
-		// No accounts should have a balance in state. If they do, bail.
-		if db.GetBalance(addr).Sign() > 0 {
-			if noCheck {
-				log.Error("account has non-zero balance in state - should never happen", "addr", addr)
-			} else {
-				log.Crit("account has non-zero balance in state - should never happen", "addr", addr)
-			}
-		}
+	for _, addr := range addresses {
+		// Balances are pre-checked not have any balances in state.
 
 		// Pull out the OVM ETH balance.
-		ovmBalance := getOVMETHBalance(db, addr)
+		ovmBalance := GetOVMETHBalance(db, addr)
 
 		// Actually perform the migration by setting the appropriate values in state.
 		db.SetBalance(addr, ovmBalance)

op-chain-ops/ether/precheck.go

 package ether
 
 import (
-	"errors"
 	"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"
@@ -11,132 +14,295 @@ import (
 	"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/ethdb"
 	"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(ldb ethdb.Database, db *state.StateDB, addresses []common.Address, allowances []*crossdomain.Allowance, chainID int, noCheck bool) ([]common.Address, error) {
+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 {
-		addrs = append(addrs, addr)
-		slotsInp[CalcOVMETHStorageKey(addr)] = 1
+		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 {
-		addrs = append(addrs, allowance.From)
-		slotsInp[CalcAllowanceStorageKey(allowance.From, allowance.To)] = 2
+		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")
-	addrs = append(addrs, sequencerEntrypointAddr)
-	slotsInp[CalcOVMETHStorageKey(sequencerEntrypointAddr)] = 1
+	entrySK := CalcOVMETHStorageKey(sequencerEntrypointAddr)
+	slotsAddrs[entrySK] = sequencerEntrypointAddr
+	slotsInp[entrySK] = 1
 
-	// Build a mapping of every storage slot in the LegacyERC20ETH contract, except the list of
-	// slots that we know we can ignore (totalSupply, name, symbol).
-	var count int
-	slotsAct := make(map[common.Hash]common.Hash)
-	progress := util.ProgressLogger(1000, "Read OVM_ETH storage slot")
-	err := db.ForEachStorage(predeploys.LegacyERC20ETHAddr, func(key, value common.Hash) bool {
-		progress()
+	// 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{})
 
-		// We can safely ignore specific slots (totalSupply, name, symbol).
-		if ignoredSlots[key] {
-			return true
+	// 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)
 		}
 
-		// Slot exists, so add it to the map.
-		slotsAct[key] = value
-		count++
-		return true
-	})
+		// 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 {
-		return nil, fmt.Errorf("cannot iterate over LegacyERC20ETHAddr: %w", err)
+			// 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")
 		}
 
-	// Log how many slots were iterated over.
-	log.Info("Iterated legacy balances", "count", count)
+		it := trie.NewIterator(st.NodeIterator(start.Bytes()))
 
-	// Iterate over the list of known slots and check that we have a slot for each one. We'll also
-	// keep track of the total balance to be migrated and throw if the total supply exceeds the
-	// expected supply delta.
-	totalFound := new(big.Int)
-	var unknown bool
-	for slot := range slotsAct {
-		slotType, ok := slotsInp[slot]
+		// 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", slot.String())
+					log.Error("ignoring unknown storage slot in state", "slot", key.String())
 				} else {
-				unknown = true
-				log.Error("unknown storage slot in state", "slot", slot.String())
-				continue
+					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:
-			// Balance slot.
-			totalFound.Add(totalFound, slotsAct[slot].Big())
+				// 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", slot, "type", slotType)
+					log.Error("unknown slot type", "slot", key, "type", slotType)
 				} else {
-				log.Crit("unknown slot type: %d", slotType)
+					log.Crit("unknown slot type %d, should never happen", slotType)
 				}
 			}
 		}
-	if unknown {
-		return nil, errors.New("unknown storage slots in state (see logs for details)")
 	}
 
+	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(params.ExpectedSupplyDelta) != 0 {
-		if noCheck {
+	if delta.Cmp(expDiff) != 0 {
 		log.Error(
 			"supply mismatch",
 			"migrated", totalFound.String(),
 			"supply", totalSupply.String(),
 			"delta", delta.String(),
-				"exp_delta", params.ExpectedSupplyDelta.String(),
-			)
-		} else {
-			log.Crit(
-				"supply mismatch",
-				"migrated", totalFound.String(),
-				"supply", totalSupply.String(),
-				"delta", delta.String(),
-				"exp_delta", params.ExpectedSupplyDelta.String(),
+			"exp_delta", expDiff.String(),
 		)
+		if !noCheck {
+			return nil, fmt.Errorf("supply mismatch: %s", delta.String())
 		}
 	}
 
@@ -146,7 +312,7 @@ func PreCheckBalances(ldb ethdb.Database, db *state.StateDB, addresses []common.
 		"migrated", totalFound.String(),
 		"supply", totalSupply.String(),
 		"delta", delta.String(),
-		"exp_delta", params.ExpectedSupplyDelta.String(),
+		"exp_delta", expDiff.String(),
 	)
 
 	// We know we have at least a superset of all addresses here since we know that we have every

op-chain-ops/ether/precheck_test.go

+package ether
+
+import (
+	"bytes"
+	"math/big"
+	"math/rand"
+	"os"
+	"sort"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/log"
+
+	"github.com/ethereum-optimism/optimism/op-chain-ops/crossdomain"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/trie"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPreCheckBalances(t *testing.T) {
+	log.Root().SetHandler(log.StreamHandler(os.Stderr, log.TerminalFormat(true)))
+
+	tests := []struct {
+		name            string
+		totalSupply     *big.Int
+		expDiff         *big.Int
+		stateBalances   map[common.Address]*big.Int
+		stateAllowances map[common.Address]common.Address
+		inputAddresses  []common.Address
+		inputAllowances []*crossdomain.Allowance
+		check           func(t *testing.T, addrs FilteredOVMETHAddresses, err error)
+	}{
+		{
+			name:        "everything matches",
+			totalSupply: big.NewInt(3),
+			expDiff:     big.NewInt(0),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+				common.HexToAddress("0x456"): big.NewInt(2),
+			},
+			stateAllowances: map[common.Address]common.Address{
+				common.HexToAddress("0x123"): common.HexToAddress("0x456"),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+				common.HexToAddress("0x456"),
+			},
+			inputAllowances: []*crossdomain.Allowance{
+				{
+					From: common.HexToAddress("0x123"),
+					To:   common.HexToAddress("0x456"),
+				},
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.NoError(t, err)
+				require.EqualValues(t, FilteredOVMETHAddresses{
+					common.HexToAddress("0x123"),
+					common.HexToAddress("0x456"),
+				}, addrs)
+			},
+		},
+		{
+			name:        "extra input addresses",
+			totalSupply: big.NewInt(1),
+			expDiff:     big.NewInt(0),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+				common.HexToAddress("0x456"),
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.NoError(t, err)
+				require.EqualValues(t, FilteredOVMETHAddresses{
+					common.HexToAddress("0x123"),
+				}, addrs)
+			},
+		},
+		{
+			name:        "extra input allowances",
+			totalSupply: big.NewInt(1),
+			expDiff:     big.NewInt(0),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+			},
+			stateAllowances: map[common.Address]common.Address{
+				common.HexToAddress("0x123"): common.HexToAddress("0x456"),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+				common.HexToAddress("0x456"),
+			},
+			inputAllowances: []*crossdomain.Allowance{
+				{
+					From: common.HexToAddress("0x123"),
+					To:   common.HexToAddress("0x456"),
+				},
+				{
+					From: common.HexToAddress("0x123"),
+					To:   common.HexToAddress("0x789"),
+				},
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.NoError(t, err)
+				require.EqualValues(t, FilteredOVMETHAddresses{
+					common.HexToAddress("0x123"),
+				}, addrs)
+			},
+		},
+		{
+			name:        "missing input addresses",
+			totalSupply: big.NewInt(2),
+			expDiff:     big.NewInt(0),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+				common.HexToAddress("0x456"): big.NewInt(1),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.Error(t, err)
+				require.ErrorContains(t, err, "unknown storage slot")
+			},
+		},
+		{
+			name:        "missing input allowances",
+			totalSupply: big.NewInt(2),
+			expDiff:     big.NewInt(0),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+			},
+			stateAllowances: map[common.Address]common.Address{
+				common.HexToAddress("0x123"): common.HexToAddress("0x456"),
+				common.HexToAddress("0x123"): common.HexToAddress("0x789"),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+			},
+			inputAllowances: []*crossdomain.Allowance{
+				{
+					From: common.HexToAddress("0x123"),
+					To:   common.HexToAddress("0x456"),
+				},
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.Error(t, err)
+				require.ErrorContains(t, err, "unknown storage slot")
+			},
+		},
+		{
+			name:        "bad supply diff",
+			totalSupply: big.NewInt(4),
+			expDiff:     big.NewInt(0),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+				common.HexToAddress("0x456"): big.NewInt(2),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+				common.HexToAddress("0x456"),
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.Error(t, err)
+				require.ErrorContains(t, err, "supply mismatch")
+			},
+		},
+		{
+			name:        "good supply diff",
+			totalSupply: big.NewInt(4),
+			expDiff:     big.NewInt(1),
+			stateBalances: map[common.Address]*big.Int{
+				common.HexToAddress("0x123"): big.NewInt(1),
+				common.HexToAddress("0x456"): big.NewInt(2),
+			},
+			inputAddresses: []common.Address{
+				common.HexToAddress("0x123"),
+				common.HexToAddress("0x456"),
+			},
+			check: func(t *testing.T, addrs FilteredOVMETHAddresses, err error) {
+				require.NoError(t, err)
+				require.EqualValues(t, FilteredOVMETHAddresses{
+					common.HexToAddress("0x123"),
+					common.HexToAddress("0x456"),
+				}, addrs)
+			},
+		},
+	}
+	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)
+		})
+	}
+}
+
+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{
+		Preimages: true,
+		Cache:     1024,
+	}), nil)
+	require.NoError(t, err)
+
+	db.CreateAccount(OVMETHAddress)
+	db.SetState(OVMETHAddress, getOVMETHTotalSupplySlot(), common.BigToHash(totalSupply))
+
+	for slot := range ignoredSlots {
+		if slot == getOVMETHTotalSupplySlot() {
+			continue
+		}
+		db.SetState(OVMETHAddress, slot, common.Hash{31: 0xff})
+	}
+	for addr, balance := range balances {
+		db.SetState(OVMETHAddress, CalcOVMETHStorageKey(addr), common.BigToHash(balance))
+	}
+	for from, to := range allowances {
+		db.SetState(OVMETHAddress, CalcAllowanceStorageKey(from, to), common.BigToHash(big.NewInt(1)))
+	}
+
+	root, err := db.Commit(false)
+	require.NoError(t, err)
+
+	err = db.Database().TrieDB().Commit(root, true)
+	require.NoError(t, err)
+
+	return db
+}
+
+// TestPreCheckBalancesRandom 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) {
+	addresses := make([]common.Address, 0)
+	stateBalances := make(map[common.Address]*big.Int)
+
+	allowances := make([]*crossdomain.Allowance, 0)
+	stateAllowances := make(map[common.Address]common.Address)
+
+	totalSupply := big.NewInt(0)
+
+	for i := 0; i < 100; i++ {
+		for i := 0; i < rand.Intn(1000); i++ {
+			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++ {
+			addr := randAddr(t)
+			to := randAddr(t)
+			allowances = append(allowances, &crossdomain.Allowance{
+				From: addr,
+				To:   to,
+			})
+			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)
+		require.NoError(t, err)
+
+		sort.Slice(outAddrs, func(i, j int) bool {
+			return bytes.Compare(outAddrs[i][:], outAddrs[j][:]) < 0
+		})
+		require.EqualValues(t, addresses, outAddrs)
+	}
+}
+
+func randAddr(t *testing.T) common.Address {
+	var addr common.Address
+	_, err := rand.Read(addr[:])
+	require.NoError(t, err)
+	return addr
+}

op-chain-ops/genesis/check.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"math/big"
 
+	"github.com/ethereum-optimism/optimism/op-chain-ops/ether"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/state"
@@ -251,7 +253,7 @@ func PostCheckPredeploys(prevDB, currDB *state.StateDB) error {
 
 		// Balances and nonces should match legacy
 		oldNonce := prevDB.GetNonce(addr)
-		oldBalance := prevDB.GetBalance(addr)
+		oldBalance := ether.GetOVMETHBalance(prevDB, addr)
 		newNonce := currDB.GetNonce(addr)
 		newBalance := currDB.GetBalance(addr)
 		if oldNonce != newNonce {
@@ -543,7 +545,7 @@ func CheckWithdrawalsAfter(db vm.StateDB, data crossdomain.MigrationData, l1Cros
 
 		// If the sender is _not_ the L2XDM, the value should not be migrated.
 		wd := wdsByOldSlot[key]
-		if wd.XDomainSender == predeploys.L2CrossDomainMessengerAddr {
+		if wd.MessageSender == predeploys.L2CrossDomainMessengerAddr {
 			// Make sure the value is abiTrue if this withdrawal should be migrated.
 			if migratedValue != abiTrue {
 				innerErr = fmt.Errorf("expected migrated value to be true, but got %s", migratedValue)
@@ -552,7 +554,7 @@ func CheckWithdrawalsAfter(db vm.StateDB, data crossdomain.MigrationData, l1Cros
 		} else {
 			// Otherwise, ensure that withdrawals from senders other than the L2XDM are _not_ migrated.
 			if migratedValue != abiFalse {
-				innerErr = fmt.Errorf("a migration from a sender other than the L2XDM was migrated")
+				innerErr = fmt.Errorf("a migration from a sender other than the L2XDM was migrated. sender: %s, migrated value: %s", wd.MessageSender, migratedValue)
 				return false
 			}
 		}

op-chain-ops/genesis/db_migration.go

@@ -82,6 +82,7 @@ func MigrateDB(ldb ethdb.Database, config *DeployConfig, l1Block *types.Block, m
 		)
 	}
 
+	dbFactory := func() (*state.StateDB, error) {
 		// Set up the backing store.
 		underlyingDB := state.NewDatabaseWithConfig(ldb, &trie.Config{
 			Preimages: true,
@@ -94,6 +95,14 @@ func MigrateDB(ldb ethdb.Database, config *DeployConfig, l1Block *types.Block, m
 			return nil, fmt.Errorf("cannot open StateDB: %w", err)
 		}
 
+		return db, nil
+	}
+
+	db, err := dbFactory()
+	if err != nil {
+		return nil, fmt.Errorf("cannot create StateDB: %w", err)
+	}
+
 	// Before we do anything else, we need to ensure that all of the input configuration is correct
 	// and nothing is missing. We'll first verify the contract configuration, then we'll verify the
 	// witness data for the migration. We operate under the assumption that the witness data is
@@ -139,7 +148,7 @@ func MigrateDB(ldb ethdb.Database, config *DeployConfig, l1Block *types.Block, m
 	// 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(ldb, db, migrationData.Addresses(), migrationData.OvmAllowances, int(config.L1ChainID), 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)
 	}