From 529e52ce89191c91e252e0f24c8be602ef0cd527 Mon Sep 17 00:00:00 2001
From: rjl493456442 <garyrong0905@gmail.com>
Date: Mon, 1 Apr 2019 16:05:48 +0800
Subject: [PATCH] core: fix state iterator more elegant

---
 core/state/statedb.go | 16 +++++++++++-----
 core/vm/interface.go  |  2 +-
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/core/state/statedb.go b/core/state/statedb.go
index a019f4e8fe..4eb942d861 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -474,27 +474,33 @@ func (self *StateDB) CreateAccount(addr common.Address) {
 	}
 }
 
-func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common.Hash) bool) {
+func (db *StateDB) ForEachStorage(addr common.Address, cb func(key, value common.Hash) bool) error {
 	so := db.getStateObject(addr)
 	if so == nil {
-		return
+		return nil
 	}
 	it := trie.NewIterator(so.getTrie(db.db).NodeIterator(nil))
+
 	for it.Next() {
 		key := common.BytesToHash(db.trie.GetKey(it.Key))
 		if value, dirty := so.dirtyStorage[key]; dirty {
-			cb(key, value)
+			if !cb(key, value) {
+				return nil
+			}
 			continue
 		}
 
 		if len(it.Value) > 0 {
 			_, content, _, err := rlp.Split(it.Value)
 			if err != nil {
-				continue
+				return err
+			}
+			if !cb(key, common.BytesToHash(content)) {
+				return nil
 			}
-			cb(key, common.BytesToHash(content))
 		}
 	}
+	return nil
 }
 
 // Copy creates a deep, independent copy of the state.
diff --git a/core/vm/interface.go b/core/vm/interface.go
index fc15082f18..dd401466ad 100644
--- a/core/vm/interface.go
+++ b/core/vm/interface.go
@@ -63,7 +63,7 @@ type StateDB interface {
 	AddLog(*types.Log)
 	AddPreimage(common.Hash, []byte)
 
-	ForEachStorage(common.Address, func(common.Hash, common.Hash) bool)
+	ForEachStorage(common.Address, func(common.Hash, common.Hash) bool) error
 }
 
 // CallContext provides a basic interface for the EVM calling conventions. The EVM