Skip to content

Commit fda09c7

Browse files
samuelarogbonlogballetrjl493456442
authored
trie: add sub-trie iterator support (#32520)
- Adds `NodeIteratorWithPrefix()` method to support iterating only nodes within a specific key prefix - Adds `NodeIteratorWithRange()` method to support iterating only nodes within a specific key range Current `NodeIterator` always traverses the entire remaining trie from a start position. For non-ethereum applications using the trie implementation, there's no way to limit iteration to just a subtree with a specific prefix. **Usage:** ```go // Only iterate nodes with prefix "key1" iter, err := trie.NodeIteratorWithPrefix([]byte("key1")) ``` Testing: Comprehensive test suite covering edge cases and boundary conditions. Closes #32484 --------- Co-authored-by: gballet <[email protected]> Co-authored-by: Gary Rong <[email protected]>
1 parent 21769f3 commit fda09c7

File tree

3 files changed

+652
-0
lines changed

3 files changed

+652
-0
lines changed

trie/iterator.go

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -836,3 +836,91 @@ func (it *unionIterator) Error() error {
836836
}
837837
return nil
838838
}
839+
840+
// subTreeIterator wraps nodeIterator to traverse a trie within a predefined
841+
// start and limit range.
842+
type subtreeIterator struct {
843+
NodeIterator
844+
845+
stopPath []byte // Precomputed hex path for stopKey (without terminator), nil means no limit
846+
exhausted bool // Flag whether the iterator has been exhausted
847+
}
848+
849+
// newSubtreeIterator creates an iterator that only traverses nodes within a subtree
850+
// defined by the given startKey and stopKey. This supports general range iteration
851+
// where startKey is inclusive and stopKey is exclusive.
852+
//
853+
// The iterator will only visit nodes whose keys k satisfy: startKey <= k < stopKey,
854+
// where comparisons are performed in lexicographic order of byte keys (internally
855+
// implemented via hex-nibble path comparisons for efficiency).
856+
//
857+
// If startKey is nil, iteration starts from the beginning. If stopKey is nil,
858+
// iteration continues to the end of the trie.
859+
func newSubtreeIterator(trie *Trie, startKey, stopKey []byte) (NodeIterator, error) {
860+
it, err := trie.NodeIterator(startKey)
861+
if err != nil {
862+
return nil, err
863+
}
864+
if startKey == nil && stopKey == nil {
865+
return it, nil
866+
}
867+
// Precompute nibble paths for efficient comparison
868+
var stopPath []byte
869+
if stopKey != nil {
870+
stopPath = keybytesToHex(stopKey)
871+
if hasTerm(stopPath) {
872+
stopPath = stopPath[:len(stopPath)-1]
873+
}
874+
}
875+
return &subtreeIterator{
876+
NodeIterator: it,
877+
stopPath: stopPath,
878+
}, nil
879+
}
880+
881+
// nextKey returns the next possible key after the given prefix.
882+
// For example, "abc" -> "abd", "ab\xff" -> "ac", etc.
883+
func nextKey(prefix []byte) []byte {
884+
if len(prefix) == 0 {
885+
return nil
886+
}
887+
// Make a copy to avoid modifying the original
888+
next := make([]byte, len(prefix))
889+
copy(next, prefix)
890+
891+
// Increment the last byte that isn't 0xff
892+
for i := len(next) - 1; i >= 0; i-- {
893+
if next[i] < 0xff {
894+
next[i]++
895+
return next
896+
}
897+
// If it's 0xff, we need to carry over
898+
// Trim trailing 0xff bytes
899+
next = next[:i]
900+
}
901+
// If all bytes were 0xff, return nil (no upper bound)
902+
return nil
903+
}
904+
905+
// newPrefixIterator creates an iterator that only traverses nodes with the given prefix.
906+
// This ensures that only keys starting with the prefix are visited.
907+
func newPrefixIterator(trie *Trie, prefix []byte) (NodeIterator, error) {
908+
return newSubtreeIterator(trie, prefix, nextKey(prefix))
909+
}
910+
911+
// Next moves the iterator to the next node. If the parameter is false, any child
912+
// nodes will be skipped.
913+
func (it *subtreeIterator) Next(descend bool) bool {
914+
if it.exhausted {
915+
return false
916+
}
917+
if !it.NodeIterator.Next(descend) {
918+
it.exhausted = true
919+
return false
920+
}
921+
if it.stopPath != nil && reachedPath(it.NodeIterator.Path(), it.stopPath) {
922+
it.exhausted = true
923+
return false
924+
}
925+
return true
926+
}

0 commit comments

Comments
 (0)