diff --git a/trie/iterator.go b/trie/iterator.go index 5ef63817a3..c63248a89c 100644 --- a/trie/iterator.go +++ b/trie/iterator.go @@ -150,11 +150,11 @@ type nodeIterator struct { // 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) } @@ -309,30 +309,30 @@ func (it *nodeIterator) Next(descend bool) bool { } // 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.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 @@ -883,13 +883,13 @@ func (it *unionIterator) Error() error { // 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. +// where startKey is inclusive and stopKey is exclusive. // // The iterator will only visit nodes whose keys k satisfy: startKey <= k < stopKey, // where comparisons are performed in lexicographic order of byte keys (internally // implemented via hex-nibble path comparisons for efficiency). // -// If startKey is nil, iteration starts from the beginning. If stopKey is nil, +// 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 { @@ -912,7 +912,7 @@ func nextKey(prefix []byte) []byte { // Make a copy to avoid modifying the original next := make([]byte, len(prefix)) copy(next, prefix) - + // Increment the last byte that isn't 0xff for i := len(next) - 1; i >= 0; i-- { if next[i] < 0xff { @@ -942,7 +942,7 @@ func newSubtreeIterator(trie *Trie, startKey, stopKey []byte, prefixMode bool) N stopPath = stopPath[:len(stopPath)-1] } } - + if trie.Hash() == types.EmptyRootHash { return &nodeIterator{ trie: trie, diff --git a/trie/iterator_test.go b/trie/iterator_test.go index f433c266d9..eb7b06f3e7 100644 --- a/trie/iterator_test.go +++ b/trie/iterator_test.go @@ -786,13 +786,13 @@ func TestPrefixIteratorEdgeCases(t *testing.T) { // Create a trie with test data trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)) testData := map[string]string{ - "abc": "value1", - "abcd": "value2", - "abce": "value3", - "abd": "value4", - "dog": "value5", - "dog\xff": "value6", // Test with 0xff byte - "dog\xff\xff": "value7", // Multiple 0xff bytes + "abc": "value1", + "abcd": "value2", + "abce": "value3", + "abd": "value4", + "dog": "value5", + "dog\xff": "value6", // Test with 0xff byte + "dog\xff\xff": "value7", // Multiple 0xff bytes } for key, value := range testData { trie.Update([]byte(key), []byte(value)) @@ -863,7 +863,7 @@ func TestPrefixIteratorEdgeCases(t *testing.T) { allFF := []byte{0xff, 0xff} trie.Update(allFF, []byte("all_ff_value")) trie.Update(append(allFF, 0x00), []byte("all_ff_plus")) - + iter, err := trie.NodeIteratorWithPrefix(allFF) if err != nil { t.Fatalf("Failed to create iterator: %v", err) @@ -905,13 +905,13 @@ func TestGeneralRangeIteration(t *testing.T) { // Create a trie with test data trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)) testData := map[string]string{ - "apple": "fruit1", + "apple": "fruit1", "apricot": "fruit2", - "banana": "fruit3", - "cherry": "fruit4", - "date": "fruit5", - "fig": "fruit6", - "grape": "fruit7", + "banana": "fruit3", + "cherry": "fruit4", + "date": "fruit5", + "fig": "fruit6", + "grape": "fruit7", } for key, value := range testData { trie.Update([]byte(key), []byte(value)) @@ -977,12 +977,12 @@ func TestPrefixIteratorWithDescend(t *testing.T) { // Create a trie with nested structure trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)) testData := map[string]string{ - "a": "value_a", - "a/b": "value_ab", - "a/b/c": "value_abc", - "a/b/d": "value_abd", - "a/e": "value_ae", - "b": "value_b", + "a": "value_a", + "a/b": "value_ab", + "a/b/c": "value_abc", + "a/b/d": "value_abd", + "a/e": "value_ae", + "b": "value_b", } for key, value := range testData { trie.Update([]byte(key), []byte(value)) @@ -994,17 +994,17 @@ func TestPrefixIteratorWithDescend(t *testing.T) { if err != nil { t.Fatalf("Failed to create iterator: %v", err) } - + // Count nodes at each level nodesVisited := 0 leafsFound := make(map[string]bool) - + // First call with descend=true to enter the "a" subtree if !iter.Next(true) { t.Fatal("Expected to find at least one node") } nodesVisited++ - + // Continue iteration, sometimes with descend=false descendPattern := []bool{false, true, false, true, true} for i := 0; iter.Next(descendPattern[i%len(descendPattern)]); i++ { @@ -1013,14 +1013,15 @@ func TestPrefixIteratorWithDescend(t *testing.T) { leafsFound[string(iter.LeafKey())] = true } } - + // We should still respect the prefix boundary even when skipping + prefix := []byte("a") for key := range leafsFound { - if !bytes.HasPrefix([]byte(key), []byte("a")) { + if !bytes.HasPrefix([]byte(key), prefix) { t.Errorf("Found key outside prefix when using descend=false: %s", key) } } - + // Should not have found "b" even if we skip some subtrees if leafsFound["b"] { t.Error("Iterator leaked outside prefix boundary with descend=false") diff --git a/trie/trie.go b/trie/trie.go index 02c62f6cf1..c5be856241 100644 --- a/trie/trie.go +++ b/trie/trie.go @@ -136,10 +136,10 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) { // 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 +// 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 +// 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.