From 6f4f78b568420ed051c169c42279bfc594276c73 Mon Sep 17 00:00:00 2001
From: Po <chenjunfeng@ethstorage.io>
Date: Thu, 26 Jun 2025 11:45:07 +0200
Subject: [PATCH] use flatReader

---
 core/state/reader.go            | 45 ++++++++++++++++++---------------
 core/state/statedb.go           | 10 ++------
 core/state/statedb_readerbal.go |  3 +--
 3 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/core/state/reader.go b/core/state/reader.go
index 7db4ee1452..4797bdf7a8 100644
--- a/core/state/reader.go
+++ b/core/state/reader.go
@@ -73,7 +73,7 @@ type StateReader interface {
 	Storage(addr common.Address, slot common.Hash) (common.Hash, error)
 
 	AccountBAL(addr common.Address) (*types.StateAccount, error)
-	StorageBAL(addr common.Address, slot common.Hash, tr *trie.StateTrie) (common.Hash, error)
+	StorageBAL(addr common.Address, slot common.Hash) (common.Hash, error)
 }
 
 // Reader defines the interface for accessing accounts, storage slots and contract
@@ -207,7 +207,7 @@ func (r *flatReader) Storage(addr common.Address, key common.Hash) (common.Hash,
 	return value, nil
 }
 
-func (r *flatReader) StorageBAL(addr common.Address, key common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+func (r *flatReader) StorageBAL(addr common.Address, key common.Hash) (common.Hash, error) {
 	return r.Storage(addr, key)
 }
 
@@ -333,15 +333,8 @@ func (r *trieReader) Storage(addr common.Address, key common.Hash) (common.Hash,
 	return value, nil
 }
 
-func (r *trieReader) StorageBAL(addr common.Address, key common.Hash, tr *trie.StateTrie) (common.Hash, error) {
-	ret, err := tr.GetStorage(addr, key.Bytes())
-	if err != nil {
-		return common.Hash{}, err
-	}
-
-	var value common.Hash
-	value.SetBytes(ret)
-	return value, nil
+func (r *trieReader) StorageBAL(addr common.Address, key common.Hash) (common.Hash, error) {
+	return r.Storage(addr, key)
 }
 
 // multiStateReader is the aggregation of a list of StateReader interface,
@@ -387,12 +380,18 @@ func (r *multiStateReader) Account(addr common.Address) (*types.StateAccount, er
 func (r *multiStateReader) AccountBAL(addr common.Address) (*types.StateAccount, error) {
 	var errs []error
 	for _, reader := range r.readers {
-		acct, err := reader.AccountBAL(addr)
-		if err == nil {
-			return acct, nil
+		switch reader.(type) {
+		case *flatReader:
+			{
+				acct, err := reader.AccountBAL(addr)
+				if err == nil {
+					return acct, nil
+				}
+				errs = append(errs, err)
+			}
 		}
-		errs = append(errs, err)
 	}
+	errs = append(errs, errors.New("no flatReader"))
 	return nil, errors.Join(errs...)
 }
 
@@ -414,15 +413,21 @@ func (r *multiStateReader) Storage(addr common.Address, slot common.Hash) (commo
 	return common.Hash{}, errors.Join(errs...)
 }
 
-func (r *multiStateReader) StorageBAL(addr common.Address, slot common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+func (r *multiStateReader) StorageBAL(addr common.Address, slot common.Hash) (common.Hash, error) {
 	var errs []error
 	for _, reader := range r.readers {
-		slot, err := reader.StorageBAL(addr, slot, tr)
-		if err == nil {
-			return slot, nil
+		switch reader.(type) {
+		case *flatReader:
+			{
+				slot, err := reader.Storage(addr, slot)
+				if err == nil {
+					return slot, nil
+				}
+				errs = append(errs, err)
+			}
 		}
-		errs = append(errs, err)
 	}
+	errs = append(errs, errors.New("no flatReader"))
 	return common.Hash{}, errors.Join(errs...)
 }
 
diff --git a/core/state/statedb.go b/core/state/statedb.go
index 3f2be3814b..d5b5abb13e 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -337,6 +337,7 @@ func (s *StateDB) SetPostBAL(postBal map[int]types.TxPostValues, bn uint64) erro
 	}
 	s.blockNumber = bn
 	s.postBAL = postBal
+	s.prefetcher = nil
 	return nil
 }
 
@@ -416,16 +417,9 @@ func (s *StateDB) prefetchBalPreblockKeys() {
 
 		lenSlots += len(acl.StorageKeys)
 
-		tr, err := trie.NewStateTrie(trie.StorageTrieID(s.originalRoot, obj.addrHash, obj.origin.Root), s.db.TrieDB())
-		if err != nil {
-			fmt.Println(s.originalRoot, obj.addrHash, obj.origin.Root)
-			log.Error("fail to create trie for", "addr:", addr, "error:", err)
-			panic("fail to create trie")
-		}
-
 		for _, key := range keys {
 			workers.Go(func() error {
-				val, err := s.reader.StorageBAL(addr, key, tr)
+				val, err := s.reader.StorageBAL(addr, key)
 				kv := &StorageKV{&addr, &key, &val}
 				if err != nil {
 					log.Error("fail to fetch storage from BALs:", addr, key)
diff --git a/core/state/statedb_readerbal.go b/core/state/statedb_readerbal.go
index 160f15dc07..d35d576fc4 100644
--- a/core/state/statedb_readerbal.go
+++ b/core/state/statedb_readerbal.go
@@ -5,7 +5,6 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
-	"github.com/ethereum/go-ethereum/trie"
 )
 
 // All states must has been cachae to statedb.
@@ -40,6 +39,6 @@ func (s *StateDB) AccountBAL(addr common.Address) (*types.StateAccount, error) {
 	panic("AccountBAL not implemented for statedb")
 }
 
-func (s *StateDB) StorageBAL(addr common.Address, slot common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+func (s *StateDB) StorageBAL(addr common.Address, slot common.Hash) (common.Hash, error) {
 	panic("StorageBAL not implemented for statedb")
 }