core/state, trie: add pre-order iterator hook, fix missing indices

core/state/iterator.go

@@ -26,6 +26,11 @@ import (
 	"github.com/ethereum/go-ethereum/trie"
 )
 
+// NodeIteratorSubtreeCallback is a callback that is invoked by the trie node
+// iterator whenever it descends into a new subtree, giving the possibility of
+// skipping the branch if it would be non-useful (e.g. already processed before).
+type NodeIteratorSubtreeCallback func(hash, parent common.Hash) bool
+
 // NodeIterator is an iterator to traverse the entire state trie post-order,
 // including all of the contract code and contract state tries.
 type NodeIterator struct {
@@ -38,6 +43,14 @@ type NodeIterator struct {
 	codeHash    common.Hash // Hash of the contract source code
 	code        []byte      // Source code associated with a contract
 
+	// Callback method enabling the user to control subtree descent.
+	//
+	// Note: This hook is invoked pre-order to allow the user to cancel the traversal
+	// of a subtree, but the iterator itself is post-order. This means that the hook
+	// may be invoked multiple times during a single iteration step (when descending)
+	// or never in multiple iteration steps (when ascending).
+	PreOrderHook NodeIteratorSubtreeCallback
+
 	Hash   common.Hash // Hash of the current entry being iterated (nil if not standalone)
 	Entry  interface{} // Current state entry being iterated (internal representation)
 	Parent common.Hash // Hash of the first full ancestor node (nil if current is the root)
@@ -66,6 +79,7 @@ func (it *NodeIterator) step() {
 	// Initialize the iterator if we've just started
 	if it.stateIt == nil {
 		it.stateIt = trie.NewNodeIterator(it.state.trie.Trie)
+		it.stateIt.PreOrderHook = trie.NodeIteratorSubtreeCallback(it.PreOrderHook)
 	}
 	// If we had data nodes previously, we surely have at least state nodes
 	if it.dataIt != nil {
@@ -99,22 +113,25 @@ func (it *NodeIterator) step() {
 	if err != nil {
 		panic(err)
 	}
+	it.accountHash = it.stateIt.Parent
+
+	if bytes.Compare(account.CodeHash, emptyCodeHash) != 0 {
+		if hash := common.BytesToHash(account.CodeHash); it.PreOrderHook == nil || it.PreOrderHook(hash, it.accountHash) {
+			it.codeHash = hash
+			if it.code, err = it.state.db.Get(account.CodeHash); err != nil {
+				panic(fmt.Sprintf("code %x: %v", account.CodeHash, err))
+			}
+		}
+	}
 	dataTrie, err := trie.New(account.Root, it.state.db)
 	if err != nil {
 		panic(err)
 	}
 	it.dataIt = trie.NewNodeIterator(dataTrie)
+	it.dataIt.PreOrderHook = trie.NodeIteratorSubtreeCallback(it.PreOrderHook)
 	if !it.dataIt.Next() {
 		it.dataIt = nil
 	}
-	if bytes.Compare(account.CodeHash, emptyCodeHash) != 0 {
-		it.codeHash = common.BytesToHash(account.CodeHash)
-		it.code, err = it.state.db.Get(account.CodeHash)
-		if err != nil {
-			panic(fmt.Sprintf("code %x: %v", account.CodeHash, err))
-		}
-	}
-	it.accountHash = it.stateIt.Parent
 }
 
 // retrieve pulls and caches the current state entry the iterator is traversing.

core/state/iterator_test.go

@@ -59,3 +59,63 @@ func TestNodeIteratorCoverage(t *testing.T) {
 		}
 	}
 }
+
+// Tests that the node iterator hook is invoked for all the nodes of the state,
+// also testing that the iterator indeed avoids visiting aborted branches.
+func TestNodeIteratorHookCoverageFull(t *testing.T)  { testNodeIteratorHookCoverage(t, false) }
+func TestNodeIteratorHookCoverageDedup(t *testing.T) { testNodeIteratorHookCoverage(t, true) }
+
+func testNodeIteratorHookCoverage(t *testing.T, dedup bool) {
+	// Create some arbitrary test state to iterate
+	db, root, _ := makeTestState(nil)
+
+	state, err := New(root, db)
+	if err != nil {
+		t.Fatalf("failed to create state trie at %x: %v", root, err)
+	}
+	// Gather all the node hashes found by the iterator
+	hashes := make(map[common.Hash]int)
+
+	it := NewNodeIterator(state)
+	it.PreOrderHook = func(hash, parent common.Hash) bool {
+		if !dedup {
+			return true
+		}
+		return hashes[hash] == 0
+	}
+	for it.Next() {
+		if it.Hash != (common.Hash{}) {
+			hashes[it.Hash]++
+		}
+	}
+	// Cross check the hashes and the database itself
+	for hash, _ := range hashes {
+		if _, err := db.Get(hash.Bytes()); err != nil {
+			t.Errorf("failed to retrieve reported node %x: %v", hash, err)
+		}
+	}
+	for _, key := range db.(*ethdb.MemDatabase).Keys() {
+		if bytes.HasPrefix(key, []byte("secure-key-")) {
+			continue
+		}
+		if bytes.HasPrefix(key, trie.ParentReferenceIndexPrefix) {
+			continue
+		}
+		if _, ok := hashes[common.BytesToHash(key)]; !ok {
+			t.Errorf("state entry not reported %x", key)
+		}
+	}
+	// Check whether duplicates were avoided or not
+	duplicates := 0
+	for hash, count := range hashes {
+		if count > 1 {
+			duplicates++
+			if dedup {
+				t.Errorf("duplicate (%d) iteration: %x", count, hash)
+			}
+		}
+	}
+	if !dedup && duplicates == 0 {
+		t.Errorf("iterator didn't traverse common subtrees")
+	}
+}

core/state/sync_test.go

@@ -44,9 +44,10 @@ func makeTestState(referrers []common.Hash) (ethdb.Database, common.Hash, []*tes
 
 	// Fill it with some arbitrary data
 	accounts := []*testAccount{}
-	for i := byte(0); i < 96; i++ {
+	for i := byte(0); i < 1; i++ {
-		obj := state.GetOrNewStateObject(common.BytesToAddress([]byte{i}))
+		for j := byte(0); j < 2; j++ {
-		acc := &testAccount{address: common.BytesToAddress([]byte{i})}
+			obj := state.GetOrNewStateObject(common.BytesToAddress([]byte{i, j}))
+			acc := &testAccount{address: common.BytesToAddress([]byte{i, j})}
 
 			obj.AddBalance(big.NewInt(int64(11 * i)))
 			acc.balance = big.NewInt(int64(11 * i))
@@ -61,6 +62,7 @@ func makeTestState(referrers []common.Hash) (ethdb.Database, common.Hash, []*tes
 			state.UpdateStateObject(obj)
 			accounts = append(accounts, acc)
 		}
+	}
 	root, _ := state.CommitIndexed(referrers)
 
 	// Remove any potentially cached data from the test state creation

trie/iterator.go

@@ -159,11 +159,24 @@ type nodeIteratorState struct {
 	child  int         // Child to be processed next
 }
 
+// NodeIteratorSubtreeCallback is a callback that is invoked by the trie node
+// iterator whenever it descends into a new subtree, giving the possibility of
+// skipping the branch if it would be non-useful (e.g. already processed before).
+type NodeIteratorSubtreeCallback func(hash, parent common.Hash) bool
+
 // NodeIterator is an iterator to traverse the trie post-order.
 type NodeIterator struct {
 	trie  *Trie                // Trie being iterated
 	stack []*nodeIteratorState // Hierarchy of trie nodes persisting the iteration state
 
+	// Callback method enabling the user to control subtree descent.
+	//
+	// Note: This hook is invoked pre-order to allow the user to cancel the traversal
+	// of a subtree, but the iterator itself is post-order. This means that the hook
+	// may be invoked multiple times during a single iteration step (when descending)
+	// or never in multiple iteration steps (when ascending).
+	PreOrderHook NodeIteratorSubtreeCallback
+
 	Hash     common.Hash // Hash of the current node being iterated (nil if not standalone)
 	Node     node        // Current node being iterated (internal representation)
 	Parent   common.Hash // Hash of the first full ancestor node (nil if current is the root)
@@ -241,7 +254,9 @@ func (it *NodeIterator) step() {
 			if err != nil {
 				panic(err)
 			}
-			it.stack = append(it.stack, &nodeIteratorState{hash: common.BytesToHash(hash), node: node, parent: ancestor, child: -1})
+			if hash := common.BytesToHash(hash); it.PreOrderHook == nil || it.PreOrderHook(hash, ancestor) {
+				it.stack = append(it.stack, &nodeIteratorState{hash: hash, node: node, parent: ancestor, child: -1})
+			}
 		} else {
 			break
 		}

trie/iterator_test.go

@@ -79,3 +79,55 @@ func TestNodeIteratorCoverage(t *testing.T) {
 		}
 	}
 }
+
+// Tests that the node iterator hook is invoked for all the nodes of the trie,
+// also testing that the iterator indeed avoids visiting aborted branches.
+func TestNodeIteratorHookCoverageFull(t *testing.T)  { testNodeIteratorHookCoverage(t, false) }
+func TestNodeIteratorHookCoverageDedup(t *testing.T) { testNodeIteratorHookCoverage(t, true) }
+
+func testNodeIteratorHookCoverage(t *testing.T, dedup bool) {
+	// Create some arbitrary test trie to iterate
+	db, trie, _ := makeTestTrie(nil)
+
+	// Gather all the node hashes found by the iterator
+	hashes := make(map[common.Hash]int)
+
+	it := NewNodeIterator(trie)
+	it.PreOrderHook = func(hash, parent common.Hash) bool {
+		if !dedup {
+			return true
+		}
+		return hashes[hash] == 0
+	}
+	for it.Next() {
+		if it.Hash != (common.Hash{}) {
+			hashes[it.Hash]++
+		}
+	}
+	// Cross check the hashes and the database itself
+	for hash, _ := range hashes {
+		if _, err := db.Get(hash.Bytes()); err != nil {
+			t.Errorf("failed to retrieve reported node %x: %v", hash, err)
+		}
+	}
+	for _, key := range db.(*ethdb.MemDatabase).Keys() {
+		if len(key) == common.HashLength {
+			if _, ok := hashes[common.BytesToHash(key)]; !ok {
+				t.Errorf("state entry not reported %x", key)
+			}
+		}
+	}
+	// Check whether duplicates were avoided or not
+	duplicates := 0
+	for hash, count := range hashes {
+		if count > 1 {
+			duplicates++
+			if dedup {
+				t.Errorf("duplicate (%d) iteration: %x", count, hash)
+			}
+		}
+	}
+	if !dedup && duplicates == 0 {
+		t.Errorf("iterator didn't traverse common subtrees")
+	}
+}

trie/sync.go

@@ -34,11 +34,19 @@ type request struct {
 	depth   int        // Depth level within the trie the node is located to prioritise DFS
 	deps    int        // Number of dependencies before allowed to commit this node
 
-	referrers []common.Hash // External parents of this node to index in addition to internal ones
+	referrers map[common.Hash]struct{} // External parents of this node to index in addition to internal ones
 
 	callback TrieSyncLeafCallback // Callback to invoke if a leaf node it reached on this branch
 }
 
+// referrer adds a new referring parent to a sync node request.
+func (r *request) referrer(hash common.Hash) {
+	if r.referrers == nil {
+		r.referrers = make(map[common.Hash]struct{})
+	}
+	r.referrers[hash] = struct{}{}
+}
+
 // SyncResult is a simple list to return missing nodes along with their request
 // hashes.
 type SyncResult struct {
@@ -72,14 +80,14 @@ func NewTrieSync(root common.Hash, database ethdb.Database, parent common.Hash,
 }
 
 // AddSubTrie registers a new trie to the sync code, rooted at the designated parent.
-func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, callback TrieSyncLeafCallback) {
+func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, callback TrieSyncLeafCallback) error {
 	// Short circuit if the trie is empty or already known
 	if root == emptyRoot {
-		return
+		return nil
 	}
 	blob, _ := s.database.Get(root.Bytes())
 	if local, err := decodeNode(blob); local != nil && err == nil {
-		return
+		return storeParentReferenceEntry(parent.Bytes(), root.Bytes(), s.database)
 	}
 	// Assemble the new sub-trie sync request
 	node := node(hashNode(root.Bytes()))
@@ -99,22 +107,23 @@ func (s *TrieSync) AddSubTrie(root common.Hash, depth int, parent common.Hash, c
 			ancestor.deps++
 			req.parents = append(req.parents, ancestor)
 		}
-		req.referrers = append(req.referrers, parent)
+		req.referrer(parent)
 	}
 	s.schedule(req)
+	return nil
 }
 
 // AddRawEntry schedules the direct retrieval of a state entry that should not be
 // interpreted as a trie node, but rather accepted and stored into the database
 // as is. This method's goal is to support misc state metadata retrievals (e.g.
 // contract code).
-func (s *TrieSync) AddRawEntry(hash common.Hash, depth int, parent common.Hash) {
+func (s *TrieSync) AddRawEntry(hash common.Hash, depth int, parent common.Hash) error {
 	// Short circuit if the entry is empty or already known
 	if hash == emptyState {
-		return
+		return nil
 	}
 	if blob, _ := s.database.Get(hash.Bytes()); blob != nil {
-		return
+		return storeParentReferenceEntry(parent.Bytes(), hash.Bytes(), s.database)
 	}
 	// Assemble the new sub-trie sync request
 	req := &request{
@@ -131,9 +140,10 @@ func (s *TrieSync) AddRawEntry(hash common.Hash, depth int, parent common.Hash)
 			ancestor.deps++
 			req.parents = append(req.parents, ancestor)
 		}
-		req.referrers = append(req.referrers, parent)
+		req.referrer(parent)
 	}
 	s.schedule(req)
+	return nil
 }
 
 // Missing retrieves the known missing nodes from the trie for retrieval.
@@ -196,6 +206,13 @@ func (s *TrieSync) schedule(req *request) {
 	// If we're already requesting this node, add a new reference and stop
 	if old, ok := s.requests[req.hash]; ok {
 		old.parents = append(old.parents, req.parents...)
+
+		for hash, _ := range req.referrers {
+			old.referrer(hash)
+		}
+		for _, parent := range req.parents {
+			old.referrer(parent.hash)
+		}
 		return
 	}
 	// Schedule the request for future retrieval
@@ -280,7 +297,7 @@ func (s *TrieSync) commit(req *request, batch ethdb.Batch) (err error) {
 			return err
 		}
 	}
-	for _, referrer := range req.referrers {
+	for referrer, _ := range req.referrers {
 		if err := storeParentReferenceEntry(referrer[:], req.hash[:], batch); err != nil {
 			return err
 		}