New Trie iterator

3 files changed, 71 insertions, 12 deletions

diff --git a/ethutil/encoding.go b/ethutil/encoding.go
index 1f661947a..9fcdf3edf 100644
--- a/ethutil/encoding.go
+++ b/ethutil/encoding.go
@@ -59,3 +59,18 @@ func CompactHexDecode(str string) []int {
 
 	return hexSlice
 }
+
+func DecodeCompact(key []int) string {
+	base := "0123456789abcdef"
+	var str string
+
+	for _, v := range key {
+		if v < 16 {
+			str += string(base[v])
+		}
+	}
+
+	res, _ := hex.DecodeString(str)
+
+	return string(res)
+}
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)
+					}
+				}
+			}
+		}
+	}
+}
diff --git a/ethutil/trie_test.go b/ethutil/trie_test.go
index d74d129ac..c89f2fbb7 100644
--- a/ethutil/trie_test.go
+++ b/ethutil/trie_test.go
@@ -154,7 +154,7 @@ func TestTrieDeleteWithValue(t *testing.T) {
 
 }
 
-func TestTrieIterator(t *testing.T) {
+func TestTriePurge(t *testing.T) {
 	_, trie := New()
 	trie.Update("c", LONG_WORD)
 	trie.Update("ca", LONG_WORD)
@@ -171,16 +171,14 @@ func TestTrieIterator(t *testing.T) {
 	}
 }
 
-func TestHashes(t *testing.T) {
+func TestTrieIt(t *testing.T) {
 	_, trie := New()
-	trie.Update("cat", "dog")
-	trie.Update("ca", "dude")
-	trie.Update("doge", "1234567890abcdefghijklmnopqrstuvwxxzABCEFGHIJKLMNOPQRSTUVWXYZ")
-	trie.Update("dog", "test")
-	trie.Update("test", "1234567890abcdefghijklmnopqrstuvwxxzABCEFGHIJKLMNOPQRSTUVWXYZ")
-	fmt.Printf("%x\n", trie.Root)
-	trie.Delete("dog")
-	fmt.Printf("%x\n", trie.Root)
-	trie.Delete("test")
-	fmt.Printf("%x\n", trie.Root)
+	trie.Update("c", LONG_WORD)
+	trie.Update("ca", LONG_WORD)
+	trie.Update("cat", LONG_WORD)
+
+	it := trie.NewIterator()
+	it.Each(func(key string, node *Value) {
+		fmt.Println(key, ":", node.Str())
+	})
 }