From 4d7c0cae7288b52decdd6b6a74f9552cddc64304 Mon Sep 17 00:00:00 2001
From: Nick Johnson <arachnid@notdot.net>
Date: Mon, 16 Jan 2017 17:37:35 +0000
Subject: [PATCH] trie, core/state: Add 'children' parameter to
 TrieIterator.next()

---
 core/state/iterator.go |  6 +++---
 trie/iterator.go       | 16 +++++++++++-----
 trie/iterator_test.go  |  2 +-
 trie/sync_test.go      |  2 +-
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/core/state/iterator.go b/core/state/iterator.go
index a58a15ad39..3e82dca099 100644
--- a/core/state/iterator.go
+++ b/core/state/iterator.go
@@ -80,7 +80,7 @@ func (it *NodeIterator) step() error {
 	}
 	// If we had data nodes previously, we surely have at least state nodes
 	if it.dataIt != nil {
-		if cont := it.dataIt.Next(); !cont {
+		if cont := it.dataIt.Next(true); !cont {
 			if it.dataIt.Error != nil {
 				return it.dataIt.Error
 			}
@@ -94,7 +94,7 @@ func (it *NodeIterator) step() error {
 		return nil
 	}
 	// Step to the next state trie node, terminating if we're out of nodes
-	if cont := it.stateIt.Next(); !cont {
+	if cont := it.stateIt.Next(true); !cont {
 		if it.stateIt.Error != nil {
 			return it.stateIt.Error
 		}
@@ -120,7 +120,7 @@ func (it *NodeIterator) step() error {
 		return err
 	}
 	it.dataIt = trie.NewNodeIterator(dataTrie)
-	if !it.dataIt.Next() {
+	if !it.dataIt.Next(true) {
 		it.dataIt = nil
 	}
 	if !bytes.Equal(account.CodeHash, emptyCodeHash) {
diff --git a/trie/iterator.go b/trie/iterator.go
index 4722b99269..75f6fe14e4 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -40,7 +40,7 @@ func NewIterator(trie *Trie) *Iterator {
 
 // Next moves the iterator forward one key-value entry.
 func (it *Iterator) Next() bool {
-	for it.nodeIt.Next() {
+	for it.nodeIt.Next(true) {
 		if it.nodeIt.Leaf {
 			it.Key = it.makeKey()
 			it.Value = it.nodeIt.LeafBlob
@@ -104,14 +104,15 @@ func NewNodeIterator(trie *Trie) *NodeIterator {
 
 // Next moves the iterator to the next node, returning whether there are any
 // further nodes. In case of an internal error this method returns false and
-// sets the Error field to the encountered failure.
-func (it *NodeIterator) Next() bool {
+// sets the Error field to the encountered failure. If `children` is false,
+// skips iterating over any subnodes of the current node.
+func (it *NodeIterator) Next(children bool) bool {
 	// If the iterator failed previously, don't do anything
 	if it.Error != nil {
 		return false
 	}
 	// Otherwise step forward with the iterator and report any errors
-	if err := it.step(); err != nil {
+	if err := it.step(children); err != nil {
 		it.Error = err
 		return false
 	}
@@ -119,7 +120,7 @@ func (it *NodeIterator) Next() bool {
 }
 
 // step moves the iterator to the next node of the trie.
-func (it *NodeIterator) step() error {
+func (it *NodeIterator) step(children bool) error {
 	if it.trie == nil {
 		// Abort if we reached the end of the iteration
 		return nil
@@ -135,6 +136,11 @@ func (it *NodeIterator) step() error {
 		return nil
 	}
 
+	if !children {
+		// If we're skipping children, pop the current node first
+		it.stack = it.stack[:len(it.stack)-1]
+	}
+
 	// Continue iteration to the next child
 	for {
 		if len(it.stack) == 0 {
diff --git a/trie/iterator_test.go b/trie/iterator_test.go
index 41672d8d4f..0582e8d752 100644
--- a/trie/iterator_test.go
+++ b/trie/iterator_test.go
@@ -101,7 +101,7 @@ func TestNodeIteratorCoverage(t *testing.T) {
 
 	// Gather all the node hashes found by the iterator
 	hashes := make(map[common.Hash]struct{})
-	for it := NewNodeIterator(trie); it.Next(); {
+	for it := NewNodeIterator(trie); it.Next(true); {
 		if it.Hash != (common.Hash{}) {
 			hashes[it.Hash] = struct{}{}
 		}
diff --git a/trie/sync_test.go b/trie/sync_test.go
index 4168c4d658..547a6938f1 100644
--- a/trie/sync_test.go
+++ b/trie/sync_test.go
@@ -81,7 +81,7 @@ func checkTrieConsistency(db Database, root common.Hash) error {
 		return nil // // Consider a non existent state consistent
 	}
 	it := NewNodeIterator(trie)
-	for it.Next() {
+	for it.Next(true) {
 	}
 	return it.Error
 }