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()) + }) }