core, eth, trie, light: clean up trie interface (#26388)

* all: cleanup trie interface

* eth, trie: address comments

core/state/database.go

@@ -73,24 +73,31 @@ type Trie interface {
 	// trie.MissingNodeError is returned.
 	TryGet(key []byte) ([]byte, error)
 
-	// TryGetAccount abstract an account read from the trie.
-	TryGetAccount(key []byte) (*types.StateAccount, error)
+	// TryGetAccount abstracts an account read from the trie. It retrieves the
+	// account blob from the trie with provided account address and decodes it
+	// with associated decoding algorithm. If the specified account is not in
+	// the trie, nil will be returned. If the trie is corrupted(e.g. some nodes
+	// are missing or the account blob is incorrect for decoding), an error will
+	// be returned.
+	TryGetAccount(address common.Address) (*types.StateAccount, error)
 
 	// TryUpdate associates key with value in the trie. If value has length zero, any
 	// existing value is deleted from the trie. The value bytes must not be modified
 	// by the caller while they are stored in the trie. If a node was not found in the
 	// database, a trie.MissingNodeError is returned.
 	TryUpdate(key, value []byte) error
 
-	// TryUpdateAccount abstract an account write to the trie.
-	TryUpdateAccount(key []byte, account *types.StateAccount) error
+	// TryUpdateAccount abstracts an account write to the trie. It encodes the
+	// provided account object with associated algorithm and then updates it
+	// in the trie with provided address.
+	TryUpdateAccount(address common.Address, account *types.StateAccount) error
 
 	// TryDelete removes any existing value for key from the trie. If a node was not
 	// found in the database, a trie.MissingNodeError is returned.
 	TryDelete(key []byte) error
 
 	// TryDeleteAccount abstracts an account deletion from the trie.
-	TryDeleteAccount(key []byte) error
+	TryDeleteAccount(address common.Address) error
 
 	// Hash returns the root hash of the trie. It does not write to the database and
 	// can be used even if the trie doesn't have one.

core/state/statedb.go

@@ -510,7 +510,7 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 	}
 	// Encode the account and update the account trie
 	addr := obj.Address()
-	if err := s.trie.TryUpdateAccount(addr[:], &obj.data); err != nil {
+	if err := s.trie.TryUpdateAccount(addr, &obj.data); err != nil {
 		s.setError(fmt.Errorf("updateStateObject (%x) error: %v", addr[:], err))
 	}
 
@@ -531,7 +531,7 @@ func (s *StateDB) deleteStateObject(obj *stateObject) {
 	}
 	// Delete the account from the trie
 	addr := obj.Address()
-	if err := s.trie.TryDeleteAccount(addr[:]); err != nil {
+	if err := s.trie.TryDeleteAccount(addr); err != nil {
 		s.setError(fmt.Errorf("deleteStateObject (%x) error: %v", addr[:], err))
 	}
 }
@@ -585,7 +585,7 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject {
 	if data == nil {
 		start := time.Now()
 		var err error
-		data, err = s.trie.TryGetAccount(addr.Bytes())
+		data, err = s.trie.TryGetAccount(addr)
 		if metrics.EnabledExpensive {
 			s.AccountReads += time.Since(start)
 		}

core/state/trie_prefetcher.go

@@ -336,11 +336,7 @@ func (sf *subfetcher) loop() {
 					if _, ok := sf.seen[string(task)]; ok {
 						sf.dups++
 					} else {
-						if len(task) == len(common.Address{}) {
-							sf.trie.TryGetAccount(task)
-						} else {
-							sf.trie.TryGet(task)
-						}
+						sf.trie.TryGet(task)
 						sf.seen[string(task)] = struct{}{}
 					}
 				}

eth/protocols/snap/handler.go

@@ -417,7 +417,7 @@ func ServiceGetStorageRangesQuery(chain *core.BlockChain, req *GetStorageRangesP
 			if err != nil {
 				return nil, nil
 			}
-			acc, err := accTrie.TryGetAccountWithPreHashedKey(account[:])
+			acc, err := accTrie.TryGetAccountByHash(account)
 			if err != nil || acc == nil {
 				return nil, nil
 			}
@@ -523,7 +523,7 @@ func ServiceGetTrieNodesQuery(chain *core.BlockChain, req *GetTrieNodesPacket, s
 			if snap == nil {
 				// We don't have the requested state snapshotted yet (or it is stale),
 				// but can look up the account via the trie instead.
-				account, err := accTrie.TryGetAccountWithPreHashedKey(pathset[0])
+				account, err := accTrie.TryGetAccountByHash(common.BytesToHash(pathset[0]))
 				loads += 8 // We don't know the exact cost of lookup, this is an estimate
 				if err != nil || account == nil {
 					break

eth/protocols/snap/sync.go

@@ -3083,7 +3083,7 @@ func (t *healRequestSort) Swap(i, j int) {
 func (t *healRequestSort) Merge() []TrieNodePathSet {
 	var result []TrieNodePathSet
 	for _, path := range t.syncPaths {
-		pathset := TrieNodePathSet([][]byte(path))
+		pathset := TrieNodePathSet(path)
 		if len(path) == 1 {
 			// It's an account reference.
 			result = append(result, pathset)

light/trie.go

@@ -115,9 +115,9 @@ func (t *odrTrie) TryGet(key []byte) ([]byte, error) {
 	return res, err
 }
 
-func (t *odrTrie) TryGetAccount(key []byte) (*types.StateAccount, error) {
-	key = crypto.Keccak256(key)
+func (t *odrTrie) TryGetAccount(address common.Address) (*types.StateAccount, error) {
 	var res types.StateAccount
+	key := crypto.Keccak256(address.Bytes())
 	err := t.do(key, func() (err error) {
 		value, err := t.trie.TryGet(key)
 		if err != nil {
@@ -131,8 +131,8 @@ func (t *odrTrie) TryGetAccount(key []byte) (*types.StateAccount, error) {
 	return &res, err
 }
 
-func (t *odrTrie) TryUpdateAccount(key []byte, acc *types.StateAccount) error {
-	key = crypto.Keccak256(key)
+func (t *odrTrie) TryUpdateAccount(address common.Address, acc *types.StateAccount) error {
+	key := crypto.Keccak256(address.Bytes())
 	value, err := rlp.EncodeToBytes(acc)
 	if err != nil {
 		return fmt.Errorf("decoding error in account update: %w", err)
@@ -157,8 +157,8 @@ func (t *odrTrie) TryDelete(key []byte) error {
 }
 
 // TryDeleteAccount abstracts an account deletion from the trie.
-func (t *odrTrie) TryDeleteAccount(key []byte) error {
-	key = crypto.Keccak256(key)
+func (t *odrTrie) TryDeleteAccount(address common.Address) error {
+	key := crypto.Keccak256(address.Bytes())
 	return t.do(key, func() error {
 		return t.trie.TryDelete(key)
 	})

trie/secure_trie.go

@@ -90,11 +90,11 @@ func (t *StateTrie) TryGet(key []byte) ([]byte, error) {
 	return t.trie.TryGet(t.hashKey(key))
 }
 
-// TryGetAccount attempts to retrieve an account with provided trie path.
+// TryGetAccount attempts to retrieve an account with provided account address.
 // If the specified account is not in the trie, nil will be returned.
 // If a trie node is not found in the database, a MissingNodeError is returned.
-func (t *StateTrie) TryGetAccount(key []byte) (*types.StateAccount, error) {
-	res, err := t.trie.TryGet(t.hashKey(key))
+func (t *StateTrie) TryGetAccount(address common.Address) (*types.StateAccount, error) {
+	res, err := t.trie.TryGet(t.hashKey(address.Bytes()))
 	if res == nil || err != nil {
 		return nil, err
 	}
@@ -103,11 +103,11 @@ func (t *StateTrie) TryGetAccount(key []byte) (*types.StateAccount, error) {
 	return ret, err
 }
 
-// TryGetAccountWithPreHashedKey does the same thing as TryGetAccount, however
-// it expects a key that is already hashed. This constitutes an abstraction leak,
-// since the client code needs to know the key format.
-func (t *StateTrie) TryGetAccountWithPreHashedKey(key []byte) (*types.StateAccount, error) {
-	res, err := t.trie.TryGet(key)
+// TryGetAccountByHash does the same thing as TryGetAccount, however
+// it expects an account hash that is the hash of address. This constitutes an
+// abstraction leak, since the client code needs to know the key format.
+func (t *StateTrie) TryGetAccountByHash(addrHash common.Hash) (*types.StateAccount, error) {
+	res, err := t.trie.TryGet(addrHash.Bytes())
 	if res == nil || err != nil {
 		return nil, err
 	}
@@ -156,16 +156,16 @@ func (t *StateTrie) TryUpdate(key, value []byte) error {
 
 // TryUpdateAccount account will abstract the write of an account to the
 // secure trie.
-func (t *StateTrie) TryUpdateAccount(key []byte, acc *types.StateAccount) error {
-	hk := t.hashKey(key)
+func (t *StateTrie) TryUpdateAccount(address common.Address, acc *types.StateAccount) error {
+	hk := t.hashKey(address.Bytes())
 	data, err := rlp.EncodeToBytes(acc)
 	if err != nil {
 		return err
 	}
 	if err := t.trie.TryUpdate(hk, data); err != nil {
 		return err
 	}
-	t.getSecKeyCache()[string(hk)] = common.CopyBytes(key)
+	t.getSecKeyCache()[string(hk)] = address.Bytes()
 	return nil
 }
 
@@ -186,8 +186,8 @@ func (t *StateTrie) TryDelete(key []byte) error {
 }
 
 // TryDeleteAccount abstracts an account deletion from the trie.
-func (t *StateTrie) TryDeleteAccount(key []byte) error {
-	hk := t.hashKey(key)
+func (t *StateTrie) TryDeleteAccount(address common.Address) error {
+	hk := t.hashKey(address.Bytes())
 	delete(t.getSecKeyCache(), string(hk))
 	return t.trie.TryDelete(hk)
 }