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 <guillaume.ballet@gmail.com>

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,
+// 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,
 		}
 	}
-
-// 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
+	it := &nodeIterator{
+		trie:   trie,
+		prefix: prefix,
+		stop:   stop,
 	}
-
-	if !pi.NodeIterator.Next(descend) {
-		pi.ended = true
-		return false
+	// 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)
 		}
-
-	// 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
+	return it
 }

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

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