Merge pull request #4 from vulcanize/dev

Concurrent traversal utilities

Author: i-norden

go.mod

@@ -0,0 +1,7 @@
+module github.com/vulcanize/go-eth-state-node-iterator
+
+go 1.13
+
+require github.com/ethereum/go-ethereum v1.9.11
+
+replace github.com/ethereum/go-ethereum v1.9.11 => github.com/vulcanize/go-ethereum v1.9.11-statediff-0.0.2

iterator.go

@@ -0,0 +1,185 @@
+//
+// Copyright © 2020 Vulcanize, Inc
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+package iterator
+
+import (
+	"bytes"
+	"math/bits"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/trie"
+)
+
+type PrefixBoundIterator struct {
+	current trie.NodeIterator
+	endKey  []byte
+}
+
+func (it *PrefixBoundIterator) Next(descend bool) bool {
+	if it.endKey == nil {
+		return it.current.Next(descend)
+	}
+	if !it.current.Next(descend) {
+		return false
+	}
+	// stop if underlying iterator went past endKey
+	cmp := bytes.Compare(it.current.Path(), it.endKey)
+	return cmp <= 0
+}
+
+func (it *PrefixBoundIterator) Error() error {
+	return it.current.Error()
+}
+func (it *PrefixBoundIterator) Hash() common.Hash {
+	return it.current.Hash()
+}
+func (it *PrefixBoundIterator) Path() []byte {
+	return it.current.Path()
+}
+func (it *PrefixBoundIterator) Leaf() bool {
+	return it.current.Leaf()
+}
+func (it *PrefixBoundIterator) LeafKey() []byte {
+	return it.current.LeafKey()
+}
+func (it *PrefixBoundIterator) LeafBlob() []byte {
+	return it.current.LeafBlob()
+}
+func (it *PrefixBoundIterator) LeafProof() [][]byte {
+	return it.current.LeafProof()
+}
+func (it *PrefixBoundIterator) Parent() common.Hash {
+	return it.current.Parent()
+}
+
+// Iterator with an upper bound value (hex path prefix)
+func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) *PrefixBoundIterator {
+	return &PrefixBoundIterator{current: it, endKey: to}
+}
+
+type prefixGenerator struct {
+	current   []byte
+	step      byte
+	stepIndex uint
+}
+
+func newPrefixGenerator(nbins uint) prefixGenerator {
+	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
+		}
+	}
+}
+
+// Generates paths that cut trie domain into `nbins` uniform conterminous bins (w/ opt. 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 {
+	var res [][]byte
+	for it := newPrefixGenerator(nbins); it.HasNext(); it.Next() {
+		next := make([]byte, len(prefix))
+		copy(next, prefix)
+		next = append(next, it.Value()...)
+		res = append(res, next)
+	}
+	return res
+}
+
+func eachPrefixRange(prefix []byte, nbins uint, callback func([]byte, []byte)) {
+	prefixes := MakePaths(prefix, nbins)
+	prefixes = append(prefixes, nil) // include tail
+	prefixes[0] = nil                // set bin 0 left bound to nil to include root
+	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])
+	}
+}
+
+// Cut a trie by path prefix, returning `nbins` iterators covering its subtries
+func SubtrieIterators(tree state.Trie, nbins uint) []trie.NodeIterator {
+	var iters []trie.NodeIterator
+	eachPrefixRange(nil, nbins, func(key []byte, endKey []byte) {
+		it := tree.NodeIterator(HexToKeyBytes(key))
+		iters = append(iters, NewPrefixBoundIterator(it, endKey))
+	})
+	return iters
+}
+
+// Factory for per-bin subtrie iterators
+type SubtrieIteratorFactory struct {
+	tree               state.Trie
+	startKeys, endKeys [][]byte
+}
+
+func (fac *SubtrieIteratorFactory) Length() int { return len(fac.startKeys) }
+
+func (fac *SubtrieIteratorFactory) IteratorAt(bin uint) *PrefixBoundIterator {
+	it := fac.tree.NodeIterator(HexToKeyBytes(fac.startKeys[bin]))
+	return NewPrefixBoundIterator(it, fac.endKeys[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, startKeys: starts, endKeys: ends}
+}

util.go

@@ -0,0 +1,49 @@
+package iterator
+
+import (
+	"bytes"
+	"github.com/ethereum/go-ethereum/trie"
+)
+
+func CompareNodes(a, b trie.NodeIterator) int {
+	if cmp := bytes.Compare(a.Path(), b.Path()); cmp != 0 {
+		return cmp
+	}
+	if a.Leaf() && !b.Leaf() {
+		return -1
+	} else if b.Leaf() && !a.Leaf() {
+		return 1
+	}
+	if cmp := bytes.Compare(a.Hash().Bytes(), b.Hash().Bytes()); cmp != 0 {
+		return cmp
+	}
+	if a.Leaf() && b.Leaf() {
+		return bytes.Compare(a.LeafBlob(), b.LeafBlob())
+	}
+	return 0
+}
+
+// hexToKeyBytes turns hex nibbles into key bytes.
+// This can only be used for keys of even length.
+func HexToKeyBytes(hex []byte) []byte {
+	if hasTerm(hex) {
+		hex = hex[:len(hex)-1]
+	}
+	if len(hex)&1 != 0 {
+		panic("can't convert hex key of odd length")
+	}
+	key := make([]byte, len(hex)/2)
+	decodeNibbles(hex, key)
+	return key
+}
+
+func decodeNibbles(nibbles []byte, bytes []byte) {
+	for bi, ni := 0, 0; ni < len(nibbles); bi, ni = bi+1, ni+2 {
+		bytes[bi] = nibbles[ni]<<4 | nibbles[ni+1]
+	}
+}
+
+// hasTerm returns whether a hex key has the terminator flag.
+func hasTerm(s []byte) bool {
+	return len(s) > 0 && s[len(s)-1] == 16
+}