Skip to content

Commit 3dfd6cd

Browse files
committed
trie: rewrite the subtreeIterator
1 parent 54568bc commit 3dfd6cd

File tree

3 files changed

+228
-136
lines changed

3 files changed

+228
-136
lines changed

trie/iterator.go

Lines changed: 49 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -146,17 +146,6 @@ type nodeIterator struct {
146146

147147
resolver NodeResolver // optional node resolver for avoiding disk hits
148148
pool []*nodeIteratorState // local pool for iterator states
149-
150-
// Fields for subtree iteration (original byte keys)
151-
startKey []byte // Start key for subtree iteration (nil for full trie)
152-
stopKey []byte // Stop key for subtree iteration (nil for full trie)
153-
154-
// Precomputed nibble paths for efficient comparison
155-
startPath []byte // Precomputed hex path for startKey (without terminator)
156-
stopPath []byte // Precomputed hex path for stopKey (without terminator)
157-
158-
// Iteration mode
159-
prefixMode bool // True if this is prefix iteration (use HasPrefix check)
160149
}
161150

162151
// errIteratorEnd is stored in nodeIterator.err when iteration is done.
@@ -307,39 +296,6 @@ func (it *nodeIterator) Next(descend bool) bool {
307296
if it.err != nil {
308297
return false
309298
}
310-
311-
// 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-
}
336-
if it.stopPath != nil && len(path) > 0 {
337-
if bytes.Compare(path, it.stopPath) >= 0 {
338-
it.err = errIteratorEnd
339-
return false
340-
}
341-
}
342-
343299
it.push(state, parentIndex, path)
344300
return true
345301
}
@@ -881,7 +837,16 @@ func (it *unionIterator) Error() error {
881837
return nil
882838
}
883839

884-
// NewSubtreeIterator creates an iterator that only traverses nodes within a subtree
840+
// subTreeIterator wraps nodeIterator to traverse a trie within a predefined
841+
// start and limit range, with optional prefix mode.
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
885850
// defined by the given startKey and stopKey. This supports general range iteration
886851
// where startKey is inclusive and stopKey is exclusive.
887852
//
@@ -890,17 +855,27 @@ func (it *unionIterator) Error() error {
890855
// implemented via hex-nibble path comparisons for efficiency).
891856
//
892857
// If startKey is nil, iteration starts from the beginning. If stopKey is nil,
893-
// iteration continues to the end of the trie. For prefix iteration, use the
894-
// Trie.NodeIteratorWithPrefix method which handles prefix semantics correctly.
895-
func NewSubtreeIterator(trie *Trie, startKey, stopKey []byte) NodeIterator {
896-
return newSubtreeIterator(trie, startKey, stopKey, false)
897-
}
898-
899-
// newPrefixIterator creates an iterator that only traverses nodes with the given prefix.
900-
// This ensures that only keys starting with the prefix are visited.
901-
func newPrefixIterator(trie *Trie, prefix []byte) NodeIterator {
902-
stopKey := nextKey(prefix)
903-
return newSubtreeIterator(trie, prefix, stopKey, true)
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
904879
}
905880

906881
// nextKey returns the next possible key after the given prefix.
@@ -927,51 +902,25 @@ func nextKey(prefix []byte) []byte {
927902
return nil
928903
}
929904

930-
func newSubtreeIterator(trie *Trie, startKey, stopKey []byte, prefixMode bool) NodeIterator {
931-
// Precompute nibble paths for efficient comparison
932-
var startPath, stopPath []byte
933-
if startKey != nil {
934-
startPath = keybytesToHex(startKey)
935-
if hasTerm(startPath) {
936-
startPath = startPath[:len(startPath)-1]
937-
}
938-
}
939-
if stopKey != nil {
940-
stopPath = keybytesToHex(stopKey)
941-
if hasTerm(stopPath) {
942-
stopPath = stopPath[:len(stopPath)-1]
943-
}
944-
}
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+
}
945910

946-
if trie.Hash() == types.EmptyRootHash {
947-
return &nodeIterator{
948-
trie: trie,
949-
err: errIteratorEnd,
950-
startKey: startKey,
951-
stopKey: stopKey,
952-
startPath: startPath,
953-
stopPath: stopPath,
954-
prefixMode: prefixMode,
955-
}
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
956916
}
957-
it := &nodeIterator{
958-
trie: trie,
959-
startKey: startKey,
960-
stopKey: stopKey,
961-
startPath: startPath,
962-
stopPath: stopPath,
963-
prefixMode: prefixMode,
964-
}
965-
// Seek to the starting position if startKey is provided
966-
if startKey != nil && len(startKey) > 0 {
967-
it.err = it.seek(startKey)
968-
} else {
969-
state, err := it.init()
970-
if err != nil {
971-
it.err = err
972-
} else {
973-
it.push(state, nil, nil)
974-
}
917+
if !it.NodeIterator.Next(descend) {
918+
it.exhausted = true
919+
return false
975920
}
976-
return it
921+
if it.stopPath != nil && reachedPath(it.NodeIterator.Path(), it.stopPath) {
922+
it.exhausted = true
923+
return false
924+
}
925+
return true
977926
}

0 commit comments

Comments
 (0)