add code 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

@@ -329,6 +329,28 @@ func (r *hookedReader) Storage(addr common.Address, slot common.Hash) (common.Ha
 	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{

core/tracing/hooks.go

@@ -204,6 +204,9 @@ type (
 
 	// 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 {
@@ -243,6 +246,7 @@ type Hooks struct {
 	// Account load
 	OnAccountLoad AccountLoadHook
 	OnStorageLoad StorageLoadHook
+	OnCodeLoad    CodeLoadHook
 }
 
 // Copy creates a new Hooks instance with all implemented hooks copied from the original.