core/state: wrap state reader to emit account load

core/state/database.go

@@ -193,7 +193,7 @@ func (db *CachingDB) Reader(stateRoot common.Hash) (Reader, error) {
 	}
 	// Set up the trie reader, which is expected to always be available
 	// as the gatekeeper unless the state is corrupted.
-	tr, err := newTrieReader(stateRoot, db.triedb, db.pointCache)
+	tr, err := newTrieReader(stateRoot, db.triedb, db, db.pointCache)
 	if err != nil {
 		return nil, err
 	}

core/state/reader.go

@@ -48,6 +48,19 @@ type Reader interface {
 	// - The returned storage slot is safe to modify after the call
 	Storage(addr common.Address, slot common.Hash) (common.Hash, error)
 
+	// Code returns the code associated with a particular account.
+	//
+	// - Returns an empty code if it does not exist
+	// - It can return an error to indicate code doesn't exist
+	// - The returned code is safe to modify after the call
+	Code(addr common.Address, codeHash common.Hash) ([]byte, error)
+
+	// CodeSize returns the size of the code associated with a particular account.
+	//
+	// - Returns 0 if the code does not exist
+	// - It can return an error to indicate code doesn't exist
+	CodeSize(addr common.Address, codeHash common.Hash) (int, error)
+
 	// Copy returns a deep-copied state reader.
 	Copy() Reader
 }
@@ -123,6 +136,16 @@ func (r *stateReader) Storage(addr common.Address, key common.Hash) (common.Hash
 	return value, nil
 }
 
