From 70b374787a724c481a508c3e389db7c21f08f61a Mon Sep 17 00:00:00 2001 From: Roy Crihfield Date: Thu, 3 Sep 2020 23:52:47 -0500 Subject: [PATCH] iterator slice instead of callback --- pkg/iterator/iterator.go | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/pkg/iterator/iterator.go b/pkg/iterator/iterator.go index efb2d4d55..8e24800b7 100644 --- a/pkg/iterator/iterator.go +++ b/pkg/iterator/iterator.go @@ -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 }