Skip to content

Commit 54568bc

Browse files
trie: add sub-trie iterator with prefix and range iteration support
1 parent a62abab commit 54568bc

File tree

3 files changed

+59
-58
lines changed

3 files changed

+59
-58
lines changed

trie/iterator.go

Lines changed: 30 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -150,11 +150,11 @@ type nodeIterator struct {
150150
// Fields for subtree iteration (original byte keys)
151151
startKey []byte // Start key for subtree iteration (nil for full trie)
152152
stopKey []byte // Stop key for subtree iteration (nil for full trie)
153-
153+
154154
// Precomputed nibble paths for efficient comparison
155155
startPath []byte // Precomputed hex path for startKey (without terminator)
156156
stopPath []byte // Precomputed hex path for stopKey (without terminator)
157-
157+
158158
// Iteration mode
159159
prefixMode bool // True if this is prefix iteration (use HasPrefix check)
160160
}
@@ -309,30 +309,30 @@ func (it *nodeIterator) Next(descend bool) bool {
309309
}
310310

311311
// Check if we're still within the subtree boundaries using precomputed paths
312-
if it.startPath != nil && len(path) > 0 {
313-
if it.prefixMode {
314-
// For prefix iteration, use HasPrefix to ensure we stay within the prefix
315-
if !bytes.HasPrefix(path, it.startPath) {
316-
it.err = errIteratorEnd
317-
return false
318-
}
319-
} else {
320-
// For range iteration, ensure we don't return nodes before the lower bound.
321-
// Advance the iterator until we reach a node at or after startPath.
322-
for bytes.Compare(path, it.startPath) < 0 {
323-
// Progress the iterator by pushing the current candidate, then peeking again.
324-
it.push(state, parentIndex, path)
325-
state, parentIndex, path, err = it.peek(descend)
326-
it.err = err
327-
if it.err != nil {
328-
return false
329-
}
330-
if len(path) == 0 {
331-
break
332-
}
333-
}
334-
}
335-
}
312+
if it.startPath != nil && len(path) > 0 {
313+
if it.prefixMode {
314+
// For prefix iteration, use HasPrefix to ensure we stay within the prefix
315+
if !bytes.HasPrefix(path, it.startPath) {
316+
it.err = errIteratorEnd
317+
return false
318+
}
319+
} else {
320+
// For range iteration, ensure we don't return nodes before the lower bound.
321+
// Advance the iterator until we reach a node at or after startPath.
322+
for bytes.Compare(path, it.startPath) < 0 {
323+
// Progress the iterator by pushing the current candidate, then peeking again.
324+
it.push(state, parentIndex, path)
325+
state, parentIndex, path, err = it.peek(descend)
326+
it.err = err
327+
if it.err != nil {
328+
return false
329+
}
330+
if len(path) == 0 {
331+
break
332+
}
333+
}
334+
}
335+
}
336336
if it.stopPath != nil && len(path) > 0 {
337337
if bytes.Compare(path, it.stopPath) >= 0 {
338338
it.err = errIteratorEnd
@@ -883,13 +883,13 @@ func (it *unionIterator) Error() error {
883883

884884
// NewSubtreeIterator creates an iterator that only traverses nodes within a subtree
885885
// defined by the given startKey and stopKey. This supports general range iteration
886-
// where startKey is inclusive and stopKey is exclusive.
886+
// where startKey is inclusive and stopKey is exclusive.
887887
//
888888
// The iterator will only visit nodes whose keys k satisfy: startKey <= k < stopKey,
889889
// where comparisons are performed in lexicographic order of byte keys (internally
890890
// implemented via hex-nibble path comparisons for efficiency).
891891
//
892-
// If startKey is nil, iteration starts from the beginning. If stopKey is nil,
892+
// If startKey is nil, iteration starts from the beginning. If stopKey is nil,
893893
// iteration continues to the end of the trie. For prefix iteration, use the
894894
// Trie.NodeIteratorWithPrefix method which handles prefix semantics correctly.
895895
func NewSubtreeIterator(trie *Trie, startKey, stopKey []byte) NodeIterator {
@@ -912,7 +912,7 @@ func nextKey(prefix []byte) []byte {
912912
// Make a copy to avoid modifying the original
913913
next := make([]byte, len(prefix))
914914
copy(next, prefix)
915-
915+
916916
// Increment the last byte that isn't 0xff
917917
for i := len(next) - 1; i >= 0; i-- {
918918
if next[i] < 0xff {
@@ -942,7 +942,7 @@ func newSubtreeIterator(trie *Trie, startKey, stopKey []byte, prefixMode bool) N
942942
stopPath = stopPath[:len(stopPath)-1]
943943
}
944944
}
945-
945+
946946
if trie.Hash() == types.EmptyRootHash {
947947
return &nodeIterator{
948948
trie: trie,

trie/iterator_test.go

Lines changed: 27 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -786,13 +786,13 @@ func TestPrefixIteratorEdgeCases(t *testing.T) {
786786
// Create a trie with test data
787787
trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
788788
testData := map[string]string{
789-
"abc": "value1",
790-
"abcd": "value2",
791-
"abce": "value3",
792-
"abd": "value4",
793-
"dog": "value5",
794-
"dog\xff": "value6", // Test with 0xff byte
795-
"dog\xff\xff": "value7", // Multiple 0xff bytes
789+
"abc": "value1",
790+
"abcd": "value2",
791+
"abce": "value3",
792+
"abd": "value4",
793+
"dog": "value5",
794+
"dog\xff": "value6", // Test with 0xff byte
795+
"dog\xff\xff": "value7", // Multiple 0xff bytes
796796
}
797797
for key, value := range testData {
798798
trie.Update([]byte(key), []byte(value))
@@ -863,7 +863,7 @@ func TestPrefixIteratorEdgeCases(t *testing.T) {
863863
allFF := []byte{0xff, 0xff}
864864
trie.Update(allFF, []byte("all_ff_value"))
865865
trie.Update(append(allFF, 0x00), []byte("all_ff_plus"))
866-
866+
867867
iter, err := trie.NodeIteratorWithPrefix(allFF)
868868
if err != nil {
869869
t.Fatalf("Failed to create iterator: %v", err)
@@ -905,13 +905,13 @@ func TestGeneralRangeIteration(t *testing.T) {
905905
// Create a trie with test data
906906
trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
907907
testData := map[string]string{
908-
"apple": "fruit1",
908+
"apple": "fruit1",
909909
"apricot": "fruit2",
910-
"banana": "fruit3",
911-
"cherry": "fruit4",
912-
"date": "fruit5",
913-
"fig": "fruit6",
914-
"grape": "fruit7",
910+
"banana": "fruit3",
911+
"cherry": "fruit4",
912+
"date": "fruit5",
913+
"fig": "fruit6",
914+
"grape": "fruit7",
915915
}
916916
for key, value := range testData {
917917
trie.Update([]byte(key), []byte(value))
@@ -977,12 +977,12 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
977977
// Create a trie with nested structure
978978
trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
979979
testData := map[string]string{
980-
"a": "value_a",
981-
"a/b": "value_ab",
982-
"a/b/c": "value_abc",
983-
"a/b/d": "value_abd",
984-
"a/e": "value_ae",
985-
"b": "value_b",
980+
"a": "value_a",
981+
"a/b": "value_ab",
982+
"a/b/c": "value_abc",
983+
"a/b/d": "value_abd",
984+
"a/e": "value_ae",
985+
"b": "value_b",
986986
}
987987
for key, value := range testData {
988988
trie.Update([]byte(key), []byte(value))
@@ -994,17 +994,17 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
994994
if err != nil {
995995
t.Fatalf("Failed to create iterator: %v", err)
996996
}
997-
997+
998998
// Count nodes at each level
999999
nodesVisited := 0
10001000
leafsFound := make(map[string]bool)
1001-
1001+
10021002
// First call with descend=true to enter the "a" subtree
10031003
if !iter.Next(true) {
10041004
t.Fatal("Expected to find at least one node")
10051005
}
10061006
nodesVisited++
1007-
1007+
10081008
// Continue iteration, sometimes with descend=false
10091009
descendPattern := []bool{false, true, false, true, true}
10101010
for i := 0; iter.Next(descendPattern[i%len(descendPattern)]); i++ {
@@ -1013,14 +1013,15 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
10131013
leafsFound[string(iter.LeafKey())] = true
10141014
}
10151015
}
1016-
1016+
10171017
// We should still respect the prefix boundary even when skipping
1018+
prefix := []byte("a")
10181019
for key := range leafsFound {
1019-
if !bytes.HasPrefix([]byte(key), []byte("a")) {
1020+
if !bytes.HasPrefix([]byte(key), prefix) {
10201021
t.Errorf("Found key outside prefix when using descend=false: %s", key)
10211022
}
10221023
}
1023-
1024+
10241025
// Should not have found "b" even if we skip some subtrees
10251026
if leafsFound["b"] {
10261027
t.Error("Iterator leaked outside prefix boundary with descend=false")

trie/trie.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,10 +136,10 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) {
136136

137137
// NodeIteratorWithPrefix returns an iterator that returns nodes of the trie whose leaf keys
138138
// start with the given prefix. Iteration includes all keys k where prefix <= k < nextKey(prefix),
139-
// effectively returning only keys that have the prefix. The iteration stops once it would
139+
// effectively returning only keys that have the prefix. The iteration stops once it would
140140
// encounter a key that doesn't start with the prefix.
141141
//
142-
// For example, with prefix "dog", the iterator will return "dog", "dogcat", "dogfish" but
142+
// For example, with prefix "dog", the iterator will return "dog", "dogcat", "dogfish" but
143143
// not "dot" or "fog". An empty prefix iterates the entire trie.
144144
func (t *Trie) NodeIteratorWithPrefix(prefix []byte) (NodeIterator, error) {
145145
// Short circuit if the trie is already committed and not usable.

0 commit comments

Comments
 (0)