trie: rewrite the subtreeIterator

trie/iterator.go

@@ -146,17 +146,6 @@ type nodeIterator struct {
 
 	resolver NodeResolver         // optional node resolver for avoiding disk hits
 	pool     []*nodeIteratorState // local pool for iterator states
-
-	// Fields for subtree iteration (original byte keys)
-	startKey []byte // Start key for subtree iteration (nil for full trie)
-	stopKey  []byte // Stop key for subtree iteration (nil for full trie)
-
-	// Precomputed nibble paths for efficient comparison
-	startPath []byte // Precomputed hex path for startKey (without terminator)
-	stopPath  []byte // Precomputed hex path for stopKey (without terminator)
-
-	// Iteration mode
-	prefixMode bool // True if this is prefix iteration (use HasPrefix check)
 }
 
 // errIteratorEnd is stored in nodeIterator.err when iteration is done.
@@ -307,39 +296,6 @@ func (it *nodeIterator) Next(descend bool) bool {
 	if it.err != nil {
 		return false
 	}
-
-	// Check if we're still within the subtree boundaries using precomputed paths
-	if it.startPath != nil && len(path) > 0 {
-		if it.prefixMode {
-			// For prefix iteration, use HasPrefix to ensure we stay within the prefix
-			if !bytes.HasPrefix(path, it.startPath) {
-				it.err = errIteratorEnd
-				return false
-			}
-		} else {
-			// For range iteration, ensure we don't return nodes before the lower bound.
-			// Advance the iterator until we reach a node at or after startPath.
-			for bytes.Compare(path, it.startPath) < 0 {
-				// Progress the iterator by pushing the current candidate, then peeking again.
-				it.push(state, parentIndex, path)
-				state, parentIndex, path, err = it.peek(descend)
-				it.err = err
-				if it.err != nil {
-					return false
-				}
-				if len(path) == 0 {
-					break
-				}
-			}
-		}
-	}
-	if it.stopPath != nil && len(path) > 0 {
-		if bytes.Compare(path, it.stopPath) >= 0 {
-			it.err = errIteratorEnd
-			return false
-		}
-	}
-
 	it.push(state, parentIndex, path)
 	return true
 }
@@ -881,7 +837,16 @@ func (it *unionIterator) Error() error {
 	return nil
 }
 
-// NewSubtreeIterator creates an iterator that only traverses nodes within a subtree
+// subTreeIterator wraps nodeIterator to traverse a trie within a predefined
+// start and limit range, with optional prefix mode.
+type subtreeIterator struct {
+	NodeIterator
+
+	stopPath  []byte // Precomputed hex path for stopKey (without terminator), nil means no limit
+	exhausted bool   // Flag whether the iterator has been exhausted
+}
+
+// newSubtreeIterator creates an iterator that only traverses nodes within a subtree
 // defined by the given startKey and stopKey. This supports general range iteration
 // where startKey is inclusive and stopKey is exclusive.
 //
