From 2fdc0da8b867c55a2d7e16b7e757eeb38060d9dc Mon Sep 17 00:00:00 2001
From: Guillaume Ballet <3272758+gballet@users.noreply.github.com>
Date: Tue, 24 Oct 2023 12:43:40 +0200
Subject: [PATCH] review feedback

---
 cmd/geth/chaincmd.go |  8 +++---
 cmd/utils/cmd.go     | 66 +++++++++++++++++++++++++-------------------
 cmd/utils/flags.go   |  5 ----
 3 files changed, 42 insertions(+), 37 deletions(-)

diff --git a/cmd/geth/chaincmd.go b/cmd/geth/chaincmd.go
index 741e2f3e29..fec7830644 100644
--- a/cmd/geth/chaincmd.go
+++ b/cmd/geth/chaincmd.go
@@ -171,8 +171,8 @@ This command dumps out the state for a given block (or latest, if none provided)
 		Action:    exportSnapshotPreimages,
 		Name:      "export-snapshot-preimages",
 		Usage:     "Export the preimage in snapshot enumeration order",
-		ArgsUsage: "<dumpfile>",
-		Flags:     flags.Merge([]cli.Flag{utils.TreeRootFlag}, utils.DatabaseFlags),
+		ArgsUsage: "<dumpfile> [<root>]",
+		Flags:     utils.DatabaseFlags,
 		Description: `
 The export-snapshot-preimages command exports hash preimages to a flat file, in exactly
 the expected order for the overlay tree migration.
@@ -428,8 +428,8 @@ func exportSnapshotPreimages(ctx *cli.Context) error {
 	chain, _ := utils.MakeChain(ctx, stack, true)
 
 	var root common.Hash
-	if ctx.String(utils.TreeRootFlag.Name) != "" {
-		rootBytes := common.FromHex(ctx.String(utils.StartKeyFlag.Name))
+	if ctx.Args().Len() > 1 {
+		rootBytes := common.FromHex(ctx.Args().Get(1))
 		if len(rootBytes) != common.HashLength {
 			return fmt.Errorf("invalid root hash length")
 		}
diff --git a/cmd/utils/cmd.go b/cmd/utils/cmd.go
index 601096bfbc..2e4004e726 100644
--- a/cmd/utils/cmd.go
+++ b/cmd/utils/cmd.go
@@ -397,45 +397,55 @@ func ExportSnapshotPreimages(chain *core.BlockChain, fn string, root common.Hash
 		root = chain.CurrentBlock().Root
 	}
 
-	accIt, err := chain.Snapshots().AccountIterator(root, common.Hash{})
-	if err != nil {
-		return err
+	type hashAndPreimageSize struct {
+		Hash common.Hash
+		Size int
 	}
-	defer accIt.Release()
+	var hashCh chan hashAndPreimageSize
 
-	count := 0
-	for accIt.Next() {
-		acc, err := types.FullAccount(accIt.Account())
+	go func() {
+		defer close(hashCh)
+		accIt, err := chain.Snapshots().AccountIterator(root, common.Hash{})
 		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)
+			log.Error("failed to create account iterator", "error", err)
+			return
 		}
+		defer accIt.Release()
 
-		if acc.Root != (common.Hash{}) && acc.Root != types.EmptyRootHash {
-			stIt, err := chain.Snapshots().StorageIterator(root, accIt.Hash(), common.Hash{})
+		count := 0
+		for accIt.Next() {
+			acc, err := types.FullAccount(accIt.Account())
 			if err != nil {
-				return fmt.Errorf("failed to create storage iterator: %w", err)
+				log.Error("failed to get full account", "error", err)
+				return
 			}
-			for stIt.Next() {
-				slotnr := rawdb.ReadPreimage(statedb.Database().DiskDB(), stIt.Hash())
-				if len(slotnr) != 32 {
-					return fmt.Errorf("slotnr not 32 len")
+			hashCh <- hashAndPreimageSize{Hash: accIt.Hash(), Size: 20}
+
+			if acc.Root != (common.Hash{}) && acc.Root != types.EmptyRootHash {
+				stIt, err := chain.Snapshots().StorageIterator(root, accIt.Hash(), common.Hash{})
+				if err != nil {
+					log.Error("failed to create storage iterator", "error", err)
+					return
 				}
-				if _, err := writer.Write(slotnr); err != nil {
-					return fmt.Errorf("failed to write slotnr preimage: %w", err)
+				for stIt.Next() {
+					hashCh <- hashAndPreimageSize{Hash: stIt.Hash(), Size: 32}
 				}
+				stIt.Release()
+			}
+			count++
+			if count%100000 == 0 {
+				log.Info("Last exported account", "account", accIt.Hash())
 			}
-			stIt.Release()
 		}
-		count++
-		if count%100000 == 0 {
-			log.Info("Last exported account", "account", accIt.Hash())
+	}()
+
+	for item := range hashCh {
+		preimage := rawdb.ReadPreimage(statedb.Database().DiskDB(), item.Hash)
+		if len(preimage) != item.Size {
+			return fmt.Errorf("invalid preimage size")
+		}
+		if _, err := writer.Write(preimage); err != nil {
+			return fmt.Errorf("failed to write preimage: %w", err)
 		}
 	}
 
diff --git a/cmd/utils/flags.go b/cmd/utils/flags.go
index c3dcf875b6..9743a7b9ca 100644
--- a/cmd/utils/flags.go
+++ b/cmd/utils/flags.go
@@ -219,11 +219,6 @@ 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 that of current block)",
-		Value: "",
-	}
 
 	defaultSyncMode = ethconfig.Defaults.SyncMode
 	SnapshotFlag    = &cli.BoolFlag{