New Trie iterator

1 files changed, 46 insertions, 0 deletions

diff --git a/ethutil/trie.go b/ethutil/trie.go
index c993e4d8f..18d0a5f0a 100644
--- a/ethutil/trie.go
+++ b/ethutil/trie.go
@@ -442,6 +442,8 @@ type TrieIterator struct {
 
 	shas   [][]byte
 	values []string
+
+	lastNode []byte
 }
 
 func (t *Trie) NewIterator() *TrieIterator {
@@ -513,3 +515,47 @@ func (it *TrieIterator) Key() string {
 func (it *TrieIterator) Value() string {
 	return ""
 }
+
+type EachCallback func(key string, node *Value)
+
+func (it *TrieIterator) Each(cb EachCallback) {
+	it.fetchNode(nil, NewValue(it.trie.Root).Bytes(), cb)
+}
+
+func (it *TrieIterator) fetchNode(key []int, node []byte, cb EachCallback) {
+	it.iterateNode(key, it.trie.cache.Get(node), cb)
+}
+
+func (it *TrieIterator) iterateNode(key []int, currentNode *Value, cb EachCallback) {
+	if currentNode.Len() == 2 {
+		k := CompactDecode(currentNode.Get(0).Str())
+
+		if currentNode.Get(1).Str() == "" {
+			it.iterateNode(key, currentNode.Get(1), cb)
+		} else {
+			pk := append(key, k...)
+
+			if k[len(k)-1] == 16 {
+				cb(DecodeCompact(pk), currentNode.Get(1))
+			} else {
+				it.fetchNode(pk, currentNode.Get(1).Bytes(), cb)
+			}
+		}
+	} else {
+		for i := 0; i < currentNode.Len(); i++ {
+			pk := append(key, i)
+			if i == 16 && currentNode.Get(i).Len() != 0 {
+				cb(DecodeCompact(pk), currentNode.Get(i))
+			} else {
+				if currentNode.Get(i).Str() == "" {
+					it.iterateNode(pk, currentNode.Get(i), cb)
+				} else {
+					val := currentNode.Get(i).Str()
+					if val != "" {
+						it.fetchNode(pk, []byte(val), cb)
+					}
+				}
+			}
+		}
+	}
+}