New Trie iterator

This commit is contained in:
obscuren
2014-05-27 01:08:51 +02:00
parent 4c7bd75c1a
commit 5cdfee5143
3 changed files with 71 additions and 12 deletions

View File

@@ -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)
}
}
}
}
}
}