From 399247b4d995e9fd9b9d1abf1f90ebadac0be578 Mon Sep 17 00:00:00 2001
From: samuel <sbayo971@gmail.com>
Date: Sat, 30 Aug 2025 02:11:40 +0100
Subject: [PATCH] trie: add sub-trie iterator support #32484

---
 trie/iterator.go             |  38 ++++++++++++
 trie/prefix_iterator_test.go | 114 +++++++++++++++++++++++++++++++++++
 trie/trie.go                 |  13 ++++
 3 files changed, 165 insertions(+)
 create mode 100644 trie/prefix_iterator_test.go

diff --git a/trie/iterator.go b/trie/iterator.go
index 80298ce48f..625182dc45 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -836,3 +836,41 @@ 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
+}
diff --git a/trie/prefix_iterator_test.go b/trie/prefix_iterator_test.go
new file mode 100644
index 0000000000..4885073298
--- /dev/null
+++ b/trie/prefix_iterator_test.go
@@ -0,0 +1,114 @@
+package trie
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/core/rawdb"
+)
+
+func TestPrefixIterator(t *testing.T) {
+	// Create a new trie
+	trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
+
+	// Insert test data
+	testData := map[string]string{
+		"key1":      "value1",
+		"key2":      "value2",
+		"key10":     "value10",
+		"key11":     "value11",
+		"different": "value_different",
+	}
+
+	for key, value := range testData {
+		trie.Update([]byte(key), []byte(value))
+	}
+
+	// Test prefix iteration for "key1" prefix
+	prefix := []byte("key1")
+	iter, err := trie.NodeIteratorWithPrefix(prefix)
+	if err != nil {
+		t.Fatalf("Failed to create prefix iterator: %v", err)
+	}
+
+	var foundKeys [][]byte
+	for iter.Next(true) {
+		if iter.Leaf() {
+			foundKeys = append(foundKeys, iter.LeafKey())
+		}
+	}
+
+	if err := iter.Error(); err != nil {
+		t.Fatalf("Iterator error: %v", err)
+	}
+
+	// Verify only keys starting with "key1" were found
+	expectedCount := 3 // "key1", "key10", "key11"
+	if len(foundKeys) != expectedCount {
+		t.Errorf("Expected %d keys, found %d", expectedCount, len(foundKeys))
+	}
+
+	for _, key := range foundKeys {
+		keyStr := string(key)
+		if !bytes.HasPrefix(key, prefix) {
+			t.Errorf("Found key %s doesn't have prefix %s", keyStr, string(prefix))
+		}
+	}
+}
+
+func TestPrefixIteratorVsFullIterator(t *testing.T) {
+	// Create a new trie with more structured data
+	trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
+
+	// Insert structured test data
+	testData := map[string]string{
+		"aaa": "value_aaa",
+		"aab": "value_aab",
+		"aba": "value_aba",
+		"bbb": "value_bbb",
+	}
+
+	for key, value := range testData {
+		trie.Update([]byte(key), []byte(value))
+	}
+
+	// Test that prefix iterator stops at boundary
+	prefix := []byte("aa")
+	prefixIter, err := trie.NodeIteratorWithPrefix(prefix)
+	if err != nil {
+		t.Fatalf("Failed to create prefix iterator: %v", err)
+	}
+
+	var prefixKeys [][]byte
+	for prefixIter.Next(true) {
+		if prefixIter.Leaf() {
+			prefixKeys = append(prefixKeys, prefixIter.LeafKey())
+		}
+	}
+
+	// Should only find "aaa" and "aab", not "aba" or "bbb"
+	if len(prefixKeys) != 2 {
+		t.Errorf("Expected 2 keys with prefix 'aa', found %d", len(prefixKeys))
+	}
+
+	// Verify no keys outside prefix were found
+	for _, key := range prefixKeys {
+		if !bytes.HasPrefix(key, prefix) {
+			t.Errorf("Prefix iterator returned key %s outside prefix %s", string(key), string(prefix))
+		}
+	}
+}
+
+func TestEmptyPrefixIterator(t *testing.T) {
+	// Test with empty trie
+	trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
+
+	iter, err := trie.NodeIteratorWithPrefix([]byte("nonexistent"))
+	if err != nil {
+		t.Fatalf("Failed to create iterator: %v", err)
+	}
+
+	if iter.Next(true) {
+		t.Error("Expected no results from empty trie")
+	}
+}
diff --git a/trie/trie.go b/trie/trie.go
index 630462f8ca..51e42ca3ee 100644
--- a/trie/trie.go
+++ b/trie/trie.go
@@ -134,6 +134,19 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) {
 	return newNodeIterator(t, start), nil
 }
 
+// NodeIteratorWithPrefix returns an iterator that returns nodes of the trie with a specific prefix.
+// Iteration starts at the key after the given prefix and stops when leaving the subtree.
+func (t *Trie) NodeIteratorWithPrefix(prefix []byte) (NodeIterator, error) {
+	// Short circuit if the trie is already committed and not usable.
+	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
+}
+
 // MustGet is a wrapper of Get and will omit any encountered error but just
 // print out an error message.
 func (t *Trie) MustGet(key []byte) []byte {