core, consensus/beacon: defer trie resolution (#31725)

Previously, the account trie for a given state root was resolved immediately 
when the stateDB was created, implying that the trie was always required
by the stateDB.

However, this assumption no longer holds, especially for path archive nodes, 
where historical states can be accessed even if the corresponding trie data 
does not exist.

Author: rjl493456442

consensus/beacon/consensus.go

@@ -391,8 +391,12 @@ func (beacon *Beacon) FinalizeAndAssemble(chain consensus.ChainHeaderReader, hea
 		if err != nil {
 			return nil, fmt.Errorf("error opening pre-state tree root: %w", err)
 		}
+		postTrie := state.GetTrie()
+		if postTrie == nil {
+			return nil, errors.New("post-state tree is not available")
+		}
 		vktPreTrie, okpre := preTrie.(*trie.VerkleTrie)
-		vktPostTrie, okpost := state.GetTrie().(*trie.VerkleTrie)
+		vktPostTrie, okpost := postTrie.(*trie.VerkleTrie)
 
 		// The witness is only attached iff both parent and current block are
 		// using verkle tree.

core/chain_makers.go

@@ -124,7 +124,7 @@ func (b *BlockGen) addTx(bc *BlockChain, vmConfig vm.Config, tx *types.Transacti
 	}
 	// Merge the tx-local access event into the "block-local" one, in order to collect
 	// all values, so that the witness can be built.
-	if b.statedb.GetTrie().IsVerkle() {
+	if b.statedb.Database().TrieDB().IsVerkle() {
 		b.statedb.AccessEvents().Merge(evm.AccessEvents)
 	}
 	b.txs = append(b.txs, tx)

core/state/dump.go

@@ -112,6 +112,9 @@ func (d iterativeDump) OnRoot(root common.Hash) {
 
 // DumpToCollector iterates the state according to the given options and inserts
 // the items into a collector for aggregation or serialization.
+//
+// The state iterator is still trie-based and can be converted to snapshot-based
+// once the state snapshot is fully integrated into database. TODO(rjl493456442).
 func (s *StateDB) DumpToCollector(c DumpCollector, conf *DumpConfig) (nextKey []byte) {
 	// Sanitize the input to allow nil configs
 	if conf == nil {
@@ -123,15 +126,20 @@ func (s *StateDB) DumpToCollector(c DumpCollector, conf *DumpConfig) (nextKey []
 		start            = time.Now()
 		logged           = time.Now()
 	)
-	log.Info("Trie dumping started", "root", s.trie.Hash())
-	c.OnRoot(s.trie.Hash())
+	log.Info("Trie dumping started", "root", s.originalRoot)
+	c.OnRoot(s.originalRoot)
 
-	trieIt, err := s.trie.NodeIterator(conf.Start)
+	tr, err := s.db.OpenTrie(s.originalRoot)
+	if err != nil {
+		return nil
+	}
+	trieIt, err := tr.NodeIterator(conf.Start)
 	if err != nil {
 		log.Error("Trie dumping error", "err", err)
 		return nil
 	}
 	it := trie.NewIterator(trieIt)
+
 	for it.Next() {
 		var data types.StateAccount
 		if err := rlp.DecodeBytes(it.Value, &data); err != nil {
@@ -147,7 +155,7 @@ func (s *StateDB) DumpToCollector(c DumpCollector, conf *DumpConfig) (nextKey []
 			}
 			address   *common.Address
 			addr      common.Address
-			addrBytes = s.trie.GetKey(it.Key)
+			addrBytes = tr.GetKey(it.Key)
 		)
 		if addrBytes == nil {
 			missingPreimages++
@@ -165,12 +173,13 @@ func (s *StateDB) DumpToCollector(c DumpCollector, conf *DumpConfig) (nextKey []
 		}
 		if !conf.SkipStorage {
 			account.Storage = make(map[common.Hash]string)
-			tr, err := obj.getTrie()
+
+			storageTr, err := s.db.OpenStorageTrie(s.originalRoot, addr, obj.Root(), tr)
 			if err != nil {
 				log.Error("Failed to load storage trie", "err", err)
 				continue
 			}
-			trieIt, err := tr.NodeIterator(nil)
+			trieIt, err := storageTr.NodeIterator(nil)
 			if err != nil {
 				log.Error("Failed to create trie iterator", "err", err)
 				continue
@@ -182,7 +191,7 @@ func (s *StateDB) DumpToCollector(c DumpCollector, conf *DumpConfig) (nextKey []
 					log.Error("Failed to decode the value returned by iterator", "error", err)
 					continue
 				}
-				key := s.trie.GetKey(storageIt.Key)
+				key := storageTr.GetKey(storageIt.Key)
 				if key == nil {
 					continue
 				}

core/state/iterator.go

@@ -32,6 +32,7 @@ import (
 // required in order to resolve the contract address.
 type nodeIterator struct {
 	state *StateDB // State being iterated
+	tr    Trie     // Primary account trie for traversal
 
 	stateIt trie.NodeIterator // Primary iterator for the global state trie
 	dataIt  trie.NodeIterator // Secondary iterator for the data trie of a contract
@@ -75,13 +76,20 @@ func (it *nodeIterator) step() error {
 	if it.state == nil {
 		return nil
 	}
+	if it.tr == nil {
+		tr, err := it.state.db.OpenTrie(it.state.originalRoot)
+		if err != nil {
+			return err
+		}
+		it.tr = tr
+	}
 	// Initialize the iterator if we've just started
-	var err error
 	if it.stateIt == nil {
-		it.stateIt, err = it.state.trie.NodeIterator(nil)
+		stateIt, err := it.tr.NodeIterator(nil)
 		if err != nil {
 			return err
 		}
+		it.stateIt = stateIt
 	}
 	// If we had data nodes previously, we surely have at least state nodes
 	if it.dataIt != nil {
@@ -116,14 +124,14 @@ func (it *nodeIterator) step() error {
 		return err
 	}
 	// Lookup the preimage of account hash
-	preimage := it.state.trie.GetKey(it.stateIt.LeafKey())
+	preimage := it.tr.GetKey(it.stateIt.LeafKey())
 	if preimage == nil {
 		return errors.New("account address is not available")
 	}
 	address := common.BytesToAddress(preimage)
 
 	// Traverse the storage slots belong to the account
-	dataTrie, err := it.state.db.OpenStorageTrie(it.state.originalRoot, address, account.Root, it.state.trie)
+	dataTrie, err := it.state.db.OpenStorageTrie(it.state.originalRoot, address, account.Root, it.tr)
 	if err != nil {
 		return err
 	}

core/state/state_object.go

@@ -124,6 +124,7 @@ func (s *stateObject) touch() {
 // subsequent reads to expand the same trie instead of reloading from disk.
 func (s *stateObject) getTrie() (Trie, error) {
 	if s.trie == nil {
+		// Assumes the primary account trie is already loaded
 		tr, err := s.db.db.OpenStorageTrie(s.db.originalRoot, s.address, s.data.Root, s.db.trie)
 		if err != nil {
 			return nil, err

core/state/state_test.go

@@ -54,8 +54,6 @@ func TestDump(t *testing.T) {
 	obj3.SetBalance(uint256.NewInt(44))
 
 	// write some of them to the trie
-	s.state.updateStateObject(obj1)
-	s.state.updateStateObject(obj2)
 	root, _ := s.state.Commit(0, false, false)
 
 	// check that DumpToCollector contains the state objects that are in trie
@@ -114,8 +112,6 @@ func TestIterativeDump(t *testing.T) {
 	obj4.AddBalance(uint256.NewInt(1337))
 
 	// write some of them to the trie
-	s.state.updateStateObject(obj1)
-	s.state.updateStateObject(obj2)
 	root, _ := s.state.Commit(0, false, false)
 	s.state, _ = New(root, tdb)
 

core/state/statedb.go

@@ -79,8 +79,8 @@ func (m *mutation) isDelete() bool {
 type StateDB struct {
 	db         Database
 	prefetcher *triePrefetcher
-	trie       Trie
 	reader     Reader
+	trie       Trie // it's resolved on first access
 
 	// originalRoot is the pre-state root, before any changes were made.
 	// It will be updated when the Commit is called.
@@ -169,13 +169,8 @@ func New(root common.Hash, db Database) (*StateDB, error) {
 // NewWithReader creates a new state for the specified state root. Unlike New,
 // this function accepts an additional Reader which is bound to the given root.
 func NewWithReader(root common.Hash, db Database, reader Reader) (*StateDB, error) {
-	tr, err := db.OpenTrie(root)
-	if err != nil {
-		return nil, err
-	}
 	sdb := &StateDB{
 		db:                   db,
-		trie:                 tr,
 		originalRoot:         root,
 		reader:               reader,
 		stateObjects:         make(map[common.Address]*stateObject),
@@ -664,7 +659,6 @@ func (s *StateDB) Copy() *StateDB {
 	// Copy all the basic fields, initialize the memory ones
 	state := &StateDB{
 		db:                   s.db,
-		trie:                 mustCopyTrie(s.trie),
 		reader:               s.reader,
 		originalRoot:         s.originalRoot,
 		stateObjects:         make(map[common.Address]*stateObject, len(s.stateObjects)),
@@ -688,6 +682,9 @@ func (s *StateDB) Copy() *StateDB {
 		transientStorage: s.transientStorage.Copy(),
 		journal:          s.journal.copy(),
 	}
+	if s.trie != nil {
+		state.trie = mustCopyTrie(s.trie)
+	}
 	if s.witness != nil {
 		state.witness = s.witness.Copy()
 	}
@@ -783,6 +780,20 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 	// Finalise all the dirty storage states and write them into the tries
 	s.Finalise(deleteEmptyObjects)
 
+	// Initialize the trie if it's not constructed yet. If the prefetch
+	// is enabled, the trie constructed below will be replaced by the
+	// prefetched one.
+	//
+	// This operation must be done before state object storage hashing,
+	// as it assumes the main trie is already loaded.
+	if s.trie == nil {
+		tr, err := s.db.OpenTrie(s.originalRoot)
+		if err != nil {
+			s.setError(err)
+			return common.Hash{}
+		}
+		s.trie = tr
+	}
 	// If there was a trie prefetcher operating, terminate it async so that the
 	// individual storage tries can be updated as soon as the disk load finishes.
 	if s.prefetcher != nil {

core/state/statedb_test.go

@@ -171,7 +171,6 @@ func TestCopy(t *testing.T) {
 	for i := byte(0); i < 255; i++ {
 		obj := orig.getOrNewStateObject(common.BytesToAddress([]byte{i}))
 		obj.AddBalance(uint256.NewInt(uint64(i)))
-		orig.updateStateObject(obj)
 	}
 	orig.Finalise(false)
 
@@ -190,10 +189,6 @@ func TestCopy(t *testing.T) {
 		origObj.AddBalance(uint256.NewInt(2 * uint64(i)))
 		copyObj.AddBalance(uint256.NewInt(3 * uint64(i)))
 		ccopyObj.AddBalance(uint256.NewInt(4 * uint64(i)))
-
-		orig.updateStateObject(origObj)
-		copy.updateStateObject(copyObj)
-		ccopy.updateStateObject(copyObj)
 	}
 
 	// Finalise the changes on all concurrently
@@ -238,7 +233,6 @@ func TestCopyWithDirtyJournal(t *testing.T) {
 		obj := orig.getOrNewStateObject(common.BytesToAddress([]byte{i}))
 		obj.AddBalance(uint256.NewInt(uint64(i)))
 		obj.data.Root = common.HexToHash("0xdeadbeef")
-		orig.updateStateObject(obj)
 	}
 	root, _ := orig.Commit(0, true, false)
 	orig, _ = New(root, db)
@@ -248,8 +242,6 @@ func TestCopyWithDirtyJournal(t *testing.T) {
 		obj := orig.getOrNewStateObject(common.BytesToAddress([]byte{i}))
 		amount := uint256.NewInt(uint64(i))
 		obj.SetBalance(new(uint256.Int).Sub(obj.Balance(), amount))
-
-		orig.updateStateObject(obj)
 	}
 	cpy := orig.Copy()
 
@@ -284,7 +276,6 @@ func TestCopyObjectState(t *testing.T) {
 		obj := orig.getOrNewStateObject(common.BytesToAddress([]byte{i}))
 		obj.AddBalance(uint256.NewInt(uint64(i)))
 		obj.data.Root = common.HexToHash("0xdeadbeef")
-		orig.updateStateObject(obj)
 	}
 	orig.Finalise(true)
 	cpy := orig.Copy()
@@ -573,7 +564,7 @@ func forEachStorage(s *StateDB, addr common.Address, cb func(key, value common.H
 	)
 
 	for it.Next() {
-		key := common.BytesToHash(s.trie.GetKey(it.Key))
+		key := common.BytesToHash(tr.GetKey(it.Key))
 		visited[key] = true
 		if value, dirty := so.dirtyStorage[key]; dirty {
 			if !cb(key, value) {

core/state_processor.go

@@ -161,10 +161,9 @@ func ApplyTransactionWithEVM(msg *Message, gp *GasPool, statedb *state.StateDB,
 
 	// Merge the tx-local access event into the "block-local" one, in order to collect
 	// all values, so that the witness can be built.
-	if statedb.GetTrie().IsVerkle() {
+	if statedb.Database().TrieDB().IsVerkle() {
 		statedb.AccessEvents().Merge(evm.AccessEvents)
 	}
-
 	return MakeReceipt(evm, result, statedb, blockNumber, blockHash, blockTime, tx, *usedGas, root), nil
 }