op-geth/trie/iterator.go

144 lines
2.7 KiB
Go

package trie
import (
"bytes"
"github.com/ethereum/go-ethereum/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
}