From 9a72c83e6f112c52aaa4c581c30eb542f85226e7 Mon Sep 17 00:00:00 2001 From: samuel Date: Sat, 30 Aug 2025 15:34:47 +0100 Subject: [PATCH] trie: refactor sub-trie iterator to extend existing nodeIterator - Remove separate prefixIterator type per reviewer feedback - Add prefix and stop fields directly to nodeIterator - Implement proper subtree boundaries with both start and stop points - Simplify maintenance by avoiding unnecessary new types Co-authored-by: gballet --- trie/iterator.go | 94 ++++++++++++++++++++++++++----------------- trie/iterator_test.go | 52 ++++++++++++++++++++++++ trie/trie.go | 6 +-- 3 files changed, 112 insertions(+), 40 deletions(-) diff --git a/trie/iterator.go b/trie/iterator.go index 625182dc45..f72b319abf 100644 --- a/trie/iterator.go +++ b/trie/iterator.go @@ -146,6 +146,10 @@ type nodeIterator struct { resolver NodeResolver // optional node resolver for avoiding disk hits pool []*nodeIteratorState // local pool for iterator states + + // Fields for subtree iteration + prefix []byte // Prefix for subtree iteration (nil for full trie) + stop []byte // Stop boundary for subtree iteration (nil for full trie) } // errIteratorEnd is stored in nodeIterator.err when iteration is done. @@ -296,6 +300,31 @@ func (it *nodeIterator) Next(descend bool) bool { if it.err != nil { return false } + + // Check if we're still within the subtree boundaries + // Note: path is already hex-encoded by the iterator + if it.prefix != nil && len(path) > 0 { + prefixHex := keybytesToHex(it.prefix) + // Remove terminator from prefix hex if present + if hasTerm(prefixHex) { + prefixHex = prefixHex[:len(prefixHex)-1] + } + if !bytes.HasPrefix(path, prefixHex) { + it.err = errIteratorEnd + return false + } + } + if it.stop != nil && len(path) > 0 { + stopHex := keybytesToHex(it.stop) + if hasTerm(stopHex) { + stopHex = stopHex[:len(stopHex)-1] + } + if bytes.Compare(path, stopHex) >= 0 { + it.err = errIteratorEnd + return false + } + } + it.push(state, parentIndex, path) return true } @@ -837,40 +866,33 @@ func (it *unionIterator) Error() error { return nil } -// prefixIterator is a wrapper around NodeIterator that stops iteration -// when it leaves a subtree with a specific prefix. -type prefixIterator struct { - NodeIterator - prefix []byte - ended bool -} - -// NewPrefixIterator creates an iterator that only traverses nodes with the given prefix. -func NewPrefixIterator(base NodeIterator, prefix []byte) NodeIterator { - return &prefixIterator{ - NodeIterator: base, - prefix: prefix, - } -} - -// Next moves the iterator to the next node within the prefix boundary. -// It returns false when no more nodes exist within the prefix. -func (pi *prefixIterator) Next(descend bool) bool { - if pi.ended { - return false - } - - if !pi.NodeIterator.Next(descend) { - pi.ended = true - return false - } - - // Check if current path is still within the prefix boundary - path := pi.Path() - if len(path) > 0 && len(pi.prefix) > 0 && !bytes.HasPrefix(path, pi.prefix) { - pi.ended = true - return false - } - - return true +// NewSubtreeIterator creates an iterator that only traverses nodes within a subtree +// defined by the given prefix and stopping point. The prefix defines where iteration +// starts, and stop defines where it ends (exclusive). +func NewSubtreeIterator(trie *Trie, prefix []byte, stop []byte) NodeIterator { + if trie.Hash() == types.EmptyRootHash { + return &nodeIterator{ + trie: trie, + err: errIteratorEnd, + prefix: prefix, + stop: stop, + } + } + it := &nodeIterator{ + trie: trie, + prefix: prefix, + stop: stop, + } + // Seek to the starting position if prefix is provided + if prefix != nil && len(prefix) > 0 { + it.err = it.seek(prefix) + } else { + state, err := it.init() + if err != nil { + it.err = err + } else { + it.push(state, nil, nil) + } + } + return it } diff --git a/trie/iterator_test.go b/trie/iterator_test.go index 74a1aa378c..83a917571d 100644 --- a/trie/iterator_test.go +++ b/trie/iterator_test.go @@ -624,6 +624,58 @@ func isTrieNode(scheme string, key, val []byte) (bool, []byte, common.Hash) { return true, path, hash } +func TestSubtreeIterator(t *testing.T) { + diskDb := rawdb.NewMemoryDatabase() + db := newTestDatabase(diskDb, rawdb.HashScheme) + tr := NewEmpty(db) + vals := []struct{ k, v string }{ + {"do", "verb"}, + {"dog", "puppy"}, + {"doge", "coin"}, + {"horse", "stallion"}, + {"house", "building"}, + {"houses", "multiple"}, + } + all := make(map[string]string) + for _, val := range vals { + all[val.k] = val.v + tr.MustUpdate([]byte(val.k), []byte(val.v)) + } + root, nodes := tr.Commit(false) + db.Update(root, types.EmptyRootHash, trienode.NewWithNodeSet(nodes)) + + // Test subtree iteration with prefix "do" + prefix := []byte("do") + stop := []byte("e") + + // We need to re-open the trie from the committed state + tr, _ = New(TrieID(root), db) + it := NewSubtreeIterator(tr, prefix, stop) + + found := make(map[string]string) + for it.Next(true) { + if it.Leaf() { + found[string(it.LeafKey())] = string(it.LeafBlob()) + } + } + + // Should find "do", "dog", "doge" but not "horse", "house", "houses" + expected := map[string]string{ + "do": "verb", + "dog": "puppy", + "doge": "coin", + } + + if len(found) != len(expected) { + t.Errorf("wrong number of values: got %d, want %d", len(found), len(expected)) + } + for k, v := range expected { + if found[k] != v { + t.Errorf("wrong value for %s: got %s, want %s", k, found[k], v) + } + } +} + func BenchmarkIterator(b *testing.B) { diskDb, srcDb, tr, _ := makeTestTrie(rawdb.HashScheme) root := tr.Hash() diff --git a/trie/trie.go b/trie/trie.go index 51e42ca3ee..2e304cc966 100644 --- a/trie/trie.go +++ b/trie/trie.go @@ -141,10 +141,8 @@ func (t *Trie) NodeIteratorWithPrefix(prefix []byte) (NodeIterator, error) { if t.committed { return nil, ErrCommitted } - baseIterator := newNodeIterator(t, prefix) - hexPrefix := keybytesToHex(prefix) - hexPrefix = hexPrefix[:len(hexPrefix)-1] // Remove terminator - return NewPrefixIterator(baseIterator, hexPrefix), nil + // Use NewSubtreeIterator with just a prefix and no stop boundary + return NewSubtreeIterator(t, prefix, nil), nil } // MustGet is a wrapper of Get and will omit any encountered error but just