Overlay Tree preimages exporting and usage

cmd/geth/chaincmd.go

@@ -143,6 +143,17 @@ It's deprecated, please use "geth db import" instead.
 		Description: `
 The export-preimages command exports hash preimages to an RLP encoded stream.
 It's deprecated, please use "geth db export" instead.
+`,
+	}
+	exportOverlayPreimagesCommand = &cli.Command{
+		Action:    exportOverlayPreimages,
+		Name:      "export-overlay-preimages",
+		Usage:     "Export the preimage in overlay tree migration order",
+		ArgsUsage: "<dumpfile>",
+		Flags:     flags.Merge([]cli.Flag{utils.TreeRootFlag}, utils.DatabasePathFlags),
+		Description: `
+The export-overlay-preimages command exports hash preimages to a flat file, in exactly
+the expected order for the overlay tree migration.
 `,
 	}
 	dumpCommand = &cli.Command{
@@ -394,6 +405,33 @@ func exportPreimages(ctx *cli.Context) error {
 	return nil
 }
 
+// exportOverlayPreimages dumps the preimage data to a flat file.
+func exportOverlayPreimages(ctx *cli.Context) error {
+	if ctx.Args().Len() < 1 {
+		utils.Fatalf("This command requires an argument.")
+	}
+	stack, _ := makeConfigNode(ctx)
+	defer stack.Close()
+
+	chain, _ := utils.MakeChain(ctx, stack)
+
+	var root common.Hash
+	if ctx.String(utils.TreeRootFlag.Name) != "" {
+		rootBytes := common.FromHex(ctx.String(utils.StartKeyFlag.Name))
+		if len(rootBytes) != common.HashLength {
+			return fmt.Errorf("invalid root hash length")
+		}
+		root = common.BytesToHash(rootBytes)
+	}
+
+	start := time.Now()
+	if err := utils.ExportOverlayPreimages(chain, ctx.Args().First(), root); err != nil {
+		utils.Fatalf("Export error: %v\n", err)
+	}
+	fmt.Printf("Export done in %v\n", time.Since(start))
+	return nil
+}
+
 func parseDumpConfig(ctx *cli.Context, stack *node.Node) (*state.DumpConfig, ethdb.Database, common.Hash, error) {
 	db := utils.MakeChainDatabase(ctx, stack, true)
 	var header *types.Header

cmd/geth/main.go

@@ -213,6 +213,7 @@ func init() {
 		exportCommand,
 		importPreimagesCommand,
 		exportPreimagesCommand,
+		exportOverlayPreimagesCommand,
 		removedbCommand,
 		dumpCommand,
 		dumpGenesisCommand,

cmd/utils/cmd.go

@@ -34,6 +34,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/rawdb"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/eth/ethconfig"
@@ -379,6 +380,75 @@ func ExportPreimages(db ethdb.Database, fn string) error {
 	return nil
 }
 
+// ExportOverlayPreimages exports all known hash preimages into the specified file,
+// in the same order as expected by the overlay tree migration.
+func ExportOverlayPreimages(chain *core.BlockChain, fn string, root common.Hash) error {
+	log.Info("Exporting preimages", "file", fn)
+
+	fh, err := os.OpenFile(fn, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.ModePerm)
+	if err != nil {
+		return err
+	}
+	defer fh.Close()
+
+	writer := bufio.NewWriter(fh)
+	defer writer.Flush()
+
+	statedb, err := chain.State()
+	if err != nil {
+		return fmt.Errorf("failed to open statedb: %w", err)
+	}
+
+	if root == (common.Hash{}) {
+		root = chain.CurrentBlock().Root()
+	}
+
+	accIt, err := statedb.Snaps().AccountIterator(root, common.Hash{})
+	if err != nil {
+		return err
+	}
+	defer accIt.Release()
+
+	count := 0
+	for accIt.Next() {
+		acc, err := snapshot.FullAccount(accIt.Account())
+		if err != nil {
+			return fmt.Errorf("invalid account encountered during traversal: %s", err)
+		}
+		addr := rawdb.ReadPreimage(statedb.Database().DiskDB(), accIt.Hash())
+		if len(addr) != 20 {
+			return fmt.Errorf("addr len is zero is not 32: %d", len(addr))
+		}
+		if _, err := writer.Write(addr); err != nil {
+			return fmt.Errorf("failed to write addr preimage: %w", err)
+		}
+
+		if acc.HasStorage() {
+			stIt, err := statedb.Snaps().StorageIterator(root, accIt.Hash(), common.Hash{})
+			if err != nil {
+				return fmt.Errorf("failed to create storage iterator: %w", err)
+			}
+			for stIt.Next() {
+				slotnr := rawdb.ReadPreimage(statedb.Database().DiskDB(), stIt.Hash())
+				if len(slotnr) != 32 {
+					return fmt.Errorf("slotnr not 32 len")
+				}
+				if _, err := writer.Write(slotnr); err != nil {
+					return fmt.Errorf("failed to write slotnr preimage: %w", err)
+				}
+			}
+			stIt.Release()
+		}
+		count++
+		if count%100000 == 0 {
+			log.Info("Last exported account", "account", accIt.Hash())
+		}
+	}
+
+	log.Info("Exported preimages", "file", fn)
+	return nil
+}
+
 // exportHeader is used in the export/import flow. When we do an export,
 // the first element we output is the exportHeader.
 // Whenever a backwards-incompatible change is made, the Version header

cmd/utils/flags.go

@@ -219,6 +219,11 @@ var (
 		Usage: "Max number of elements (0 = no limit)",
 		Value: 0,
 	}
+	TreeRootFlag = &cli.StringFlag{
+		Name:  "roothash",
+		Usage: "Root hash of the tree (if empty, use the latest)",
+		Value: "",
+	}
 
 	defaultSyncMode = ethconfig.Defaults.SyncMode
 	SyncModeFlag    = &flags.TextMarshalerFlag{

core/state/database.go

@@ -26,6 +26,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/trie"
@@ -73,18 +74,24 @@ type Database interface {
 
 	Transitioned() bool
 
-	SetCurrentAccountHash(hash common.Hash)
+	SetCurrentSlotHash(hash common.Hash)
 
-	GetCurrentAccountHash() common.Hash
+	GetCurrentAccountAddress() *common.Address
 
-	SetCurrentSlotHash(hash common.Hash)
+	SetCurrentAccountAddress(common.Address)
+
+	GetCurrentAccountHash() common.Hash
 
 	GetCurrentSlotHash() common.Hash
 
 	SetStorageProcessed(bool)
 
 	GetStorageProcessed() bool
 
+	GetCurrentPreimageOffset() int64
+
+	SetCurrentPreimageOffset(int64)
+
 	AddRootTranslation(originalRoot, translatedRoot common.Hash)
 }
 
@@ -246,9 +253,10 @@ type cachingDB struct {
 
 	addrToPoint *utils.PointCache
 
-	baseRoot           common.Hash // hash of the read-only base tree
-	CurrentAccountHash common.Hash // hash of the last translated account
-	CurrentSlotHash    common.Hash // hash of the last translated storage slot
+	baseRoot              common.Hash     // hash of the read-only base tree
+	CurrentAccountAddress *common.Address // addresss of the last translated account
+	CurrentSlotHash       common.Hash     // hash of the last translated storage slot
+	CurrentPreimageOffset int64           // next byte to read from the preimage file
 
 	// Mark whether the storage for an account has been processed. This is useful if the
 	// maximum number of leaves of the conversion is reached before the whole storage is
@@ -409,12 +417,28 @@ func (db *cachingDB) GetTreeKeyHeader(addr []byte) *verkle.Point {
 	return db.addrToPoint.GetTreeKeyHeader(addr)
 }
 
-func (db *cachingDB) SetCurrentAccountHash(hash common.Hash) {
-	db.CurrentAccountHash = hash
+func (db *cachingDB) SetCurrentAccountAddress(addr common.Address) {
+	db.CurrentAccountAddress = &addr
 }
 
 func (db *cachingDB) GetCurrentAccountHash() common.Hash {
-	return db.CurrentAccountHash
+	var addrHash common.Hash
+	if db.CurrentAccountAddress != nil {
+		addrHash = crypto.Keccak256Hash(db.CurrentAccountAddress[:])
+	}
+	return addrHash
+}
+
+func (db *cachingDB) GetCurrentAccountAddress() *common.Address {
+	return db.CurrentAccountAddress
+}
+
+func (db *cachingDB) GetCurrentPreimageOffset() int64 {
+	return db.CurrentPreimageOffset
+}
+
+func (db *cachingDB) SetCurrentPreimageOffset(offset int64) {
+	db.CurrentPreimageOffset = offset
 }
 
 func (db *cachingDB) SetCurrentSlotHash(hash common.Hash) {