trie: add sub-trie iterator support #32484

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
+}

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")
+	}
+}

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 {