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
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

View file

@ -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))
@ -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))
@ -1015,8 +1015,9 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
}
// 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)
}
}