core/state, core/vm, light: documented AccountRef, implemented state iterator

core/state/statedb.go

@@ -474,6 +474,27 @@ func (self *StateDB) CreateAccount(addr common.Address) {
 	}
 }
 
+func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common.Hash) bool) {
+	so := db.getStateObject(addr)
+	if so == nil {
+		return
+	}
+
+	// When iterating over the storage check the cache first
+	for h, value := range so.cachedStorage {
+		cb(h, value)
+	}
+
+	it := so.getTrie(db.db).Iterator()
+	for it.Next() {
+		// ignore cached values
+		key := common.BytesToHash(db.trie.GetKey(it.Key))
+		if _, ok := so.cachedStorage[key]; !ok {
+			cb(key, common.BytesToHash(it.Value))
+		}
+	}
+}
+
 // Copy creates a deep, independent copy of the state.
 // Snapshots of the copied state cannot be applied to the copy.
 func (self *StateDB) Copy() *StateDB {

core/vm/contract.go

@@ -27,9 +27,19 @@ type ContractRef interface {
 	Address() common.Address
 }
 
+// AccountRef implements ContractRef.
+//
+// Account references are used during EVM initialisation and
+// it's primary use is to fetch addresses. Removing this object
+// proves difficult because of the cached jump destinations which
+// are fetched from the parent contract (i.e. the caller), which
+// is a ContractRef.
 type AccountRef common.Address
 
-func (ar AccountRef) Address() common.Address  { return (common.Address)(ar) }
+// Address casts AccountRef to a Address
+func (ar AccountRef) Address() common.Address { return (common.Address)(ar) }
+
+// Reference returns a casted Address to AccountRef
 func Reference(addr common.Address) AccountRef { return (AccountRef)(addr) }
 
 // Contract represents an ethereum contract in the state database. It contains

core/vm/interface.go

@@ -60,6 +60,8 @@ type StateDB interface {
 
 	AddLog(*types.Log)
 	AddPreimage(common.Hash, []byte)
+
+	ForEachStorage(common.Address, func(common.Hash, common.Hash) bool)
 }
 
 // CallContext provides a basic interface for the EVM calling conventions. The EVM EVM

core/vm/logger.go

@@ -145,13 +145,11 @@ func (l *StructLogger) CaptureState(env *EVM, pc uint64, op OpCode, gas, cost *b
 			// Get the contract account and loop over each storage entry. This may involve looping over
 			// the trie and is a very expensive process.
 
-			/* TODO
-			env.StateDB.GetAccount(contract.Address()).ForEachStorage(func(key, value common.Hash) bool {
+			env.StateDB.ForEachStorage(contract.Address(), func(key, value common.Hash) bool {
 				storage[key] = value
 				// Return true, indicating we'd like to continue.
 				return true
 			})
-			*/
 		} else {
 			// copy a snapshot of the current storage to a new container.
 			storage = l.changedValues[contract.Address()].Copy()

light/state.go

@@ -268,6 +268,26 @@ func (self *LightState) CreateStateObject(ctx context.Context, addr common.Addre
 	return newSo, nil
 }
 
+// ForEachStorage calls a callback function for every key/value pair found
+// in the local storage cache. Note that unlike core/state.StateObject,
+// light.StateObject only returns cached values and doesn't download the
+// entire storage tree.
+func (self *LightState) ForEachStorage(ctx context.Context, addr common.Address, cb func(key, value common.Hash) bool) error {
+	so, err := self.GetStateObject(ctx, addr)
+	if err != nil {
+		return err
+	}
+
+	if so == nil {
+		return nil
+	}
+
+	for h, v := range so.storage {
+		cb(h, v)
+	}
+	return nil
+}
+
 //
 // Setting, copying of the state methods
 //

light/vm_env.go

@@ -81,6 +81,15 @@ func (s *VMState) SubBalance(addr common.Address, amount *big.Int) {
 	s.errHandler(err)
 }
 
+// ForEachStorage calls a callback function for every key/value pair found
+// in the local storage cache. Note that unlike core/state.StateObject,
+// light.StateObject only returns cached values and doesn't download the
+// entire storage tree.
+func (s *VMState) ForEachStorage(addr common.Address, cb func(key, value common.Hash) bool) {
+	err := s.state.ForEachStorage(s.ctx, addr, cb)
+	s.errHandler(err)
+}
+
 // GetBalance retrieves the balance from the given address or 0 if the account does
 // not exist
 func (s *VMState) GetBalance(addr common.Address) *big.Int {