iterator package

* stripped down Nodeiterator interface

* prefix-bound iterator for traversing a specific range of keys

* visitor over subtries divided at a specific depth

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

pkg/iterator/iterator.go

@@ -0,0 +1,115 @@
+//
+// 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 (
+	// "fmt"
+	"bytes"
+
+	// "github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/state"
+)
+
+type NodeIterator interface {
+	Next(bool) bool
+	Error() error
+	Hash() common.Hash
+	Path() []byte
+	Leaf() bool
+}
+
+type prefixBoundIterator struct {
+	current trie.NodeIterator
+	endKey []byte
+	// endHash common.Hash
+}
+
+func (it *prefixBoundIterator) Next(descend bool) bool {
+	if it.endKey == nil {		// todo: want?
+		return it.current.Next(descend)
+	}
+	cmp := bytes.Compare(it.current.Path(), it.endKey) // stop before endKey
+	if cmp >= 0 {
+		return false
+	}
+	return it.current.Next(descend)
+}
+
+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()
+}
+
+// Iterator with an upper bound value (hex path prefix)
+func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) NodeIterator {
+	return &prefixBoundIterator{current: it, endKey: to}
+}
+
+// array of 0..f with prefix
+func prefixedNibbles(prefix []byte) [][]byte {
+	var ret [][]byte
+	for i := byte(0); i < 16; i++ {
+		elem := make([]byte, len(prefix))
+		copy(elem, prefix)
+		elem = append(elem, i)
+		ret = append(ret, elem)
+	}
+	return ret
+}
+
+// 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]]
+func MakePaths(prefix []byte, length int) [][]byte {
+	paths := [][]byte{prefix}
+	for depth := 0; depth < length; depth++ {
+		var newPaths [][]byte
+		for _, path := range paths {
+			for _, newPath := range prefixedNibbles(path) {
+				newPaths = append(newPaths, newPath)
+			}
+		}
+		paths = newPaths
+	}
+	return paths
+}
+
+// Apply a function to 16^cutdepth subtries divided by path prefix
+func VisitSubtries(tree state.Trie, cutDepth int, callback func(NodeIterator)) {
+	prefixes := MakePaths(nil, cutDepth)
+	// pre- and postpend nil to include root & tail
+	prefixes = append(prefixes, nil)
+	prefixes = append([][]byte{nil}, prefixes...)
+
+	for i := 0; i < len(prefixes)-1; i++ {
+		key := prefixes[i]
+		if len(key) % 2 != 0 {
+			key = append(key, 0)
+		}
+		it := tree.NodeIterator(HexToKeyBytes(key))
+		callback(NewPrefixBoundIterator(it, prefixes[i+1]))
+	}
+}

pkg/iterator/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
+}