trie: add sub-trie iterator with prefix and range iteration support

This commit is contained in:
samuel 2025-09-10 14:25:58 +01:00 committed by Gary Rong
parent f125dbd9c7
commit 2db2a18e83
3 changed files with 59 additions and 58 deletions

View file

@ -309,30 +309,30 @@ func (it *nodeIterator) Next(descend bool) bool {
} }
// Check if we're still within the subtree boundaries using precomputed paths // Check if we're still within the subtree boundaries using precomputed paths
if it.startPath != nil && len(path) > 0 { if it.startPath != nil && len(path) > 0 {
if it.prefixMode { if it.prefixMode {
// For prefix iteration, use HasPrefix to ensure we stay within the prefix // For prefix iteration, use HasPrefix to ensure we stay within the prefix
if !bytes.HasPrefix(path, it.startPath) { if !bytes.HasPrefix(path, it.startPath) {
it.err = errIteratorEnd it.err = errIteratorEnd
return false return false
} }
} else { } else {
// For range iteration, ensure we don't return nodes before the lower bound. // 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. // Advance the iterator until we reach a node at or after startPath.
for bytes.Compare(path, it.startPath) < 0 { for bytes.Compare(path, it.startPath) < 0 {
// Progress the iterator by pushing the current candidate, then peeking again. // Progress the iterator by pushing the current candidate, then peeking again.
it.push(state, parentIndex, path) it.push(state, parentIndex, path)
state, parentIndex, path, err = it.peek(descend) state, parentIndex, path, err = it.peek(descend)
it.err = err it.err = err
if it.err != nil { if it.err != nil {
return false return false
} }
if len(path) == 0 { if len(path) == 0 {
break break
} }
} }
} }
} }
if it.stopPath != nil && len(path) > 0 { if it.stopPath != nil && len(path) > 0 {
if bytes.Compare(path, it.stopPath) >= 0 { if bytes.Compare(path, it.stopPath) >= 0 {
it.err = errIteratorEnd it.err = errIteratorEnd

View file

@ -786,13 +786,13 @@ func TestPrefixIteratorEdgeCases(t *testing.T) {
// Create a trie with test data // Create a trie with test data
trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)) trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
testData := map[string]string{ testData := map[string]string{
"abc": "value1", "abc": "value1",
"abcd": "value2", "abcd": "value2",
"abce": "value3", "abce": "value3",
"abd": "value4", "abd": "value4",
"dog": "value5", "dog": "value5",
"dog\xff": "value6", // Test with 0xff byte "dog\xff": "value6", // Test with 0xff byte
"dog\xff\xff": "value7", // Multiple 0xff bytes "dog\xff\xff": "value7", // Multiple 0xff bytes
} }
for key, value := range testData { for key, value := range testData {
trie.Update([]byte(key), []byte(value)) trie.Update([]byte(key), []byte(value))
@ -905,13 +905,13 @@ func TestGeneralRangeIteration(t *testing.T) {
// Create a trie with test data // Create a trie with test data
trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)) trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
testData := map[string]string{ testData := map[string]string{
"apple": "fruit1", "apple": "fruit1",
"apricot": "fruit2", "apricot": "fruit2",
"banana": "fruit3", "banana": "fruit3",
"cherry": "fruit4", "cherry": "fruit4",
"date": "fruit5", "date": "fruit5",
"fig": "fruit6", "fig": "fruit6",
"grape": "fruit7", "grape": "fruit7",
} }
for key, value := range testData { for key, value := range testData {
trie.Update([]byte(key), []byte(value)) trie.Update([]byte(key), []byte(value))
@ -977,12 +977,12 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
// Create a trie with nested structure // Create a trie with nested structure
trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)) trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
testData := map[string]string{ testData := map[string]string{
"a": "value_a", "a": "value_a",
"a/b": "value_ab", "a/b": "value_ab",
"a/b/c": "value_abc", "a/b/c": "value_abc",
"a/b/d": "value_abd", "a/b/d": "value_abd",
"a/e": "value_ae", "a/e": "value_ae",
"b": "value_b", "b": "value_b",
} }
for key, value := range testData { for key, value := range testData {
trie.Update([]byte(key), []byte(value)) trie.Update([]byte(key), []byte(value))
@ -1015,8 +1015,9 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
} }
// We should still respect the prefix boundary even when skipping // We should still respect the prefix boundary even when skipping
prefix := []byte("a")
for key := range leafsFound { 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) t.Errorf("Found key outside prefix when using descend=false: %s", key)
} }
} }