cmd/geth,trie: add --contract flag to inspect trie

cmd/geth/dbcmd.go

@@ -68,6 +68,10 @@ var (
 		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,
@@ -118,6 +122,7 @@ Remove blockchain and state databases`,
 			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, 
@@ -488,6 +493,12 @@ func inspectTrie(ctx *cli.Context) error {
 	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)
 }

trie/inspect.go

@@ -31,13 +31,19 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
+	"time"
 
 	"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/internal/tablewriter"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/triedb/database"
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/semaphore"
 )
 
@@ -117,6 +123,157 @@ func Inspect(triedb database.NodeDatabase, root common.Hash, config *InspectConf
 	return Summarize(config.DumpPath, config)
 }
 
+// InspectContract inspects the on-disk footprint of a single contract.
+// It reports snapshot storage (slot count + size) and storage trie node
+// statistics (node type breakdown and per-depth distribution).
+func InspectContract(triedb database.NodeDatabase, db ethdb.Database, stateRoot common.Hash, address common.Address) error {
+	// Resolve account from the state trie.
+	accountHash := crypto.Keccak256Hash(address.Bytes())
+	accountTrie, err := New(TrieID(stateRoot), triedb)
+	if err != nil {
+		return fmt.Errorf("failed to open account trie: %w", err)
+	}
+	accountRLP, err := accountTrie.Get(crypto.Keccak256(address.Bytes()))
+	if err != nil {
+		return fmt.Errorf("failed to read account: %w", err)
+	}
+	if accountRLP == nil {
+		return fmt.Errorf("account not found: %s", address)
+	}
+	var account types.StateAccount
+	if err := rlp.DecodeBytes(accountRLP, &account); err != nil {
+		return fmt.Errorf("failed to decode account: %w", err)
+	}
+	if account.Root == (common.Hash{}) || account.Root == types.EmptyRootHash {
+		return fmt.Errorf("account %s has no storage", address)
+	}
+
+	// Look up account snapshot.
+	accountData := rawdb.ReadAccountSnapshot(db, accountHash)
+
+	// Run trie walk + snap iteration in parallel.
+	var (
+		snapSlots atomic.Uint64
+		snapSize  atomic.Uint64
+		g         errgroup.Group
+		start     = time.Now()
+	)
+
+	// Goroutine 1: Snapshot storage iteration.
+	g.Go(func() error {
+		prefix := append(rawdb.SnapshotStoragePrefix, accountHash.Bytes()...)
+		it := db.NewIterator(prefix, nil)
+		defer it.Release()
+
+		for it.Next() {
+			if !bytes.HasPrefix(it.Key(), prefix) {
+				break
+			}
+			snapSlots.Add(1)
+			snapSize.Add(uint64(len(it.Key()) + len(it.Value())))
+		}
+		return it.Error()
+	})
+
+	// Goroutine 2: Storage trie walk using the existing inspector.
+	var storageStat *LevelStats
+	g.Go(func() error {
+		owner := accountHash
+		storage, err := New(StorageTrieID(stateRoot, owner, account.Root), triedb)
+		if err != nil {
+			return fmt.Errorf("failed to open storage trie: %w", err)
+		}
+		storageStat = NewLevelStats()
+
+		in := &inspector{
+			triedb:      triedb,
+			root:        stateRoot,
+			config:      &InspectConfig{NoStorage: true},
+			accountStat: NewLevelStats(), // unused, but needed by inspector
+			sem:         semaphore.NewWeighted(inspectParallelism),
+			dumpBuf:     bufio.NewWriter(io.Discard),
+		}
+
+		in.recordRootSize(storage, account.Root, storageStat)
+		in.inspect(storage, storage.root, 0, []byte{}, storageStat)
+
+		if err := in.getError(); err != nil {
+			return err
+		}
+		return nil
+	})
+
+	// Progress reporter.
+	done := make(chan struct{})
+	go func() {
+		ticker := time.NewTicker(8 * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				log.Info("Inspecting contract",
+					"snapSlots", snapSlots.Load(),
+					"elapsed", common.PrettyDuration(time.Since(start)))
+			case <-done:
+				return
+			}
+		}
+	}()
+
+	if err := g.Wait(); err != nil {
+		close(done)
+		return err
+	}
+	close(done)
+
+	// Display results.
+	fmt.Printf("\n=== Contract Inspection: %s ===\n", address)
+	fmt.Printf("Account hash: %s\n\n", accountHash)
+
+	if len(accountData) == 0 {
+		fmt.Println("Account snapshot: not found")
+	} else {
+		fmt.Printf("Account snapshot: %s\n", common.StorageSize(len(accountData)))
+	}
+
+	fmt.Printf("Snapshot storage: %d slots (%s)\n",
+		snapSlots.Load(), common.StorageSize(snapSize.Load()))
+
+	// Compute trie totals from LevelStats.
+	var trieTotal, trieSize uint64
+	for i := 0; i < trieStatLevels; i++ {
+		short, full, value, size := storageStat.level[i].load()
+		trieTotal += short + full + value
+		trieSize += size
+	}
+	fmt.Printf("Storage trie:     %d nodes (%s)\n", trieTotal, common.StorageSize(trieSize))
+
+	// Depth distribution table with node type columns.
+	fmt.Println("\nStorage Trie Depth Distribution:")
+	b := new(strings.Builder)
+	table := tablewriter.NewWriter(b)
+	table.SetHeader([]string{"Depth", "Short", "Full", "Value", "Nodes", "Size"})
+	for i := 0; i < trieStatLevels; i++ {
+		short, full, value, size := storageStat.level[i].load()
+		total := short + full + value
+		if total == 0 && size == 0 {
+			continue
+		}
+		table.AppendBulk([][]string{{
+			fmt.Sprint(i),
+			fmt.Sprint(short),
+			fmt.Sprint(full),
+			fmt.Sprint(value),
+			fmt.Sprint(total),
+			common.StorageSize(size).String(),
+		}})
+	}
+	table.Render()
+	fmt.Print(b.String())
+
+	return nil
+}
+
 func normalizeInspectConfig(config *InspectConfig) *InspectConfig {
 	if config == nil {
 		config = &InspectConfig{}

trie/inspect_test.go

@@ -24,6 +24,7 @@ import (
 	"reflect"
 	"testing"
 
+	"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"
@@ -156,6 +157,59 @@ func assertAccountTotalsMatchLevels(t *testing.T, account storageStats) {
 	}
 }
 
+// TestInspectContract tests the InspectContract function on a single account
+// with storage and snapshot data.
+func TestInspectContract(t *testing.T) {
+	diskdb := rawdb.NewMemoryDatabase()
+	db := newTestDatabase(diskdb, rawdb.HashScheme)
+
+	// Create a contract address and its storage trie.
+	address := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678")
+	accountHash := crypto.Keccak256Hash(address.Bytes())
+
+	// Build a storage trie with some entries.
+	storageTrie := NewEmpty(db)
+	storageSlots := make(map[common.Hash][]byte)
+	for i := 0; i < 10; i++ {
+		k := crypto.Keccak256Hash([]byte{byte(i)})
+		v := []byte{byte(i + 1)}
+		storageTrie.MustUpdate(k.Bytes(), v)
+		storageSlots[k] = v
+	}
+	storageRoot, storageNodes := storageTrie.Commit(true)
+	db.Update(storageRoot, types.EmptyRootHash, trienode.NewWithNodeSet(storageNodes))
+	db.Commit(storageRoot)
+
+	// Build the account trie with the contract account.
+	account := types.StateAccount{
+		Nonce:    1,
+		Balance:  uint256.NewInt(1000),
+		Root:     storageRoot,
+		CodeHash: crypto.Keccak256(nil),
+	}
+	accountRLP, err := rlp.EncodeToBytes(&account)
+	if err != nil {
+		t.Fatalf("failed to encode account: %v", err)
+	}
+
+	accountTrie := NewEmpty(db)
+	accountTrie.MustUpdate(crypto.Keccak256(address.Bytes()), accountRLP)
+	stateRoot, accountNodes := accountTrie.Commit(true)
+	db.Update(stateRoot, types.EmptyRootHash, trienode.NewWithNodeSet(accountNodes))
+	db.Commit(stateRoot)
+
+	// Write snapshot data for the account and its storage slots.
+	rawdb.WriteAccountSnapshot(diskdb, accountHash, accountRLP)
+	for k, v := range storageSlots {
+		rawdb.WriteStorageSnapshot(diskdb, accountHash, k, v)
+	}
+
+	// InspectContract should succeed without error.
+	if err := InspectContract(db, diskdb, stateRoot, address); err != nil {
+		t.Fatalf("InspectContract failed: %v", err)
+	}
+}
+
 func makeAccountsWithStorage(db *testDb, size int, storage bool) (addresses [][20]byte, accounts [][]byte) {
 	// Make the random benchmark deterministic
 	random := rand.New(rand.NewSource(0))