+// Code implements Reader, retrieving the code associated with a particular account.
+func (r *stateReader) Code(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	return nil, nil
+}
+
+// CodeSize implements Reader, returning the size of the code associated with a particular account.
+func (r *stateReader) CodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	return 0, nil
+}
+
 // Copy implements Reader, returning a deep-copied snap reader.
 func (r *stateReader) Copy() Reader {
 	return &stateReader{
@@ -134,17 +157,18 @@ func (r *stateReader) Copy() Reader {
 // trieReader implements the Reader interface, providing functions to access
 // state from the referenced trie.
 type trieReader struct {
-	root     common.Hash                    // State root which uniquely represent a state
-	db       *triedb.Database               // Database for loading trie
-	buff     crypto.KeccakState             // Buffer for keccak256 hashing
-	mainTrie Trie                           // Main trie, resolved in constructor
-	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
+	root       common.Hash                    // State root which uniquely represent a state
+	db         *triedb.Database               // Database for loading trie
+	contractDB Database                       // Database for loading code
+	buff       crypto.KeccakState             // Buffer for keccak256 hashing
+	mainTrie   Trie                           // Main trie, resolved in constructor
+	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
 }
 
 // trieReader constructs a trie reader of the specific state. An error will be
 // returned if the associated trie specified by root is not existent.
-func newTrieReader(root common.Hash, db *triedb.Database, cache *utils.PointCache) (*trieReader, error) {
+func newTrieReader(root common.Hash, db *triedb.Database, contractDB Database, cache *utils.PointCache) (*trieReader, error) {
 	var (
 		tr  Trie
 		err error
@@ -158,12 +182,13 @@ func newTrieReader(root common.Hash, db *triedb.Database, cache *utils.PointCach
 		return nil, err
 	}
 	return &trieReader{
-		root:     root,
-		db:       db,
-		buff:     crypto.NewKeccakState(),
-		mainTrie: tr,
-		subRoots: make(map[common.Address]common.Hash),
-		subTries: make(map[common.Address]Trie),
+		root:       root,
+		db:         db,
+		contractDB: contractDB,
+		buff:       crypto.NewKeccakState(),
+		mainTrie:   tr,
+		subRoots:   make(map[common.Address]common.Hash),
+		subTries:   make(map[common.Address]Trie),
 	}, nil
 }
 
@@ -227,6 +252,16 @@ func (r *trieReader) Storage(addr common.Address, key common.Hash) (common.Hash,
 	return value, nil
 }
 
+// Code implements Reader, retrieving the code associated with a particular account.
+func (r *trieReader) Code(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	return r.contractDB.ContractCode(addr, codeHash)
+}
+
+// CodeSize implements Reader, returning the size of the code associated with a particular account.
+func (r *trieReader) CodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	return r.contractDB.ContractCodeSize(addr, codeHash)
+}
+
 // Copy implements Reader, returning a deep-copied trie reader.
 func (r *trieReader) Copy() Reader {
 	tries := make(map[common.Address]Trie)
@@ -298,6 +333,30 @@ func (r *multiReader) Storage(addr common.Address, slot common.Hash) (common.Has
 	return common.Hash{}, errors.Join(errs...)
 }
 
+// Code implements Reader, retrieving the code associated with a particular account.
+func (r *multiReader) Code(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	var errs []error
+	for _, reader := range r.readers {
+		code, err := reader.Code(addr, codeHash)
+		if err == nil {
+			return code, nil
+		}
+	}
+	return nil, errors.Join(errs...)
+}
+
+// CodeSize implements Reader, returning the size of the code associated with a particular account.
+func (r *multiReader) CodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	var errs []error
+	for _, reader := range r.readers {
+		size, err := reader.CodeSize(addr, codeHash)
+		if err == nil {
+			return size, nil
+		}
+	}
+	return 0, errors.Join(errs...)
+}
+
 // Copy implementing Reader interface, returning a deep-copied state reader.
 func (r *multiReader) Copy() Reader {
 	var readers []Reader

core/state/state_object.go

@@ -510,7 +510,7 @@ func (s *stateObject) Code() []byte {
 	if bytes.Equal(s.CodeHash(), types.EmptyCodeHash.Bytes()) {
 		return nil
 	}
-	code, err := s.db.db.ContractCode(s.address, common.BytesToHash(s.CodeHash()))
+	code, err := s.db.reader.Code(s.address, common.BytesToHash(s.CodeHash()))
 	if err != nil {
 		s.db.setError(fmt.Errorf("can't load code hash %x: %v", s.CodeHash(), err))
 	}
@@ -528,7 +528,7 @@ func (s *stateObject) CodeSize() int {
 	if bytes.Equal(s.CodeHash(), types.EmptyCodeHash.Bytes()) {
 		return 0
 	}
-	size, err := s.db.db.ContractCodeSize(s.address, common.BytesToHash(s.CodeHash()))
+	size, err := s.db.reader.CodeSize(s.address, common.BytesToHash(s.CodeHash()))
 	if err != nil {
 		s.db.setError(fmt.Errorf("can't load code size %x: %v", s.CodeHash(), err))
 	}

core/state/statedb_hooked.go

@@ -42,9 +42,17 @@ func NewHookedState(stateDb *StateDB, hooks *tracing.Hooks) *hookedStateDB {
 	if s.hooks == nil {
 		s.hooks = new(tracing.Hooks)
 	}
+	s.inner.reader = newHookedReader(s.inner.reader, s.hooks)
 	return s
 }
 
+// Close unwraps the inner reader.
+func (s *hookedStateDB) Close() {
+	if rd, ok := s.inner.reader.(*hookedReader); ok {
+		s.inner.reader = rd.inner
+	}
+}
+
 func (s *hookedStateDB) CreateAccount(addr common.Address) {
 	s.inner.CreateAccount(addr)
 }
@@ -288,3 +296,65 @@ func (s *hookedStateDB) Finalise(deleteEmptyObjects bool) {
 		}
 	}
 }
+
+// hookedReader wraps a Reader and invokes hooks when accounts and storage are loaded
+type hookedReader struct {
+	inner Reader
+	hooks *tracing.Hooks
+}
+
+// newHookedReader creates a new hookedReader that wraps the given reader
+func newHookedReader(reader Reader, hooks *tracing.Hooks) *hookedReader {
+	return &hookedReader{
+		inner: reader,
+		hooks: hooks,
+	}
+}
+
+// Account implements Reader, retrieving the account and invoking OnAccountLoad hook
+func (r *hookedReader) Account(addr common.Address) (*types.StateAccount, error) {
+	acct, err := r.inner.Account(addr)
+	if err == nil && r.hooks.OnAccountLoad != nil {
+		r.hooks.OnAccountLoad(addr, acct)
+	}
+	return acct, err
+}
+
+// Storage implements Reader, retrieving storage and invoking OnStorageLoad hook
+func (r *hookedReader) Storage(addr common.Address, slot common.Hash) (common.Hash, error) {
+	value, err := r.inner.Storage(addr, slot)
+	if err == nil && r.hooks.OnStorageLoad != nil {
+		r.hooks.OnStorageLoad(addr, slot, value)
+	}
+	return value, err
+}
+
+// Code implements Reader, retrieving the code associated with a particular account.
+func (r *hookedReader) Code(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	code, err := r.inner.Code(addr, codeHash)
+	if err == nil && r.hooks.OnCodeLoad != nil {
+		r.hooks.OnCodeLoad(addr, code)
+	}
+	return code, err
+}
+
+// CodeSize implements Reader, returning the size of the code associated with a particular account.
+func (r *hookedReader) CodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	size, err := r.inner.CodeSize(addr, codeHash)
+	if err != nil {
+		return 0, err
+	}
+	code, err := r.inner.Code(addr, codeHash)
+	if err == nil && r.hooks.OnCodeLoad != nil {
+		r.hooks.OnCodeLoad(addr, code)
+	}
+	return size, err
+}
+
+// Copy implements Reader
+func (r *hookedReader) Copy() Reader {
+	return &hookedReader{
+		inner: r.inner.Copy(),
+		hooks: r.hooks,
+	}
+}

core/state/statedb_hooked_test.go

@@ -46,6 +46,7 @@ func TestBurn(t *testing.T) {
 			}
 		},
 	})
+	defer hooked.Close()
 	createAndDestroy := func(addr common.Address) {
 		hooked.AddBalance(addr, uint256.NewInt(100), tracing.BalanceChangeUnspecified)
 		hooked.CreateContract(addr)
@@ -135,6 +136,7 @@ func TestHooks(t *testing.T) {
 			emitF("%v.code hash read: %v", addr, hash)
 		},
 	})
