improve prefix gen

use stateful iterator, and allow splitting into powers of 2 rather than 16
This commit is contained in:
Roy Crihfield 2020-08-31 11:06:03 -05:00
parent c3c3195991
commit d82b49b8ba

View File

@ -17,10 +17,9 @@
package iterator package iterator
import ( import (
// "fmt"
"bytes" "bytes"
"math/bits"
// "github.com/ethereum/go-ethereum/common"
"github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/common"
"github.com/ethereum/go-ethereum/core/state" "github.com/ethereum/go-ethereum/core/state"
"github.com/ethereum/go-ethereum/trie" "github.com/ethereum/go-ethereum/trie"
@ -69,37 +68,73 @@ func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) NodeIterator {
return &prefixBoundIterator{current: it, endKey: to} return &prefixBoundIterator{current: it, endKey: to}
} }
// array of 0..f with prefix type prefixGenerator struct {
func prefixedNibbles(prefix []byte) [][]byte { current []byte
var ret [][]byte step byte
for i := byte(0); i < 16; i++ { stepIndex uint
elem := make([]byte, len(prefix)) }
copy(elem, prefix)
elem = append(elem, i) func newPrefixGenerator(nbins uint) prefixGenerator {
ret = append(ret, elem) if bits.OnesCount(nbins) != 1 {
panic("nbins must be a power of 2")
}
// determine step dist. and path index at which to step
var step byte
var stepIndex uint
for ; nbins != 0; stepIndex++ {
divisor := byte(nbins & 0xf)
if divisor != 0 {
step = 0x10 / divisor
}
nbins = nbins >> 4
}
return prefixGenerator{
current: make([]byte, stepIndex),
step: step,
stepIndex: stepIndex-1,
}
}
func (gen *prefixGenerator) Value() []byte {
return gen.current
}
func (gen *prefixGenerator) HasNext() bool {
return gen.current[0] <= 0xf
}
func (gen *prefixGenerator) Next() {
gen.current[gen.stepIndex] += gen.step
overflow := false
for ix := 0; ix < len(gen.current); ix++ {
rix := len(gen.current)-1-ix // reverse
if overflow {
gen.current[rix]++
overflow = false
}
if rix != 0 && gen.current[rix] > 0xf {
gen.current[rix] = 0
overflow = true
}
} }
return ret
} }
// Generates ordered cartesian product of all nibbles of given length, w/ optional prefix // Generates ordered cartesian product of all nibbles of given length, w/ optional prefix
// eg. MakePaths([4], 2) => [[4 0 0] [4 0 1] ... [4 f f]] // eg. MakePaths([4], 2) => [[4 0 0] [4 0 1] ... [4 f f]]
func MakePaths(prefix []byte, length int) [][]byte { func MakePaths(prefix []byte, nbins uint) [][]byte {
paths := [][]byte{prefix} var res [][]byte
for depth := 0; depth < length; depth++ { for it := newPrefixGenerator(nbins); it.HasNext(); it.Next() {
var newPaths [][]byte next := make([]byte, len(prefix))
for _, path := range paths { copy(next, prefix)
for _, newPath := range prefixedNibbles(path) { next = append(next, it.Value()...)
newPaths = append(newPaths, newPath) res = append(res, next)
} }
} return res
paths = newPaths
}
return paths
} }
// Apply a function to 16^cutdepth subtries divided by path prefix // Apply a function to 16^cutdepth subtries divided by path prefix
func VisitSubtries(tree state.Trie, cutDepth int, callback func(NodeIterator)) { func VisitSubtries(tree state.Trie, nbins uint, callback func(NodeIterator)) {
prefixes := MakePaths(nil, cutDepth) prefixes := MakePaths(nil, nbins)
// pre- and postpend nil to include root & tail // pre- and postpend nil to include root & tail
prefixes = append(prefixes, nil) prefixes = append(prefixes, nil)
prefixes = append([][]byte{nil}, prefixes...) prefixes = append([][]byte{nil}, prefixes...)