core/state: separate trie reader to mptReader and ubtReader (#34763)

This PR separates the trie reader to mptTrieReader and ubtTrieReader for
improved readability and extensibility.

---------

Co-authored-by: Guillaume Ballet <3272758+gballet@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

core/state/database_mpt.go

@@ -81,7 +81,7 @@ func (db *MPTDatabase) StateReader(stateRoot common.Hash) (StateReader, error) {
 	}
 	// Configure the trie reader, which is expected to be available as the
 	// gatekeeper unless the state is corrupted.
-	tr, err := newTrieReader(stateRoot, db.triedb)
+	tr, err := newMPTTrieReader(stateRoot, db.triedb)
 	if err != nil {
 		return nil, err
 	}

core/state/database_ubt.go

@@ -61,7 +61,7 @@ func (db *UBTDatabase) StateReader(stateRoot common.Hash) (StateReader, error) {
 	}
 	// Configure the trie reader, which is expected to be available as the
 	// gatekeeper unless the state is corrupted.
-	tr, err := newTrieReader(stateRoot, db.triedb)
+	tr, err := newUBTTrieReader(stateRoot, db.triedb)
 	if err != nil {
 		return nil, err
 	}

core/state/reader.go

@@ -148,70 +148,28 @@ func (r *flatReader) Storage(addr common.Address, key common.Hash) (common.Hash,
 	return value, nil
 }
 
-// trieReader implements the StateReader interface, providing functions to access
+// mptTrieReader implements the StateReader interface, providing functions to
-// state from the referenced trie.
+// access state from the referenced Merkle-Patricia-tree.
 //
-// trieReader is safe for concurrent read.
+// mptTrieReader is safe for concurrent read.
-type trieReader struct {
+type mptTrieReader struct {
-	root common.Hash      // State root which uniquely represent a state
+	root common.Hash      // State root which uniquely represents a state
 	db   *triedb.Database // Database for loading trie
 
-	// Main trie, resolved in constructor. Note either the Merkle-Patricia-tree
+	mainTrie Trie                           // Main trie, resolved in constructor, not thread-safe
-	// or Verkle-tree is not safe for concurrent read.
-	mainTrie Trie
-
 	subRoots map[common.Address]common.Hash // Set of storage roots, cached when the account is resolved
 	subTries map[common.Address]Trie        // Group of storage tries, cached when it's resolved
 	lock     sync.Mutex                     // Lock for protecting concurrent read
 }
 
-// newTrieReader constructs a trie reader of the specific state. An error will be
+// newMPTTrieReader constructs a Merkle-Patricia-tree reader of the specific state.
-// returned if the associated trie specified by root is not existent.
+// An error will be returned if the associated trie specified by root is not existent.
-func newTrieReader(root common.Hash, db *triedb.Database) (*trieReader, error) {
+func newMPTTrieReader(root common.Hash, db *triedb.Database) (*mptTrieReader, error) {
-	var (
+	tr, err := trie.NewStateTrie(trie.StateTrieID(root), db)
-		tr  Trie
-		err error
-	)
-	if !db.IsUBT() {
-		tr, err = trie.NewStateTrie(trie.StateTrieID(root), db)
-	} else {
-		// When IsUBT() is true, create a BinaryTrie wrapped in TransitionTrie
-		binTrie, binErr := bintrie.NewBinaryTrie(root, db)
-		if binErr != nil {
-			return nil, binErr
-		}
-
-		// Based on the transition status, determine if the overlay
-		// tree needs to be created, or if a single, target tree is
-		// to be picked.
-		ts := overlay.LoadTransitionState(db.Disk(), root, true)
-		if ts.InTransition() {
-			mpt, err := trie.NewStateTrie(trie.StateTrieID(ts.BaseRoot), db)
-			if err != nil {
-				return nil, err
-			}
-			tr = transitiontrie.NewTransitionTrie(mpt, binTrie, false)
-		} else {
-			// HACK: Use TransitionTrie with nil base as a wrapper to make BinaryTrie
-			// satisfy the Trie interface. This works around the import cycle between
-			// trie and trie/bintrie packages.
-			//
-			// TODO: In future PRs, refactor the package structure to avoid this hack:
-			// - Option 1: Move common interfaces (Trie, NodeIterator) to a separate
-			//   package that both trie and trie/bintrie can import
-			// - Option 2: Create a factory function in the trie package that returns
-			//   BinaryTrie as a Trie interface without direct import
-			// - Option 3: Move BinaryTrie to the main trie package
-			//
-			// The current approach works but adds unnecessary overhead and complexity
-			// by using TransitionTrie when there's no actual transition happening.
-			tr = transitiontrie.NewTransitionTrie(nil, binTrie, false)
-		}
-	}
 	if err != nil {
 		return nil, err
 	}
-	return &trieReader{
+	return &mptTrieReader{
 		root:     root,
 		db:       db,
 		mainTrie: tr,
@@ -221,7 +179,7 @@ func newTrieReader(root common.Hash, db *triedb.Database) (*trieReader, error) {
 }
 
 // account is the inner version of Account and assumes the r.lock is already held.
-func (r *trieReader) account(addr common.Address) (*types.StateAccount, error) {
+func (r *mptTrieReader) account(addr common.Address) (*types.StateAccount, error) {
 	account, err := r.mainTrie.GetAccount(addr)
 	if err != nil {
 		return nil, err
@@ -236,9 +194,9 @@ func (r *trieReader) account(addr common.Address) (*types.StateAccount, error) {
 
 // Account implements StateReader, retrieving the account specified by the address.
 //
-// An error will be returned if the trie state is corrupted. An nil account
+// An error will be returned if the trie state is corrupted. A nil account
 // will be returned if it's not existent in the trie.
-func (r *trieReader) Account(addr common.Address) (*types.StateAccount, error) {
+func (r *mptTrieReader) Account(addr common.Address) (*types.StateAccount, error) {
 	r.lock.Lock()
 	defer r.lock.Unlock()
 
@@ -250,43 +208,118 @@ func (r *trieReader) Account(addr common.Address) (*types.StateAccount, error) {
 //
 // An error will be returned if the trie state is corrupted. An empty storage
 // slot will be returned if it's not existent in the trie.
-func (r *trieReader) Storage(addr common.Address, key common.Hash) (common.Hash, error) {
+func (r *mptTrieReader) Storage(addr common.Address, key common.Hash) (common.Hash, error) {
 	r.lock.Lock()
 	defer r.lock.Unlock()
 
-	var (
+	tr, found := r.subTries[addr]
-		tr    Trie
+	if !found {
-		found bool
+		root, ok := r.subRoots[addr]
-		value common.Hash
-	)
-	if r.db.IsUBT() {
-		tr = r.mainTrie
-	} else {
-		tr, found = r.subTries[addr]
-		if !found {
-			root, ok := r.subRoots[addr]
 
-			// The storage slot is accessed without account caching. It's unexpected
+		// The storage slot is accessed without account caching. It's unexpected
-			// behavior but try to resolve the account first anyway.
+		// behavior but try to resolve the account first anyway.
-			if !ok {
+		if !ok {
-				_, err := r.account(addr)
+			_, err := r.account(addr)
-				if err != nil {
-					return common.Hash{}, err
-				}
-				root = r.subRoots[addr]
-			}
-			var err error
-			tr, err = trie.NewStateTrie(trie.StorageTrieID(r.root, crypto.Keccak256Hash(addr.Bytes()), root), r.db)
 			if err != nil {
 				return common.Hash{}, err
 			}
-			r.subTries[addr] = tr
+			root = r.subRoots[addr]
 		}
+		var err error
+		tr, err = trie.NewStateTrie(trie.StorageTrieID(r.root, crypto.Keccak256Hash(addr.Bytes()), root), r.db)
+		if err != nil {
+			return common.Hash{}, err
+		}
+		r.subTries[addr] = tr
 	}
 	ret, err := tr.GetStorage(addr, key.Bytes())
 	if err != nil {
 		return common.Hash{}, err
 	}
+	var value common.Hash
+	value.SetBytes(ret)
+	return value, nil
+}
+
+// ubtTrieReader implements the StateReader interface, providing functions to access
+// state from the referenced Unified-binary-trie.
+//
+// ubtTrieReader is safe for concurrent read.
+type ubtTrieReader struct {
+	root common.Hash      // State root which uniquely represents a state
+	db   *triedb.Database // Database for loading trie
+	tr   Trie             // Referenced unified binary trie
+	lock sync.Mutex       // Lock for protecting concurrent read
+}
+
+// newUBTTrieReader constructs a Unified-binary-trie reader of the specific state.
+// An error will be returned if the associated trie specified by root is not existent.
+func newUBTTrieReader(root common.Hash, db *triedb.Database) (*ubtTrieReader, error) {
+	binTrie, binErr := bintrie.NewBinaryTrie(root, db)
+	if binErr != nil {
+		return nil, binErr
+	}
+	// Based on the transition status, determine if the overlay
+	// tree needs to be created, or if a single, target tree is
+	// to be picked.
+	var (
+		tr Trie
+		ts = overlay.LoadTransitionState(db.Disk(), root, true)
+	)
+	if ts.InTransition() {
+		mpt, err := trie.NewStateTrie(trie.StateTrieID(ts.BaseRoot), db)
+		if err != nil {
+			return nil, err
+		}
+		tr = transitiontrie.NewTransitionTrie(mpt, binTrie, false)
+	} else {
+		// HACK: Use TransitionTrie with nil base as a wrapper to make BinaryTrie
+		// satisfy the Trie interface. This works around the import cycle between
+		// trie and trie/bintrie packages.
+		//
+		// TODO: In future PRs, refactor the package structure to avoid this hack:
+		// - Option 1: Move common interfaces (Trie, NodeIterator) to a separate
+		//   package that both trie and trie/bintrie can import
+		// - Option 2: Create a factory function in the trie package that returns
+		//   BinaryTrie as a Trie interface without direct import
+		// - Option 3: Move BinaryTrie to the main trie package
+		//
+		// The current approach works but adds unnecessary overhead and complexity
+		// by using TransitionTrie when there's no actual transition happening.
+		tr = transitiontrie.NewTransitionTrie(nil, binTrie, false)
+	}
+	return &ubtTrieReader{
+		root: root,
+		db:   db,
+		tr:   tr,
+	}, nil
+}
+
+// Account implements StateReader, retrieving the account specified by the address.
+//
+// An error will be returned if the trie state is corrupted. A nil account
+// will be returned if it's not existent in the trie.
+func (r *ubtTrieReader) Account(addr common.Address) (*types.StateAccount, error) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	return r.tr.GetAccount(addr)
+}
+
+// Storage implements StateReader, retrieving the storage slot specified by the
+// address and slot key.
+//
+// An error will be returned if the trie state is corrupted. An empty storage
+// slot will be returned if it's not existent in the trie.
+func (r *ubtTrieReader) Storage(addr common.Address, key common.Hash) (common.Hash, error) {
+	r.lock.Lock()
+	defer r.lock.Unlock()
+
+	ret, err := r.tr.GetStorage(addr, key.Bytes())
+	if err != nil {
+		return common.Hash{}, err
+	}
+	var value common.Hash
 	value.SetBytes(ret)
 	return value, nil
 }