cmd/geth: add inspect trie tool to analysis trie storage

cmd/geth/dbcmd.go

@@ -19,6 +19,7 @@ package main
 import (
 	"bytes"
 	"fmt"
+	"math"
 	"os"
 	"os/signal"
 	"path/filepath"
@@ -37,6 +38,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/internal/tablewriter"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
@@ -53,6 +55,23 @@ var (
 		Name:  "remove.chain",
 		Usage: "If set, selects the state data for removal",
 	}
+	inspectTrieTopFlag = &cli.IntFlag{
+		Name:  "top",
+		Usage: "Print the top N results per ranking category",
+		Value: 10,
+	}
+	inspectTrieDumpPathFlag = &cli.StringFlag{
+		Name:  "dump-path",
+		Usage: "Path for the trie statistics dump file",
+	}
+	inspectTrieSummarizeFlag = &cli.StringFlag{
+		Name:  "summarize",
+		Usage: "Summarize an existing trie dump file (skip trie traversal)",
+	}
+	inspectTrieContractFlag = &cli.StringFlag{
+		Name:  "contract",
+		Usage: "Inspect only the storage of the given contract address (skips full account trie walk)",
+	}
 
 	removedbCommand = &cli.Command{
 		Action:    removeDB,
@@ -74,6 +93,7 @@ Remove blockchain and state databases`,
 			dbCompactCmd,
 			dbGetCmd,
 			dbDeleteCmd,
+			dbInspectTrieCmd,
 			dbPutCmd,
 			dbGetSlotsCmd,
 			dbDumpFreezerIndex,
@@ -92,6 +112,22 @@ Remove blockchain and state databases`,
 		Usage:       "Inspect the storage size for each type of data in the database",
 		Description: `This commands iterates the entire database. If the optional 'prefix' and 'start' arguments are provided, then the iteration is limited to the given subset of data.`,
 	}
+	dbInspectTrieCmd = &cli.Command{
+		Action:    inspectTrie,
+		Name:      "inspect-trie",
+		ArgsUsage: "<blocknum>",
+		Flags: slices.Concat([]cli.Flag{
+			utils.ExcludeStorageFlag,
+			inspectTrieTopFlag,
+			utils.OutputFileFlag,
+			inspectTrieDumpPathFlag,
+			inspectTrieSummarizeFlag,
+			inspectTrieContractFlag,
+		}, utils.NetworkFlags, utils.DatabaseFlags),
+		Usage: "Print detailed trie information about the structure of account trie and storage tries.",
+		Description: `This commands iterates the entrie trie-backed state. If the 'blocknum' is not specified, 
+the latest block number will be used by default.`,
+	}
 	dbCheckStateContentCmd = &cli.Command{
 		Action:    checkStateContent,
 		Name:      "check-state-content",
@@ -385,6 +421,88 @@ func checkStateContent(ctx *cli.Context) error {
 	return nil
 }
 
+func inspectTrie(ctx *cli.Context) error {
+	topN := ctx.Int(inspectTrieTopFlag.Name)
+	if topN <= 0 {
+		return fmt.Errorf("invalid --%s value %d (must be > 0)", inspectTrieTopFlag.Name, topN)
+	}
+	config := &trie.InspectConfig{
+		NoStorage: ctx.Bool(utils.ExcludeStorageFlag.Name),
+		TopN:      topN,
+		Path:      ctx.String(utils.OutputFileFlag.Name),
+	}
+
+	if summarizePath := ctx.String(inspectTrieSummarizeFlag.Name); summarizePath != "" {
+		if ctx.NArg() > 0 {
+			return fmt.Errorf("block number argument is not supported with --%s", inspectTrieSummarizeFlag.Name)
+		}
+		config.DumpPath = summarizePath
+		log.Info("Summarizing trie dump", "path", summarizePath, "top", topN)
+		return trie.Summarize(summarizePath, config)
+	}
+	if ctx.NArg() > 1 {
+		return fmt.Errorf("excessive number of arguments: %v", ctx.Command.ArgsUsage)
+	}
+
+	stack, _ := makeConfigNode(ctx)
+	db := utils.MakeChainDatabase(ctx, stack, false)
+	defer stack.Close()
+	defer db.Close()
+
+	var (
+		trieRoot common.Hash
+		hash     common.Hash
+		number   uint64
+	)
+	switch {
+	case ctx.NArg() == 0 || ctx.Args().Get(0) == "latest":
+		head := rawdb.ReadHeadHeaderHash(db)
+		n, ok := rawdb.ReadHeaderNumber(db, head)
+		if !ok {
+			return fmt.Errorf("could not load head block hash")
+		}
+		number = n
+	case ctx.Args().Get(0) == "snapshot":
+		trieRoot = rawdb.ReadSnapshotRoot(db)
+		number = math.MaxUint64
+	default:
+		var err error
+		number, err = strconv.ParseUint(ctx.Args().Get(0), 10, 64)
+		if err != nil {
+			return fmt.Errorf("failed to parse blocknum, Args[0]: %v, err: %v", ctx.Args().Get(0), err)
+		}
+	}
+
+	if number != math.MaxUint64 {
+		hash = rawdb.ReadCanonicalHash(db, number)
+		if hash == (common.Hash{}) {
+			return fmt.Errorf("canonical hash for block %d not found", number)
+		}
+		blockHeader := rawdb.ReadHeader(db, hash, number)
+		trieRoot = blockHeader.Root
+	}
+	if trieRoot == (common.Hash{}) {
+		log.Error("Empty root hash")
+	}
+
+	config.DumpPath = ctx.String(inspectTrieDumpPathFlag.Name)
+	if config.DumpPath == "" {
+		config.DumpPath = stack.ResolvePath("trie-dump.bin")
+	}
+
+	triedb := utils.MakeTrieDatabase(ctx, stack, db, false, true, false)
+	defer triedb.Close()
+
+	if contractAddr := ctx.String(inspectTrieContractFlag.Name); contractAddr != "" {
+		address := common.HexToAddress(contractAddr)
+		log.Info("Inspecting contract", "address", address, "root", trieRoot, "block", number)
+		return trie.InspectContract(triedb, db, trieRoot, address)
+	}
+
+	log.Info("Inspecting trie", "root", trieRoot, "block", number, "dump", config.DumpPath, "top", topN)
+	return trie.Inspect(triedb, trieRoot, config)
+}
+
 func showDBStats(db ethdb.KeyValueStater) {
 	stats, err := db.Stat()
 	if err != nil {
@@ -759,7 +877,7 @@ func showMetaData(ctx *cli.Context) error {
 		data = append(data, []string{"headHeader.Root", fmt.Sprintf("%v", h.Root)})
 		data = append(data, []string{"headHeader.Number", fmt.Sprintf("%d (%#x)", h.Number, h.Number)})
 	}
-	table := rawdb.NewTableWriter(os.Stdout)
+	table := tablewriter.NewWriter(os.Stdout)
 	table.SetHeader([]string{"Field", "Value"})
 	table.AppendBulk(data)
 	table.Render()

cmd/utils/flags.go

@@ -218,6 +218,16 @@ var (
 		Usage: "Max number of elements (0 = no limit)",
 		Value: 0,
 	}
+	TopFlag = &cli.IntFlag{
+		Name:  "top",
+		Usage: "Print the top N results",
+		Value: 5,
+	}
+	OutputFileFlag = &cli.StringFlag{
+		Name:  "output",
+		Usage: "Writes the result in json to the output",
+		Value: "",
+	}
 
 	SnapshotFlag = &cli.BoolFlag{
 		Name:     "snapshot",

core/rawdb/database.go

@@ -35,6 +35,7 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/ethdb/memorydb"
+	"github.com/ethereum/go-ethereum/internal/tablewriter"
 	"github.com/ethereum/go-ethereum/log"
 	"golang.org/x/sync/errgroup"
 )
@@ -650,7 +651,7 @@ func InspectDatabase(db ethdb.Database, keyPrefix, keyStart []byte) error {
 		total.Add(uint64(ancient.size()))
 	}
 
-	table := NewTableWriter(os.Stdout)
+	table := tablewriter.NewWriter(os.Stdout)
 	table.SetHeader([]string{"Database", "Category", "Size", "Items"})
 	table.SetFooter([]string{"", "Total", common.StorageSize(total.Load()).String(), fmt.Sprintf("%d", count.Load())})
 	table.AppendBulk(stats)

core/stateless/stats.go

@@ -27,6 +27,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/metrics"
+	"github.com/ethereum/go-ethereum/trie"
 )
 
 var accountTrieLeavesAtDepth [16]*metrics.Counter
@@ -41,59 +42,68 @@ func init() {
 
 // WitnessStats aggregates statistics for account and storage trie accesses.
 type WitnessStats struct {
-	accountTrieLeaves [16]int64
-	storageTrieLeaves [16]int64
+	accountTrie *trie.LevelStats
+	storageTrie *trie.LevelStats
 }
 
 // NewWitnessStats creates a new WitnessStats collector.
 func NewWitnessStats() *WitnessStats {
-	return &WitnessStats{}
+	return &WitnessStats{
+		accountTrie: trie.NewLevelStats(),
+		storageTrie: trie.NewLevelStats(),
+	}
+}
+
+func (s *WitnessStats) init() {
+	if s.accountTrie == nil {
+		s.accountTrie = trie.NewLevelStats()
+	}
+	if s.storageTrie == nil {
+		s.storageTrie = trie.NewLevelStats()
+	}
 }
 
 // Add records trie access depths from the given node paths.
 // If `owner` is the zero hash, accesses are attributed to the account trie;
 // otherwise, they are attributed to the storage trie of that account.
 func (s *WitnessStats) Add(nodes map[string][]byte, owner common.Hash) {
-	// Extract paths from the nodes map
+	s.init()
+
+	// Extract paths from the nodes map.
 	paths := slices.Collect(maps.Keys(nodes))
 	sort.Strings(paths)
 
+	ownerStat := s.accountTrie
+	if owner != (common.Hash{}) {
+		ownerStat = s.storageTrie
+	}
+
 	for i, path := range paths {
 		// If current path is a prefix of the next path, it's not a leaf.
 		// The last path is always a leaf.
 		if i == len(paths)-1 || !strings.HasPrefix(paths[i+1], paths[i]) {
-			depth := len(path)
-			if owner == (common.Hash{}) {
-				if depth >= len(s.accountTrieLeaves) {
-					depth = len(s.accountTrieLeaves) - 1
-				}
-				s.accountTrieLeaves[depth] += 1
-			} else {
-				if depth >= len(s.storageTrieLeaves) {
-					depth = len(s.storageTrieLeaves) - 1
-				}
-				s.storageTrieLeaves[depth] += 1
-			}
+			ownerStat.AddLeaf(len(path))
 		}
 	}
 }
 
 // ReportMetrics reports the collected statistics to the global metrics registry.
 func (s *WitnessStats) ReportMetrics(blockNumber uint64) {
-	// Encode the metrics as JSON for easier consumption
-	accountLeavesJson, _ := json.Marshal(s.accountTrieLeaves)
-	storageLeavesJson, _ := json.Marshal(s.storageTrieLeaves)
-
-	// Log account trie depth statistics
-	log.Info("Account trie depth stats",
-		"block", blockNumber,
-		"leavesAtDepth", string(accountLeavesJson))
-	log.Info("Storage trie depth stats",
-		"block", blockNumber,
-		"leavesAtDepth", string(storageLeavesJson))
+	s.init()
 
-	for i := 0; i < 16; i++ {
-		accountTrieLeavesAtDepth[i].Inc(s.accountTrieLeaves[i])
-		storageTrieLeavesAtDepth[i].Inc(s.storageTrieLeaves[i])
+	accountTrieLeaves := s.accountTrie.LeafDepths()
+	storageTrieLeaves := s.storageTrie.LeafDepths()
+
+	// Encode the metrics as JSON for easier consumption.
+	accountLeavesJSON, _ := json.Marshal(accountTrieLeaves)
+	storageLeavesJSON, _ := json.Marshal(storageTrieLeaves)
+
+	// Log account trie depth statistics.
+	log.Info("Account trie depth stats", "block", blockNumber, "leavesAtDepth", string(accountLeavesJSON))
+	log.Info("Storage trie depth stats", "block", blockNumber, "leavesAtDepth", string(storageLeavesJSON))
+
+	for i := 0; i < len(accountTrieLeavesAtDepth); i++ {
+		accountTrieLeavesAtDepth[i].Inc(accountTrieLeaves[i])
+		storageTrieLeavesAtDepth[i].Inc(storageTrieLeaves[i])
 	}
 }