core/state: fix trie prefetcher for verkle

core/state/statedb.go

@@ -882,8 +882,10 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 			return nil
 		})
 	}
-	// If witness building is enabled, gather all the read-only accesses
-	if s.witness != nil {
+	// If witness building is enabled, gather all the read-only accesses.
+	// Skip witness collection in Verkle mode, they will be gathered
+	// together at the end.
+	if s.witness != nil && !s.db.TrieDB().IsVerkle() {
 		// Pull in anything that has been accessed before destruction
 		for _, obj := range s.stateObjectsDestruct {
 			// Skip any objects that haven't touched their storage
@@ -924,7 +926,7 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 	// only a single trie is used for state hashing. Replacing a non-nil verkle tree
 	// here could result in losing uncommitted changes from storage.
 	start = time.Now()
-	if s.prefetcher != nil && (s.trie == nil || !s.trie.IsVerkle()) {
+	if s.prefetcher != nil {
 		if trie := s.prefetcher.trie(common.Hash{}, s.originalRoot); trie == nil {
 			log.Error("Failed to retrieve account pre-fetcher trie")
 		} else {

core/state/trie_prefetcher.go

@@ -40,6 +40,7 @@ var (
 //
 // Note, the prefetcher's API is not thread safe.
 type triePrefetcher struct {
+	verkle   bool                   // Flag whether the prefetcher is in verkle mode
 	db       Database               // Database to fetch trie nodes through
 	root     common.Hash            // Root hash of the account trie for metrics
 	fetchers map[string]*subfetcher // Subfetchers for each trie
@@ -66,6 +67,7 @@ type triePrefetcher struct {
 func newTriePrefetcher(db Database, root common.Hash, namespace string, noreads bool) *triePrefetcher {
 	prefix := triePrefetchMetricsPrefix + namespace
 	return &triePrefetcher{
+		verkle:   db.TrieDB().IsVerkle(),
 		db:       db,
 		root:     root,
 		fetchers: make(map[string]*subfetcher), // Active prefetchers use the fetchers map
@@ -196,12 +198,18 @@ func (p *triePrefetcher) trie(owner common.Hash, root common.Hash) Trie {
 func (p *triePrefetcher) used(owner common.Hash, root common.Hash, used [][]byte) {
 	if fetcher := p.fetchers[p.trieID(owner, root)]; fetcher != nil {
 		fetcher.wait() // ensure the fetcher's idle before poking in its internals
-		fetcher.used = used
+		fetcher.used = append(fetcher.used, used...)
 	}
 }
 
 // trieID returns an unique trie identifier consists the trie owner and root hash.
 func (p *triePrefetcher) trieID(owner common.Hash, root common.Hash) string {
+	// The trie in verkle is only identified by state root
+	if p.verkle {
+		return p.root.Hex()
+	}
+	// The trie in merkle is either identified by state root (account trie),
+	// or identified by the owner and trie root (storage trie)
 	trieID := make([]byte, common.HashLength*2)
 	copy(trieID, owner.Bytes())
 	copy(trieID[common.HashLength:], root.Bytes())
@@ -320,29 +328,46 @@ func (sf *subfetcher) terminate(async bool) {
 	<-sf.term
 }
 
+// openTrie resolves the target trie from database for prefetching.
+func (sf *subfetcher) openTrie() error {
+	// Open the verkle tree if the sub-fetcher is in verkle mode
+	if sf.db.TrieDB().IsVerkle() {
+		tr, err := sf.db.OpenTrie(sf.state)
+		if err != nil {
+			log.Warn("Trie prefetcher failed opening verkle trie", "root", sf.root, "err", err)
+			return err
+		}
+		sf.trie = tr
+		return nil
+	}
+	// Open the merkle tree if the sub-fetcher is in merkle mode
+	if sf.owner == (common.Hash{}) {
+		tr, err := sf.db.OpenTrie(sf.state)
+		if err != nil {
+			log.Warn("Trie prefetcher failed opening account trie", "root", sf.root, "err", err)
+			return err
+		}
+		sf.trie = tr
+		return nil
+	}
+	tr, err := sf.db.OpenStorageTrie(sf.state, sf.addr, sf.root, nil)
+	if err != nil {
+		log.Warn("Trie prefetcher failed opening storage trie", "root", sf.root, "err", err)
+		return err
+	}
+	sf.trie = tr
+	return nil
+}
+
 // loop loads newly-scheduled trie tasks as they are received and loads them, stopping
 // when requested.
 func (sf *subfetcher) loop() {
 	// No matter how the loop stops, signal anyone waiting that it's terminated
 	defer close(sf.term)
 
-	// Start by opening the trie and stop processing if it fails
-	if sf.owner == (common.Hash{}) {
-		trie, err := sf.db.OpenTrie(sf.root)
-		if err != nil {
-			log.Warn("Trie prefetcher failed opening trie", "root", sf.root, "err", err)
-			return
-		}
-		sf.trie = trie
-	} else {
-		trie, err := sf.db.OpenStorageTrie(sf.state, sf.addr, sf.root, nil)
-		if err != nil {
-			log.Warn("Trie prefetcher failed opening trie", "root", sf.root, "err", err)
-			return
-		}
-		sf.trie = trie
+	if err := sf.openTrie(); err != nil {
+		return
 	}
-	// Trie opened successfully, keep prefetching items
 	for {
 		select {
 		case <-sf.wake:

core/state/trie_prefetcher_test.go

@@ -24,6 +24,9 @@ import (
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/tracing"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/internal/testrand"
+	"github.com/ethereum/go-ethereum/triedb"
 	"github.com/holiman/uint256"
 )
 
@@ -62,3 +65,44 @@ func TestUseAfterTerminate(t *testing.T) {
 		t.Errorf("Prefetcher returned nil trie after terminate")
 	}
 }
+
+func TestVerklePrefetcher(t *testing.T) {
+	db := NewDatabaseWithConfig(rawdb.NewMemoryDatabase(), triedb.VerkleDefaults)
+	state, err := New(types.EmptyRootHash, db, nil)
+	if err != nil {
+		t.Fatalf("failed to initialize state: %v", err)
+	}
+	// Create an account and check if the retrieved balance is correct
+	addr := testrand.Address()
+	skey := testrand.Hash()
+	sval := testrand.Hash()
+
+	state.SetBalance(addr, uint256.NewInt(42), tracing.BalanceChangeUnspecified) // Change the account trie
+	state.SetCode(addr, []byte("hello"))                                         // Change an external metadata
+	state.SetState(addr, skey, sval)                                             // Change the storage trie
+	root, _ := state.Commit(0, true)
+
+	state, _ = New(root, db, nil)
+	sRoot := state.GetStorageRoot(addr)
+	fetcher := newTriePrefetcher(db, root, "", false)
+
+	// Read account
+	fetcher.prefetch(common.Hash{}, root, common.Address{}, [][]byte{
+		addr.Bytes(),
+	}, false)
+
+	// Read storage slot
+	fetcher.prefetch(crypto.Keccak256Hash(addr.Bytes()), sRoot, addr, [][]byte{
+		skey.Bytes(),
+	}, false)
+
+	fetcher.terminate(false)
+	accountTrie := fetcher.trie(common.Hash{}, root)
+	storageTrie := fetcher.trie(crypto.Keccak256Hash(addr.Bytes()), sRoot)
+
+	rootA := accountTrie.Hash()
+	rootB := storageTrie.Hash()
+	if rootA != rootB {
+		t.Fatal("Two different tries are retrieved")
+	}
+}