aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--trie/iterator.go138
-rw-r--r--trie/iterator_test.go97
2 files changed, 199 insertions, 36 deletions
diff --git a/trie/iterator.go b/trie/iterator.go
index ddc674d2b..42149a7d3 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -18,7 +18,7 @@ package trie
import (
"bytes"
-
+ "container/heap"
"github.com/ethereum/go-ethereum/common"
)
@@ -268,6 +268,26 @@ outer:
return nil
}
+func compareNodes(a, b NodeIterator) int {
+ cmp := bytes.Compare(a.Path(), b.Path())
+ if cmp != 0 {
+ return cmp
+ }
+
+ if a.Leaf() && !b.Leaf() {
+ return -1
+ } else if b.Leaf() && !a.Leaf() {
+ return 1
+ }
+
+ cmp = bytes.Compare(a.Hash().Bytes(), b.Hash().Bytes())
+ if cmp != 0 {
+ return cmp
+ }
+
+ return bytes.Compare(a.LeafBlob(), b.LeafBlob())
+}
+
type differenceIterator struct {
a, b NodeIterator // Nodes returned are those in b - a.
eof bool // Indicates a has run out of elements
@@ -321,8 +341,7 @@ func (it *differenceIterator) Next(bool) bool {
}
for {
- apath, bpath := it.a.Path(), it.b.Path()
- switch bytes.Compare(apath, bpath) {
+ switch compareNodes(it.a, it.b) {
case -1:
// b jumped past a; advance a
if !it.a.Next(true) {
@@ -334,15 +353,6 @@ func (it *differenceIterator) Next(bool) bool {
// b is before a
return true
case 0:
- if it.a.Hash() != it.b.Hash() || it.a.Leaf() != it.b.Leaf() {
- // Keys are identical, but hashes or leaf status differs
- return true
- }
- if it.a.Leaf() && it.b.Leaf() && !bytes.Equal(it.a.LeafBlob(), it.b.LeafBlob()) {
- // Both are leaf nodes, but with different values
- return true
- }
-
// a and b are identical; skip this whole subtree if the nodes have hashes
hasHash := it.a.Hash() == common.Hash{}
if !it.b.Next(hasHash) {
@@ -364,3 +374,107 @@ func (it *differenceIterator) Error() error {
}
return it.b.Error()
}
+
+type nodeIteratorHeap []NodeIterator
+
+func (h nodeIteratorHeap) Len() int { return len(h) }
+func (h nodeIteratorHeap) Less(i, j int) bool { return compareNodes(h[i], h[j]) < 0 }
+func (h nodeIteratorHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
+func (h *nodeIteratorHeap) Push(x interface{}) { *h = append(*h, x.(NodeIterator)) }
+func (h *nodeIteratorHeap) Pop() interface{} {
+ n := len(*h)
+ x := (*h)[n-1]
+ *h = (*h)[0 : n-1]
+ return x
+}
+
+type unionIterator struct {
+ items *nodeIteratorHeap // Nodes returned are the union of the ones in these iterators
+ count int // Number of nodes scanned across all tries
+ err error // The error, if one has been encountered
+}
+
+// NewUnionIterator constructs a NodeIterator that iterates over elements in the union
+// of the provided NodeIterators. Returns the iterator, and a pointer to an integer
+// recording the number of nodes visited.
+func NewUnionIterator(iters []NodeIterator) (NodeIterator, *int) {
+ h := make(nodeIteratorHeap, len(iters))
+ copy(h, iters)
+ heap.Init(&h)
+
+ ui := &unionIterator{
+ items: &h,
+ }
+ return ui, &ui.count
+}
+
+func (it *unionIterator) Hash() common.Hash {
+ return (*it.items)[0].Hash()
+}
+
+func (it *unionIterator) Parent() common.Hash {
+ return (*it.items)[0].Parent()
+}
+
+func (it *unionIterator) Leaf() bool {
+ return (*it.items)[0].Leaf()
+}
+
+func (it *unionIterator) LeafBlob() []byte {
+ return (*it.items)[0].LeafBlob()
+}
+
+func (it *unionIterator) Path() []byte {
+ return (*it.items)[0].Path()
+}
+
+// Next returns the next node in the union of tries being iterated over.
+//
+// It does this by maintaining a heap of iterators, sorted by the iteration
+// order of their next elements, with one entry for each source trie. Each
+// time Next() is called, it takes the least element from the heap to return,
+// advancing any other iterators that also point to that same element. These
+// iterators are called with descend=false, since we know that any nodes under
+// these nodes will also be duplicates, found in the currently selected iterator.
+// Whenever an iterator is advanced, it is pushed back into the heap if it still
+// has elements remaining.
+//
+// In the case that descend=false - eg, we're asked to ignore all subnodes of the
+// current node - we also advance any iterators in the heap that have the current
+// path as a prefix.
+func (it *unionIterator) Next(descend bool) bool {
+ if len(*it.items) == 0 {
+ return false
+ }
+
+ // Get the next key from the union
+ least := heap.Pop(it.items).(NodeIterator)
+
+ // Skip over other nodes as long as they're identical, or, if we're not descending, as
+ // long as they have the same prefix as the current node.
+ for len(*it.items) > 0 && ((!descend && bytes.HasPrefix((*it.items)[0].Path(), least.Path())) || compareNodes(least, (*it.items)[0]) == 0) {
+ skipped := heap.Pop(it.items).(NodeIterator)
+ // Skip the whole subtree if the nodes have hashes; otherwise just skip this node
+ if skipped.Next(skipped.Hash() == common.Hash{}) {
+ it.count += 1
+ // If there are more elements, push the iterator back on the heap
+ heap.Push(it.items, skipped)
+ }
+ }
+
+ if least.Next(descend) {
+ it.count += 1
+ heap.Push(it.items, least)
+ }
+
+ return len(*it.items) > 0
+}
+
+func (it *unionIterator) Error() error {
+ for i := 0; i < len(*it.items); i++ {
+ if err := (*it.items)[i].Error(); err != nil {
+ return err
+ }
+ }
+ return nil
+}
diff --git a/trie/iterator_test.go b/trie/iterator_test.go
index 0ad9711ed..c101bb7b0 100644
--- a/trie/iterator_test.go
+++ b/trie/iterator_test.go
@@ -117,36 +117,38 @@ func TestNodeIteratorCoverage(t *testing.T) {
}
}
+var testdata1 = []struct{ k, v string }{
+ {"bar", "b"},
+ {"barb", "ba"},
+ {"bars", "bb"},
+ {"bard", "bc"},
+ {"fab", "z"},
+ {"foo", "a"},
+ {"food", "ab"},
+ {"foos", "aa"},
+}
+
+var testdata2 = []struct{ k, v string }{
+ {"aardvark", "c"},
+ {"bar", "b"},
+ {"barb", "bd"},
+ {"bars", "be"},
+ {"fab", "z"},
+ {"foo", "a"},
+ {"foos", "aa"},
+ {"food", "ab"},
+ {"jars", "d"},
+}
+
func TestDifferenceIterator(t *testing.T) {
triea := newEmpty()
- valsa := []struct{ k, v string }{
- {"bar", "b"},
- {"barb", "ba"},
- {"bars", "bb"},
- {"bard", "bc"},
- {"fab", "z"},
- {"foo", "a"},
- {"food", "ab"},
- {"foos", "aa"},
- }
- for _, val := range valsa {
+ for _, val := range testdata1 {
triea.Update([]byte(val.k), []byte(val.v))
}
triea.Commit()
trieb := newEmpty()
- valsb := []struct{ k, v string }{
- {"aardvark", "c"},
- {"bar", "b"},
- {"barb", "bd"},
- {"bars", "be"},
- {"fab", "z"},
- {"foo", "a"},
- {"foos", "aa"},
- {"food", "ab"},
- {"jars", "d"},
- }
- for _, val := range valsb {
+ for _, val := range testdata2 {
trieb.Update([]byte(val.k), []byte(val.v))
}
trieb.Commit()
@@ -166,10 +168,57 @@ func TestDifferenceIterator(t *testing.T) {
}
for _, item := range all {
if found[item.k] != item.v {
- t.Errorf("iterator value mismatch for %s: got %q want %q", item.k, found[item.k], item.v)
+ t.Errorf("iterator value mismatch for %s: got %v want %v", item.k, found[item.k], item.v)
}
}
if len(found) != len(all) {
t.Errorf("iterator count mismatch: got %d values, want %d", len(found), len(all))
}
}
+
+func TestUnionIterator(t *testing.T) {
+ triea := newEmpty()
+ for _, val := range testdata1 {
+ triea.Update([]byte(val.k), []byte(val.v))
+ }
+ triea.Commit()
+
+ trieb := newEmpty()
+ for _, val := range testdata2 {
+ trieb.Update([]byte(val.k), []byte(val.v))
+ }
+ trieb.Commit()
+
+ di, _ := NewUnionIterator([]NodeIterator{NewNodeIterator(triea), NewNodeIterator(trieb)})
+ it := NewIteratorFromNodeIterator(di)
+
+ all := []struct{ k, v string }{
+ {"aardvark", "c"},
+ {"barb", "bd"},
+ {"barb", "ba"},
+ {"bard", "bc"},
+ {"bars", "bb"},
+ {"bars", "be"},
+ {"bar", "b"},
+ {"fab", "z"},
+ {"food", "ab"},
+ {"foos", "aa"},
+ {"foo", "a"},
+ {"jars", "d"},
+ }
+
+ for i, kv := range all {
+ if !it.Next() {
+ t.Errorf("Iterator ends prematurely at element %d", i)
+ }
+ if kv.k != string(it.Key) {
+ t.Errorf("iterator value mismatch for element %d: got key %s want %s", i, it.Key, kv.k)
+ }
+ if kv.v != string(it.Value) {
+ t.Errorf("iterator value mismatch for element %d: got value %s want %s", i, it.Value, kv.v)
+ }
+ }
+ if it.Next() {
+ t.Errorf("Iterator returned extra values.")
+ }
+}