diff --git a/ethtrie/iterator.go b/ethtrie/iterator.go new file mode 100644 index 0000000000..e168162290 --- /dev/null +++ b/ethtrie/iterator.go @@ -0,0 +1,143 @@ +package ethtrie + +import ( + "bytes" + + "github.com/ethereum/eth-go/ethutil" +) + +type NodeType byte + +const ( + EmptyNode NodeType = iota + BranchNode + LeafNode + ExtNode +) + +func getType(node *ethutil.Value) NodeType { + if node.Len() == 0 { + return EmptyNode + } + + if node.Len() == 2 { + k := CompactDecode(node.Get(0).Str()) + if HasTerm(k) { + return LeafNode + } + + return ExtNode + } + + return BranchNode +} + +type Iterator struct { + Path [][]byte + trie *Trie + + Key []byte + Value *ethutil.Value +} + +func NewIterator(trie *Trie) *Iterator { + return &Iterator{trie: trie} +} + +func (self *Iterator) key(node *ethutil.Value, path [][]byte) []byte { + switch getType(node) { + case LeafNode: + k := RemTerm(CompactDecode(node.Get(0).Str())) + + self.Path = append(path, k) + self.Value = node.Get(1) + + return k + case BranchNode: + if node.Get(16).Len() > 0 { + return []byte{16} + } + + for i := byte(0); i < 16; i++ { + o := self.key(self.trie.getNode(node.Get(int(i)).Raw()), append(path, []byte{i})) + if o != nil { + return append([]byte{i}, o...) + } + } + case ExtNode: + currKey := node.Get(0).Bytes() + + return self.key(self.trie.getNode(node.Get(1).Raw()), append(path, currKey)) + } + + return nil +} + +func (self *Iterator) next(node *ethutil.Value, key []byte, path [][]byte) []byte { + switch typ := getType(node); typ { + case EmptyNode: + return nil + case BranchNode: + if len(key) > 0 { + subNode := self.trie.getNode(node.Get(int(key[0])).Raw()) + + o := self.next(subNode, key[1:], append(path, key[:1])) + if o != nil { + return append([]byte{key[0]}, o...) + } + } + + var r byte = 0 + if len(key) > 0 { + r = key[0] + 1 + } + + for i := r; i < 16; i++ { + subNode := self.trie.getNode(node.Get(int(i)).Raw()) + o := self.key(subNode, append(path, []byte{i})) + if o != nil { + return append([]byte{i}, o...) + } + } + case LeafNode, ExtNode: + k := RemTerm(CompactDecode(node.Get(0).Str())) + if typ == LeafNode { + if bytes.Compare([]byte(k), []byte(key)) > 0 { + self.Value = node.Get(1) + self.Path = append(path, k) + + return k + } + } else { + subNode := self.trie.getNode(node.Get(1).Raw()) + subKey := key[len(k):] + var ret []byte + if BeginsWith(key, k) { + ret = self.next(subNode, subKey, append(path, k)) + } else if bytes.Compare(k, key[:len(k)]) > 0 { + ret = self.key(node, append(path, k)) + } else { + ret = nil + } + + if ret != nil { + return append(k, ret...) + } + } + } + + return nil +} + +// Get the next in keys +func (self *Iterator) Next(key string) []byte { + self.trie.mut.Lock() + defer self.trie.mut.Unlock() + + k := RemTerm(CompactHexDecode(key)) + n := self.next(self.trie.getNode(self.trie.Root), k, nil) + + self.Key = []byte(DecodeCompact(n)) + + return self.Key +}