+	defer sdb.Close()
 	sdb.AddBalance(common.Address{0xaa}, uint256.NewInt(100), tracing.BalanceChangeUnspecified)
 	sdb.SubBalance(common.Address{0xaa}, uint256.NewInt(50), tracing.BalanceChangeTransfer)
 	sdb.SetNonce(common.Address{0xaa}, 1337)

core/state_processor.go

@@ -76,7 +76,9 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg
 	// Apply pre-execution system calls.
 	var tracingStateDB = vm.StateDB(statedb)
 	if hooks := cfg.Tracer; hooks != nil {
-		tracingStateDB = state.NewHookedState(statedb, hooks)
+		tsd := state.NewHookedState(statedb, hooks)
+		defer tsd.Close()
+		tracingStateDB = tsd
 	}
 	context = NewEVMBlockContext(header, p.chain, nil)
 	evm := vm.NewEVM(context, tracingStateDB, p.config, cfg)

core/tracing/hooks.go

@@ -198,6 +198,15 @@ type (
 
 	// BlockHashReadHook is called when EVM reads the blockhash of a block.
 	BlockHashReadHook = func(blockNumber uint64, hash common.Hash)
+
+	// AccountLoadHook is called when an account is loaded from the state.
+	AccountLoadHook = func(addr common.Address, acc *types.StateAccount)
+
+	// StorageLoadHook is called when a storage slot is loaded from the state.
+	StorageLoadHook = func(addr common.Address, slot common.Hash, value common.Hash)
+
+	// CodeLoadHook is called when the code of an account is loaded from the state.
+	CodeLoadHook = func(addr common.Address, code []byte)
 )
 
 type Hooks struct {
@@ -234,6 +243,10 @@ type Hooks struct {
 	OnStorageRead  StorageReadHook
 	// Block hash read
 	OnBlockHashRead BlockHashReadHook
+	// Account load
+	OnAccountLoad AccountLoadHook
+	OnStorageLoad StorageLoadHook
+	OnCodeLoad    CodeLoadHook
 }
 
 // Copy creates a new Hooks instance with all implemented hooks copied from the original.

eth/tracers/internal/tracetest/calltrace_test.go

@@ -126,7 +126,9 @@ func testCallTracer(tracerName string, dirPath string, t *testing.T) {
 			}
 			logState := vm.StateDB(st.StateDB)
 			if tracer.Hooks != nil {
-				logState = state.NewHookedState(st.StateDB, tracer.Hooks)
+				tsd := state.NewHookedState(st.StateDB, tracer.Hooks)
+				defer tsd.Close()
+				logState = tsd
 			}
 			msg, err := core.TransactionToMessage(tx, signer, context.BaseFee)
 			if err != nil {
@@ -360,7 +362,9 @@ func TestInternals(t *testing.T) {
 
 			logState := vm.StateDB(st.StateDB)
 			if hooks := tc.tracer.Hooks; hooks != nil {
-				logState = state.NewHookedState(st.StateDB, hooks)
+				tsd := state.NewHookedState(st.StateDB, hooks)
+				defer tsd.Close()
+				logState = tsd
 			}
 
 			tx, err := types.SignNewTx(key, signer, &types.LegacyTx{

internal/ethapi/simulate.go

@@ -187,7 +187,9 @@ func (sim *simulator) processBlock(ctx context.Context, block *simBlock, header,
 	)
 	tracingStateDB := vm.StateDB(sim.state)
 	if hooks := tracer.Hooks(); hooks != nil {
-		tracingStateDB = state.NewHookedState(sim.state, hooks)
+		tsd := state.NewHookedState(sim.state, hooks)
+		defer tsd.Close()
+		tracingStateDB = tsd
 	}
 	evm := vm.NewEVM(blockContext, tracingStateDB, sim.chainConfig, *vmConfig)
 	// It is possible to override precompiles with EVM bytecode, or