@@ -890,17 +855,27 @@ func (it *unionIterator) Error() error {
 // implemented via hex-nibble path comparisons for efficiency).
 //
 // If startKey is nil, iteration starts from the beginning. If stopKey is nil,
-// iteration continues to the end of the trie. For prefix iteration, use the
-// Trie.NodeIteratorWithPrefix method which handles prefix semantics correctly.
-func NewSubtreeIterator(trie *Trie, startKey, stopKey []byte) NodeIterator {
-	return newSubtreeIterator(trie, startKey, stopKey, false)
-}
-
-// newPrefixIterator creates an iterator that only traverses nodes with the given prefix.
-// This ensures that only keys starting with the prefix are visited.
-func newPrefixIterator(trie *Trie, prefix []byte) NodeIterator {
-	stopKey := nextKey(prefix)
-	return newSubtreeIterator(trie, prefix, stopKey, true)
+// iteration continues to the end of the trie.
+func newSubtreeIterator(trie *Trie, startKey, stopKey []byte) (NodeIterator, error) {
+	it, err := trie.NodeIterator(startKey)
+	if err != nil {
+		return nil, err
+	}
+	if startKey == nil && stopKey == nil {
+		return it, nil
+	}
+	// Precompute nibble paths for efficient comparison
+	var stopPath []byte
+	if stopKey != nil {
+		stopPath = keybytesToHex(stopKey)
+		if hasTerm(stopPath) {
+			stopPath = stopPath[:len(stopPath)-1]
+		}
+	}
+	return &subtreeIterator{
+		NodeIterator: it,
+		stopPath:     stopPath,
+	}, nil
 }
 
 // nextKey returns the next possible key after the given prefix.
@@ -927,51 +902,25 @@ func nextKey(prefix []byte) []byte {
 	return nil
 }
 
-func newSubtreeIterator(trie *Trie, startKey, stopKey []byte, prefixMode bool) NodeIterator {
-	// Precompute nibble paths for efficient comparison
-	var startPath, stopPath []byte
-	if startKey != nil {
-		startPath = keybytesToHex(startKey)
-		if hasTerm(startPath) {
-			startPath = startPath[:len(startPath)-1]
-		}
-	}
-	if stopKey != nil {
-		stopPath = keybytesToHex(stopKey)
-		if hasTerm(stopPath) {
-			stopPath = stopPath[:len(stopPath)-1]
-		}
-	}
-
-	if trie.Hash() == types.EmptyRootHash {
-		return &nodeIterator{
-			trie:       trie,
-			err:        errIteratorEnd,
-			startKey:   startKey,
-			stopKey:    stopKey,
-			startPath:  startPath,
-			stopPath:   stopPath,
-			prefixMode: prefixMode,
-		}
-	}
-	it := &nodeIterator{
-		trie:       trie,
-		startKey:   startKey,
-		stopKey:    stopKey,
-		startPath:  startPath,
-		stopPath:   stopPath,
-		prefixMode: prefixMode,
-	}
-	// Seek to the starting position if startKey is provided
-	if startKey != nil && len(startKey) > 0 {
-		it.err = it.seek(startKey)
-	} else {
-		state, err := it.init()
-		if err != nil {
-			it.err = err
-		} else {
-			it.push(state, nil, nil)
-		}
-	}
-	return it
+// newPrefixIterator creates an iterator that only traverses nodes with the given prefix.
+// This ensures that only keys starting with the prefix are visited.
+func newPrefixIterator(trie *Trie, prefix []byte) (NodeIterator, error) {
+	return newSubtreeIterator(trie, prefix, nextKey(prefix))
+}
+
+// Next moves the iterator to the next node. If the parameter is false, any child
+// nodes will be skipped.
+func (it *subtreeIterator) Next(descend bool) bool {
+	if it.exhausted {
+		return false
+	}
+	if !it.NodeIterator.Next(descend) {
+		it.exhausted = true
+		return false
+	}
+	if it.stopPath != nil && reachedPath(it.NodeIterator.Path(), it.stopPath) {
+		it.exhausted = true
+		return false
+	}
+	return true
 }

trie/iterator_test.go

@@ -19,7 +19,9 @@ package trie
 import (
 	"bytes"
 	"fmt"
+	"maps"
 	"math/rand"
+	"slices"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -625,16 +627,22 @@ func isTrieNode(scheme string, key, val []byte) (bool, []byte, common.Hash) {
 }
 
 func TestSubtreeIterator(t *testing.T) {
-	diskDb := rawdb.NewMemoryDatabase()
-	db := newTestDatabase(diskDb, rawdb.HashScheme)
-	tr := NewEmpty(db)
+	var (
+		db = newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme)
+		tr = NewEmpty(db)
+	)
 	vals := []struct{ k, v string }{
 		{"do", "verb"},
 		{"dog", "puppy"},
 		{"doge", "coin"},
+		{"dog\xff", "value6"},
+		{"dog\xff\xff", "value7"},
 		{"horse", "stallion"},
 		{"house", "building"},
 		{"houses", "multiple"},
+		{"xyz", "value"},
+		{"xyz\xff", "value"},
+		{"xyz\xff\xff", "value"},
 	}
 	all := make(map[string]string)
 	for _, val := range vals {
@@ -644,33 +652,156 @@ func TestSubtreeIterator(t *testing.T) {
 	root, nodes := tr.Commit(false)
 	db.Update(root, types.EmptyRootHash, trienode.NewWithNodeSet(nodes))
 
-	// Test prefix iteration using NewPrefixIterator
-	prefix := []byte("do")
-
-	// We need to re-open the trie from the committed state
+	allNodes := make(map[string][]byte)
 	tr, _ = New(TrieID(root), db)
-	it := newPrefixIterator(tr, prefix)
-
-	found := make(map[string]string)
+	it, err := tr.NodeIterator(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
 	for it.Next(true) {
-		if it.Leaf() {
-			found[string(it.LeafKey())] = string(it.LeafBlob())
+		allNodes[string(it.Path())] = it.NodeBlob()
+	}
+	allKeys := slices.Collect(maps.Keys(all))
+
+	suites := []struct {
+		start    []byte
+		end      []byte
+		expected []string
+	}{
+		// entire key range
+		{
+			start:    nil,
+			end:      nil,
+			expected: allKeys,
+		},
+		{
+			start:    nil,
+			end:      bytes.Repeat([]byte{0xff}, 32),
+			expected: allKeys,
+		},
+		{
+			start:    bytes.Repeat([]byte{0x0}, 32),
+			end:      bytes.Repeat([]byte{0xff}, 32),
+			expected: allKeys,
+		},
+		// key range with start
+		{
+			start:    []byte("do"),
+			end:      nil,
+			expected: allKeys,
+		},
+		{
+			start:    []byte("doe"),
+			end:      nil,
+			expected: allKeys[1:],
+		},
+		{
+			start:    []byte("dog"),
+			end:      nil,
+			expected: allKeys[1:],
+		},
+		{
+			start:    []byte("doge"),
+			end:      nil,
+			expected: allKeys[2:],
+		},
+		{
+			start:    []byte("dog\xff"),
+			end:      nil,
+			expected: allKeys[3:],
+		},
+		{
+			start:    []byte("dog\xff\xff"),
+			end:      nil,
+			expected: allKeys[4:],
+		},
+		{
+			start:    []byte("dog\xff\xff\xff"),
+			end:      nil,
+			expected: allKeys[5:],
+		},
+		// key range with limit
+		{
+			start:    nil,
+			end:      []byte("xyz"),
+			expected: allKeys[:len(allKeys)-3],
+		},
+		{
+			start:    nil,
+			end:      []byte("xyz\xff"),
+			expected: allKeys[:len(allKeys)-2],
+		},
+		{
+			start:    nil,
+			end:      []byte("xyz\xff\xff"),
+			expected: allKeys[:len(allKeys)-1],
+		},
+		{
+			start:    nil,
+			end:      []byte("xyz\xff\xff\xff"),
+			expected: allKeys,
+		},
+	}
+	for _, suite := range suites {
+		// We need to re-open the trie from the committed state
+		tr, _ = New(TrieID(root), db)
+		it, err := newSubtreeIterator(tr, suite.start, suite.end)
+		if err != nil {
+			t.Fatal(err)
 		}
-	}
 
-	// Should find "do", "dog", "doge" but not "horse", "house", "houses"
-	expected := map[string]string{
-		"do":   "verb",
-		"dog":  "puppy",
-		"doge": "coin",
-	}
+		found := make(map[string]string)
+		for it.Next(true) {
+			if it.Leaf() {
+				found[string(it.LeafKey())] = string(it.LeafBlob())
+			}
+		}
+		if len(found) != len(suite.expected) {
+			t.Errorf("wrong number of values: got %d, want %d", len(found), len(suite.expected))
+		}
+		for k, v := range found {
+			if all[k] != v {
+				t.Errorf("wrong value for %s: got %s, want %s", k, found[k], all[k])
+			}
+		}
 
-	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)
+		expectedNodes := make(map[string][]byte)
+		for path, blob := range allNodes {
+			if suite.start != nil {
+				hexStart := keybytesToHex(suite.start)
+				hexStart = hexStart[:len(hexStart)-1]
+				if !reachedPath([]byte(path), hexStart) {
+					continue
+				}
+			}
+			if suite.end != nil {
+				hexEnd := keybytesToHex(suite.end)
+				hexEnd = hexEnd[:len(hexEnd)-1]
+				if reachedPath([]byte(path), hexEnd) {
+					continue
+				}
+			}
+			expectedNodes[path] = bytes.Clone(blob)
+		}
+
+		// Compare the result yield from the subtree iterator
+		var (
+			subCount int
+			subIt, _ = newSubtreeIterator(tr, suite.start, suite.end)
+		)
+		for subIt.Next(true) {
+			blob, ok := expectedNodes[string(subIt.Path())]
+			if !ok {
+				t.Errorf("Unexpected node iterated, path: %v", subIt.Path())
+			}
+			subCount++
+
+			if !bytes.Equal(blob, subIt.NodeBlob()) {
+				t.Errorf("Unexpected node blob, path: %v, want: %v, got: %v", subIt.Path(), blob, subIt.NodeBlob())
+			}
+		}
+		if subCount != len(expectedNodes) {
+			t.Errorf("Unexpected node being iterated, want: %d, got: %d", len(expectedNodes), subCount)
 		}
 	}
 }
@@ -875,7 +1006,7 @@ func TestPrefixIteratorEdgeCases(t *testing.T) {
 			}
 		}
 		// Should find at least the allFF key itself
-		if count < 1 {
+		if count != 2 {
 			t.Errorf("Expected at least 1 result for all-0xff prefix, got %d", count)
 		}
 	})
@@ -919,7 +1050,7 @@ func TestGeneralRangeIteration(t *testing.T) {
 
 	// Test range iteration from "banana" to "fig" (exclusive)
 	t.Run("RangeIteration", func(t *testing.T) {
-		iter := NewSubtreeIterator(trie, []byte("banana"), []byte("fig"))
+		iter, _ := newSubtreeIterator(trie, []byte("banana"), []byte("fig"))
 		found := make(map[string]bool)
 		for iter.Next(true) {
 			if iter.Leaf() {
@@ -937,7 +1068,7 @@ func TestGeneralRangeIteration(t *testing.T) {
 
 	// Test with nil stopKey (iterate to end)
 	t.Run("NilStopKey", func(t *testing.T) {
-		iter := NewSubtreeIterator(trie, []byte("date"), nil)
+		iter, _ := newSubtreeIterator(trie, []byte("date"), nil)
 		found := make(map[string]bool)
 		for iter.Next(true) {
 			if iter.Leaf() {
@@ -955,7 +1086,7 @@ func TestGeneralRangeIteration(t *testing.T) {
 
 	// Test with nil startKey (iterate from beginning)
 	t.Run("NilStartKey", func(t *testing.T) {
-		iter := NewSubtreeIterator(trie, nil, []byte("cherry"))
+		iter, _ := newSubtreeIterator(trie, nil, []byte("cherry"))
 		found := make(map[string]bool)
 		for iter.Next(true) {
 			if iter.Leaf() {

trie/trie.go

@@ -134,20 +134,32 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) {
 	return newNodeIterator(t, start), nil
 }
 
-// NodeIteratorWithPrefix returns an iterator that returns nodes of the trie whose leaf keys
-// start with the given prefix. Iteration includes all keys k where prefix <= k < nextKey(prefix),
-// effectively returning only keys that have the prefix. The iteration stops once it would
-// encounter a key that doesn't start with the prefix.
+// NodeIteratorWithPrefix returns an iterator that returns nodes of the trie
+// whose leaf keys start with the given prefix. Iteration includes all keys
+// where prefix <= k < nextKey(prefix), effectively returning only keys that
+// have the prefix. The iteration stops once it would encounter a key that
+// doesn't start with the prefix.
 //
-// For example, with prefix "dog", the iterator will return "dog", "dogcat", "dogfish" but
-// not "dot" or "fog". An empty prefix iterates the entire trie.
+// For example, with prefix "dog", the iterator will return "dog", "dogcat",
+// "dogfish" but not "dot" or "fog". An empty prefix iterates the entire trie.
 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
 	}
 	// Use the dedicated prefix iterator which handles prefix checking correctly
-	return newPrefixIterator(t, prefix), nil
+	return newPrefixIterator(t, prefix)
+}
+
+// NodeIteratorWithRange returns an iterator over trie nodes whose leaf keys
+// fall within the specified range. It includes all keys where start <= k < end.
+// Iteration stops once a key beyond the end boundary is encountered.
+func (t *Trie) NodeIteratorWithRange(start, end []byte) (NodeIterator, error) {
+	// Short circuit if the trie is already committed and not usable.
+	if t.committed {
+		return nil, ErrCommitted
+	}
+	return newSubtreeIterator(t, start, end)
 }
 
 // MustGet is a wrapper of Get and will omit any encountered error but just