From 45ee2fbe19c2170f47b0659af72a9d37c7f9494f Mon Sep 17 00:00:00 2001 From: Gary Rong Date: Mon, 15 Sep 2025 10:19:40 +0800 Subject: [PATCH] trie: rewrite the subtreeIterator --- trie/iterator.go | 155 ++++++++++++---------------------- trie/iterator_test.go | 189 +++++++++++++++++++++++++++++++++++------- trie/trie.go | 26 ++++-- 3 files changed, 231 insertions(+), 139 deletions(-) diff --git a/trie/iterator.go b/trie/iterator.go index c63248a89c..ca382c0243 100644 --- a/trie/iterator.go +++ b/trie/iterator.go @@ -146,17 +146,6 @@ type nodeIterator struct { resolver NodeResolver // optional node resolver for avoiding disk hits pool []*nodeIteratorState // local pool for iterator states - - // Fields for subtree iteration (original byte keys) - startKey []byte // Start key for subtree iteration (nil for full trie) - stopKey []byte // Stop key for subtree iteration (nil for full trie) - - // Precomputed nibble paths for efficient comparison - startPath []byte // Precomputed hex path for startKey (without terminator) - stopPath []byte // Precomputed hex path for stopKey (without terminator) - - // Iteration mode - prefixMode bool // True if this is prefix iteration (use HasPrefix check) } // errIteratorEnd is stored in nodeIterator.err when iteration is done. @@ -307,39 +296,6 @@ func (it *nodeIterator) Next(descend bool) bool { if it.err != nil { return false } - - // Check if we're still within the subtree boundaries using precomputed paths - if it.startPath != nil && len(path) > 0 { - if it.prefixMode { - // For prefix iteration, use HasPrefix to ensure we stay within the prefix - if !bytes.HasPrefix(path, it.startPath) { - it.err = errIteratorEnd - return false - } - } else { - // For range iteration, ensure we don't return nodes before the lower bound. - // Advance the iterator until we reach a node at or after startPath. - for bytes.Compare(path, it.startPath) < 0 { - // Progress the iterator by pushing the current candidate, then peeking again. - it.push(state, parentIndex, path) - state, parentIndex, path, err = it.peek(descend) - it.err = err - if it.err != nil { - return false - } - if len(path) == 0 { - break - } - } - } - } - if it.stopPath != nil && len(path) > 0 { - if bytes.Compare(path, it.stopPath) >= 0 { - it.err = errIteratorEnd - return false - } - } - it.push(state, parentIndex, path) return true } @@ -881,7 +837,16 @@ func (it *unionIterator) Error() error { return nil } -// NewSubtreeIterator creates an iterator that only traverses nodes within a subtree +// subTreeIterator wraps nodeIterator to traverse a trie within a predefined +// start and limit range, with optional prefix mode. +type subtreeIterator struct { + NodeIterator + + stopPath []byte // Precomputed hex path for stopKey (without terminator), nil means no limit + exhausted bool // Flag whether the iterator has been exhausted +} + +// newSubtreeIterator creates an iterator that only traverses nodes within a subtree // defined by the given startKey and stopKey. This supports general range iteration // where startKey is inclusive and stopKey is exclusive. // @@ -890,17 +855,27 @@ func (it *unionIterator) Error() error { // implemented via hex-nibble path comparisons for efficiency). // // If startKey is nil, iteration starts from the beginning. If stopKey is nil, -// iteration continues to the end of the trie. For prefix iteration, use the -// Trie.NodeIteratorWithPrefix method which handles prefix semantics correctly. -func NewSubtreeIterator(trie *Trie, startKey, stopKey []byte) NodeIterator { - return newSubtreeIterator(trie, startKey, stopKey, false) -} - -// newPrefixIterator creates an iterator that only traverses nodes with the given prefix. -// This ensures that only keys starting with the prefix are visited. -func newPrefixIterator(trie *Trie, prefix []byte) NodeIterator { - stopKey := nextKey(prefix) - return newSubtreeIterator(trie, prefix, stopKey, true) +// iteration continues to the end of the trie. +func newSubtreeIterator(trie *Trie, startKey, stopKey []byte) (NodeIterator, error) { + it, err := trie.NodeIterator(startKey) + if err != nil { + return nil, err + } + if startKey == nil && stopKey == nil { + return it, nil + } + // Precompute nibble paths for efficient comparison + var stopPath []byte + if stopKey != nil { + stopPath = keybytesToHex(stopKey) + if hasTerm(stopPath) { + stopPath = stopPath[:len(stopPath)-1] + } + } + return &subtreeIterator{ + NodeIterator: it, + stopPath: stopPath, + }, nil } // nextKey returns the next possible key after the given prefix. @@ -927,51 +902,25 @@ func nextKey(prefix []byte) []byte { return nil } -func newSubtreeIterator(trie *Trie, startKey, stopKey []byte, prefixMode bool) NodeIterator { - // Precompute nibble paths for efficient comparison - var startPath, stopPath []byte - if startKey != nil { - startPath = keybytesToHex(startKey) - if hasTerm(startPath) { - startPath = startPath[:len(startPath)-1] - } - } - if stopKey != nil { - stopPath = keybytesToHex(stopKey) - if hasTerm(stopPath) { - stopPath = stopPath[:len(stopPath)-1] - } - } - - if trie.Hash() == types.EmptyRootHash { - return &nodeIterator{ - trie: trie, - err: errIteratorEnd, - startKey: startKey, - stopKey: stopKey, - startPath: startPath, - stopPath: stopPath, - prefixMode: prefixMode, - } - } - it := &nodeIterator{ - trie: trie, - startKey: startKey, - stopKey: stopKey, - startPath: startPath, - stopPath: stopPath, - prefixMode: prefixMode, - } - // Seek to the starting position if startKey is provided - if startKey != nil && len(startKey) > 0 { - it.err = it.seek(startKey) - } else { - state, err := it.init() - if err != nil { - it.err = err - } else { - it.push(state, nil, nil) - } - } - return it +// newPrefixIterator creates an iterator that only traverses nodes with the given prefix. +// This ensures that only keys starting with the prefix are visited. +func newPrefixIterator(trie *Trie, prefix []byte) (NodeIterator, error) { + return newSubtreeIterator(trie, prefix, nextKey(prefix)) +} + +// Next moves the iterator to the next node. If the parameter is false, any child +// nodes will be skipped. +func (it *subtreeIterator) Next(descend bool) bool { + if it.exhausted { + return false + } + if !it.NodeIterator.Next(descend) { + it.exhausted = true + return false + } + if it.stopPath != nil && reachedPath(it.NodeIterator.Path(), it.stopPath) { + it.exhausted = true + return false + } + return true } diff --git a/trie/iterator_test.go b/trie/iterator_test.go index eb7b06f3e7..f1451cef90 100644 --- a/trie/iterator_test.go +++ b/trie/iterator_test.go @@ -19,7 +19,9 @@ package trie import ( "bytes" "fmt" + "maps" "math/rand" + "slices" "testing" "github.com/ethereum/go-ethereum/common" @@ -625,16 +627,22 @@ func isTrieNode(scheme string, key, val []byte) (bool, []byte, common.Hash) { } func TestSubtreeIterator(t *testing.T) { - diskDb := rawdb.NewMemoryDatabase() - db := newTestDatabase(diskDb, rawdb.HashScheme) - tr := NewEmpty(db) + var ( + db = newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme) + tr = NewEmpty(db) + ) vals := []struct{ k, v string }{ {"do", "verb"}, {"dog", "puppy"}, {"doge", "coin"}, + {"dog\xff", "value6"}, + {"dog\xff\xff", "value7"}, {"horse", "stallion"}, {"house", "building"}, {"houses", "multiple"}, + {"xyz", "value"}, + {"xyz\xff", "value"}, + {"xyz\xff\xff", "value"}, } all := make(map[string]string) for _, val := range vals { @@ -644,33 +652,156 @@ func TestSubtreeIterator(t *testing.T) { root, nodes := tr.Commit(false) db.Update(root, types.EmptyRootHash, trienode.NewWithNodeSet(nodes)) - // Test prefix iteration using NewPrefixIterator - prefix := []byte("do") - - // We need to re-open the trie from the committed state + allNodes := make(map[string][]byte) tr, _ = New(TrieID(root), db) - it := newPrefixIterator(tr, prefix) - - found := make(map[string]string) + it, err := tr.NodeIterator(nil) + if err != nil { + t.Fatal(err) + } for it.Next(true) { - if it.Leaf() { - found[string(it.LeafKey())] = string(it.LeafBlob()) + allNodes[string(it.Path())] = it.NodeBlob() + } + allKeys := slices.Collect(maps.Keys(all)) + + suites := []struct { + start []byte + end []byte + expected []string + }{ + // entire key range + { + start: nil, + end: nil, + expected: allKeys, + }, + { + start: nil, + end: bytes.Repeat([]byte{0xff}, 32), + expected: allKeys, + }, + { + start: bytes.Repeat([]byte{0x0}, 32), + end: bytes.Repeat([]byte{0xff}, 32), + expected: allKeys, + }, + // key range with start + { + start: []byte("do"), + end: nil, + expected: allKeys, + }, + { + start: []byte("doe"), + end: nil, + expected: allKeys[1:], + }, + { + start: []byte("dog"), + end: nil, + expected: allKeys[1:], + }, + { + start: []byte("doge"), + end: nil, + expected: allKeys[2:], + }, + { + start: []byte("dog\xff"), + end: nil, + expected: allKeys[3:], + }, + { + start: []byte("dog\xff\xff"), + end: nil, + expected: allKeys[4:], + }, + { + start: []byte("dog\xff\xff\xff"), + end: nil, + expected: allKeys[5:], + }, + // key range with limit + { + start: nil, + end: []byte("xyz"), + expected: allKeys[:len(allKeys)-3], + }, + { + start: nil, + end: []byte("xyz\xff"), + expected: allKeys[:len(allKeys)-2], + }, + { + start: nil, + end: []byte("xyz\xff\xff"), + expected: allKeys[:len(allKeys)-1], + }, + { + start: nil, + end: []byte("xyz\xff\xff\xff"), + expected: allKeys, + }, + } + for _, suite := range suites { + // We need to re-open the trie from the committed state + tr, _ = New(TrieID(root), db) + it, err := newSubtreeIterator(tr, suite.start, suite.end) + if err != nil { + t.Fatal(err) } - } - // Should find "do", "dog", "doge" but not "horse", "house", "houses" - expected := map[string]string{ - "do": "verb", - "dog": "puppy", - "doge": "coin", - } + found := make(map[string]string) + for it.Next(true) { + if it.Leaf() { + found[string(it.LeafKey())] = string(it.LeafBlob()) + } + } + if len(found) != len(suite.expected) { + t.Errorf("wrong number of values: got %d, want %d", len(found), len(suite.expected)) + } + for k, v := range found { + if all[k] != v { + t.Errorf("wrong value for %s: got %s, want %s", k, found[k], all[k]) + } + } - if len(found) != len(expected) { - t.Errorf("wrong number of values: got %d, want %d", len(found), len(expected)) - } - for k, v := range expected { - if found[k] != v { - t.Errorf("wrong value for %s: got %s, want %s", k, found[k], v) + expectedNodes := make(map[string][]byte) + for path, blob := range allNodes { + if suite.start != nil { + hexStart := keybytesToHex(suite.start) + hexStart = hexStart[:len(hexStart)-1] + if !reachedPath([]byte(path), hexStart) { + continue + } + } + if suite.end != nil { + hexEnd := keybytesToHex(suite.end) + hexEnd = hexEnd[:len(hexEnd)-1] + if reachedPath([]byte(path), hexEnd) { + continue + } + } + expectedNodes[path] = bytes.Clone(blob) + } + + // Compare the result yield from the subtree iterator + var ( + subCount int + subIt, _ = newSubtreeIterator(tr, suite.start, suite.end) + ) + for subIt.Next(true) { + blob, ok := expectedNodes[string(subIt.Path())] + if !ok { + t.Errorf("Unexpected node iterated, path: %v", subIt.Path()) + } + subCount++ + + if !bytes.Equal(blob, subIt.NodeBlob()) { + t.Errorf("Unexpected node blob, path: %v, want: %v, got: %v", subIt.Path(), blob, subIt.NodeBlob()) + } + } + if subCount != len(expectedNodes) { + t.Errorf("Unexpected node being iterated, want: %d, got: %d", len(expectedNodes), subCount) } } } @@ -875,7 +1006,7 @@ func TestPrefixIteratorEdgeCases(t *testing.T) { } } // Should find at least the allFF key itself - if count < 1 { + if count != 2 { t.Errorf("Expected at least 1 result for all-0xff prefix, got %d", count) } }) @@ -919,7 +1050,7 @@ func TestGeneralRangeIteration(t *testing.T) { // Test range iteration from "banana" to "fig" (exclusive) t.Run("RangeIteration", func(t *testing.T) { - iter := NewSubtreeIterator(trie, []byte("banana"), []byte("fig")) + iter, _ := newSubtreeIterator(trie, []byte("banana"), []byte("fig")) found := make(map[string]bool) for iter.Next(true) { if iter.Leaf() { @@ -937,7 +1068,7 @@ func TestGeneralRangeIteration(t *testing.T) { // Test with nil stopKey (iterate to end) t.Run("NilStopKey", func(t *testing.T) { - iter := NewSubtreeIterator(trie, []byte("date"), nil) + iter, _ := newSubtreeIterator(trie, []byte("date"), nil) found := make(map[string]bool) for iter.Next(true) { if iter.Leaf() { @@ -955,7 +1086,7 @@ func TestGeneralRangeIteration(t *testing.T) { // Test with nil startKey (iterate from beginning) t.Run("NilStartKey", func(t *testing.T) { - iter := NewSubtreeIterator(trie, nil, []byte("cherry")) + iter, _ := newSubtreeIterator(trie, nil, []byte("cherry")) found := make(map[string]bool) for iter.Next(true) { if iter.Leaf() { diff --git a/trie/trie.go b/trie/trie.go index c5be856241..36cc732ee8 100644 --- a/trie/trie.go +++ b/trie/trie.go @@ -134,20 +134,32 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) { return newNodeIterator(t, start), nil } -// NodeIteratorWithPrefix returns an iterator that returns nodes of the trie whose leaf keys -// start with the given prefix. Iteration includes all keys k where prefix <= k < nextKey(prefix), -// effectively returning only keys that have the prefix. The iteration stops once it would -// encounter a key that doesn't start with the prefix. +// NodeIteratorWithPrefix returns an iterator that returns nodes of the trie +// whose leaf keys start with the given prefix. Iteration includes all keys +// where prefix <= k < nextKey(prefix), effectively returning only keys that +// have the prefix. The iteration stops once it would encounter a key that +// doesn't start with the prefix. // -// For example, with prefix "dog", the iterator will return "dog", "dogcat", "dogfish" but -// not "dot" or "fog". An empty prefix iterates the entire trie. +// For example, with prefix "dog", the iterator will return "dog", "dogcat", +// "dogfish" but not "dot" or "fog". An empty prefix iterates the entire trie. func (t *Trie) NodeIteratorWithPrefix(prefix []byte) (NodeIterator, error) { // Short circuit if the trie is already committed and not usable. if t.committed { return nil, ErrCommitted } // Use the dedicated prefix iterator which handles prefix checking correctly - return newPrefixIterator(t, prefix), nil + return newPrefixIterator(t, prefix) +} + +// NodeIteratorWithRange returns an iterator over trie nodes whose leaf keys +// fall within the specified range. It includes all keys where start <= k < end. +// Iteration stops once a key beyond the end boundary is encountered. +func (t *Trie) NodeIteratorWithRange(start, end []byte) (NodeIterator, error) { + // Short circuit if the trie is already committed and not usable. + if t.committed { + return nil, ErrCommitted + } + return newSubtreeIterator(t, start, end) } // MustGet is a wrapper of Get and will omit any encountered error but just