fix edge behavior and constructor

This commit is contained in:
Roy Crihfield 2022-02-21 16:30:47 +08:00
parent c532c17b05
commit c134d25de5

View File

@ -27,22 +27,32 @@ import (
type PrefixBoundIterator struct { type PrefixBoundIterator struct {
trie.NodeIterator trie.NodeIterator
EndPath []byte EndPath []byte
primed bool
} }
// Next advances until reaching EndPath, or the underlying iterator becomes invalid
func (it *PrefixBoundIterator) Next(descend bool) bool { func (it *PrefixBoundIterator) Next(descend bool) bool {
if it.EndPath == nil { // NodeIterator starts in an invalid state and must be advanced once before accessing values.
return it.NodeIterator.Next(descend) // Since this begins valid (pointing to the lower bound element), the first Next must be a no-op.
if !it.primed {
it.primed = true
return (it.EndPath == nil) || bytes.Compare(it.Path(), it.EndPath) < 0
} }
if !it.NodeIterator.Next(descend) { if !it.NodeIterator.Next(descend) {
return false return false
} }
// stop if underlying iterator went past endKey return (it.EndPath == nil) || bytes.Compare(it.Path(), it.EndPath) < 0
cmp := bytes.Compare(it.Path(), it.EndPath)
return cmp <= 0
} }
// Iterator with an upper bound value (hex path prefix) // Iterator with an upper bound value (hex path prefix)
func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) *PrefixBoundIterator { func NewPrefixBoundIterator(tree state.Trie, from, to []byte) *PrefixBoundIterator {
it := tree.NodeIterator(nil)
for it.Next(true) {
if bytes.Compare(from, it.Path()) <= 0 {
break
}
}
return &PrefixBoundIterator{NodeIterator: it, EndPath: to} return &PrefixBoundIterator{NodeIterator: it, EndPath: to}
} }
@ -99,7 +109,7 @@ func (gen *prefixGenerator) Next() {
} }
} }
// Generates paths that cut trie domain into `nbins` uniform conterminous bins (w/ opt. prefix) // MakePaths generates paths that cut trie domain into `nbins` uniform conterminous bins (w/ opt. prefix)
// eg. MakePaths([], 2) => [[0] [8]] // eg. MakePaths([], 2) => [[0] [8]]
// MakePaths([4], 32) => [[4 0 0] [4 0 8] [4 1 0]... [4 f 8]] // MakePaths([4], 32) => [[4 0 0] [4 0 8] [4 1 0]... [4 f 8]]
func MakePaths(prefix []byte, nbins uint) [][]byte { func MakePaths(prefix []byte, nbins uint) [][]byte {
@ -113,49 +123,28 @@ func MakePaths(prefix []byte, nbins uint) [][]byte {
return res return res
} }
func eachPrefixRange(prefix []byte, nbins uint, callback func([]byte, []byte)) { // truncate zeros from end of a path
prefixes := MakePaths(prefix, nbins) func truncateZeros(path []byte) []byte {
prefixes = append(prefixes, nil) // include tail l := len(path)
prefixes[0] = nil // set bin 0 left bound to nil to include root for ; l > 0 && path[l-1] == 0; l-- {
for i := 0; i < len(prefixes)-1; i++ {
key := prefixes[i]
if len(key)%2 != 0 { // zero-pad for odd-length keys
key = append(key, 0)
} }
callback(key, prefixes[i+1]) return path[:l]
}
func eachInterval(prefix []byte, nbins uint, cb func([]byte, []byte)) {
paths := MakePaths(prefix, nbins)
paths = append(paths, nil) // include tail
paths[0] = nil // set bin 0 left bound to nil to include root
for i := 0; i < len(paths)-1; i++ {
cb(truncateZeros(paths[i]), truncateZeros(paths[i+1]))
} }
} }
// Cut a trie by path prefix, returning `nbins` iterators covering its subtries // SubtrieIterators cuts a trie by path prefix, returning `nbins` iterators covering its subtries
func SubtrieIterators(tree state.Trie, nbins uint) []trie.NodeIterator { func SubtrieIterators(tree state.Trie, nbins uint) []trie.NodeIterator {
var iters []trie.NodeIterator var iters []trie.NodeIterator
eachPrefixRange(nil, nbins, func(from []byte, to []byte) { eachInterval(nil, nbins, func(from, to []byte) {
it := tree.NodeIterator(HexToKeyBytes(from)) iters = append(iters, NewPrefixBoundIterator(tree, from, to))
iters = append(iters, NewPrefixBoundIterator(it, to))
}) })
return iters return iters
} }
// Factory for per-bin subtrie iterators
type SubtrieIteratorFactory struct {
tree state.Trie
startPaths, endPaths [][]byte
}
func (fac *SubtrieIteratorFactory) Length() int { return len(fac.startPaths) }
func (fac *SubtrieIteratorFactory) IteratorAt(bin uint) *PrefixBoundIterator {
it := fac.tree.NodeIterator(HexToKeyBytes(fac.startPaths[bin]))
return NewPrefixBoundIterator(it, fac.endPaths[bin])
}
// Cut a trie by path prefix, returning `nbins` iterators covering its subtries
func NewSubtrieIteratorFactory(tree state.Trie, nbins uint) SubtrieIteratorFactory {
starts := make([][]byte, 0, nbins)
ends := make([][]byte, 0, nbins)
eachPrefixRange(nil, nbins, func(key []byte, endKey []byte) {
starts = append(starts, key)
ends = append(ends, endKey)
})
return SubtrieIteratorFactory{tree: tree, startPaths: starts, endPaths: ends}
}