iterator slice instead of callback

This commit is contained in:
Roy Crihfield 2020-09-03 23:52:47 -05:00
parent 8b8f718505
commit 70b374787a

View File

@ -122,7 +122,7 @@ func (gen *prefixGenerator) Next() {
}
}
// Generates paths that cut the trie domain into "nbins" bins, w/ optional prefix
// Generates paths that cut the trie domain into "nbins" uniform bins, w/ optional prefix
// eg. MakePaths([], 2) => [[0] [8]]
// MakePaths([4], 32) => [[4 0 0] [4 0 8] [4 1 0]... [4 f 8]]
func MakePaths(prefix []byte, nbins uint) [][]byte {
@ -136,19 +136,21 @@ func MakePaths(prefix []byte, nbins uint) [][]byte {
return res
}
// Apply a function to nbins subtries divided according to path prefix
func VisitSubtries(tree state.Trie, nbins uint, callback func(NodeIterator)) {
// Cut a trie by path prefix, returning `nbins` iterators covering its subtries
func SubtrieIterators(tree state.Trie, nbins uint) []NodeIterator {
prefixes := MakePaths(nil, nbins)
// pre- and postpend nil to include root & tail
prefixes = append(prefixes, nil)
prefixes = append([][]byte{nil}, prefixes...)
var iters []NodeIterator
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)
}
it := tree.NodeIterator(HexToKeyBytes(key))
callback(NewPrefixBoundIterator(it, prefixes[i+1]))
iters = append(iters, NewPrefixBoundIterator(it, prefixes[i+1]))
}
return iters
}