trie: add sub-trie iterator with prefix and range iteration support

trie/iterator.go

@@ -150,11 +150,11 @@ type nodeIterator struct {
 	// 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)
 }
@@ -309,30 +309,30 @@ func (it *nodeIterator) Next(descend bool) bool {
 	}
 
 	// Check if we're still within the subtree boundaries using precomputed paths
-    if it.startPath != nil && len(path) > 0 {
+	if it.startPath != nil && len(path) > 0 {
-        if it.prefixMode {
+		if it.prefixMode {
-            // For prefix iteration, use HasPrefix to ensure we stay within the prefix
+			// For prefix iteration, use HasPrefix to ensure we stay within the prefix
-            if !bytes.HasPrefix(path, it.startPath) {
+			if !bytes.HasPrefix(path, it.startPath) {
-                it.err = errIteratorEnd
+				it.err = errIteratorEnd
-                return false
+				return false
-            }
+			}
-        } else {
+		} else {
-            // For range iteration, ensure we don't return nodes before the lower bound.
+			// 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.
+			// Advance the iterator until we reach a node at or after startPath.
-            for bytes.Compare(path, it.startPath) < 0 {
+			for bytes.Compare(path, it.startPath) < 0 {
-                // Progress the iterator by pushing the current candidate, then peeking again.
+				// Progress the iterator by pushing the current candidate, then peeking again.
-                it.push(state, parentIndex, path)
+				it.push(state, parentIndex, path)
-                state, parentIndex, path, err = it.peek(descend)
+				state, parentIndex, path, err = it.peek(descend)
-                it.err = err
+				it.err = err
-                if it.err != nil {
+				if it.err != nil {
-                    return false
+					return false
-                }
+				}
-                if len(path) == 0 {
+				if len(path) == 0 {
-                    break
+					break
-                }
+				}
-            }
+			}
-        }
+		}
-    }
+	}
 	if it.stopPath != nil && len(path) > 0 {
 		if bytes.Compare(path, it.stopPath) >= 0 {
 			it.err = errIteratorEnd
@@ -883,13 +883,13 @@ func (it *unionIterator) Error() error {
 
 // 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. 
+// where startKey is inclusive and stopKey is exclusive.
 //
 // The iterator will only visit nodes whose keys k satisfy: startKey <= k < stopKey,
 // where comparisons are performed in lexicographic order of byte keys (internally
 // implemented via hex-nibble path comparisons for efficiency).
 //
-// If startKey is nil, iteration starts from the beginning. If stopKey is nil, 
+// 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 {
@@ -912,7 +912,7 @@ func nextKey(prefix []byte) []byte {
 	// Make a copy to avoid modifying the original
 	next := make([]byte, len(prefix))
 	copy(next, prefix)
-	
+
 	// Increment the last byte that isn't 0xff
 	for i := len(next) - 1; i >= 0; i-- {
 		if next[i] < 0xff {
@@ -942,7 +942,7 @@ func newSubtreeIterator(trie *Trie, startKey, stopKey []byte, prefixMode bool) N
 			stopPath = stopPath[:len(stopPath)-1]
 		}
 	}
-	
+
 	if trie.Hash() == types.EmptyRootHash {
 		return &nodeIterator{
 			trie:       trie,

trie/iterator_test.go

@@ -786,13 +786,13 @@ func TestPrefixIteratorEdgeCases(t *testing.T) {
 	// Create a trie with test data
 	trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
 	testData := map[string]string{
-		"abc":      "value1",
+		"abc":         "value1",
-		"abcd":     "value2",
+		"abcd":        "value2",
-		"abce":     "value3",
+		"abce":        "value3",
-		"abd":      "value4",
+		"abd":         "value4",
-		"dog":      "value5",
+		"dog":         "value5",
-		"dog\xff":  "value6",  // Test with 0xff byte
+		"dog\xff":     "value6", // Test with 0xff byte
-		"dog\xff\xff": "value7",  // Multiple 0xff bytes
+		"dog\xff\xff": "value7", // Multiple 0xff bytes
 	}
 	for key, value := range testData {
 		trie.Update([]byte(key), []byte(value))
@@ -863,7 +863,7 @@ func TestPrefixIteratorEdgeCases(t *testing.T) {
 		allFF := []byte{0xff, 0xff}
 		trie.Update(allFF, []byte("all_ff_value"))
 		trie.Update(append(allFF, 0x00), []byte("all_ff_plus"))
-		
+
 		iter, err := trie.NodeIteratorWithPrefix(allFF)
 		if err != nil {
 			t.Fatalf("Failed to create iterator: %v", err)
@@ -905,13 +905,13 @@ func TestGeneralRangeIteration(t *testing.T) {
 	// Create a trie with test data
 	trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
 	testData := map[string]string{
-		"apple":  "fruit1",
+		"apple":   "fruit1",
 		"apricot": "fruit2",
-		"banana": "fruit3",
+		"banana":  "fruit3",
-		"cherry": "fruit4",
+		"cherry":  "fruit4",
-		"date":   "fruit5",
+		"date":    "fruit5",
-		"fig":    "fruit6",
+		"fig":     "fruit6",
-		"grape":  "fruit7",
+		"grape":   "fruit7",
 	}
 	for key, value := range testData {
 		trie.Update([]byte(key), []byte(value))
@@ -977,12 +977,12 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
 	// Create a trie with nested structure
 	trie := NewEmpty(newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.HashScheme))
 	testData := map[string]string{
-		"a":       "value_a",
+		"a":     "value_a",
-		"a/b":     "value_ab",
+		"a/b":   "value_ab",
-		"a/b/c":   "value_abc",
+		"a/b/c": "value_abc",
-		"a/b/d":   "value_abd",
+		"a/b/d": "value_abd",
-		"a/e":     "value_ae",
+		"a/e":   "value_ae",
-		"b":       "value_b",
+		"b":     "value_b",
 	}
 	for key, value := range testData {
 		trie.Update([]byte(key), []byte(value))
@@ -994,17 +994,17 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
 		if err != nil {
 			t.Fatalf("Failed to create iterator: %v", err)
 		}
-		
+
 		// Count nodes at each level
 		nodesVisited := 0
 		leafsFound := make(map[string]bool)
-		
+
 		// First call with descend=true to enter the "a" subtree
 		if !iter.Next(true) {
 			t.Fatal("Expected to find at least one node")
 		}
 		nodesVisited++
-		
+
 		// Continue iteration, sometimes with descend=false
 		descendPattern := []bool{false, true, false, true, true}
 		for i := 0; iter.Next(descendPattern[i%len(descendPattern)]); i++ {
@@ -1013,14 +1013,15 @@ func TestPrefixIteratorWithDescend(t *testing.T) {
 				leafsFound[string(iter.LeafKey())] = true
 			}
 		}
-		
+
 		// We should still respect the prefix boundary even when skipping
+		prefix := []byte("a")
 		for key := range leafsFound {
-			if !bytes.HasPrefix([]byte(key), []byte("a")) {
+			if !bytes.HasPrefix([]byte(key), prefix) {
 				t.Errorf("Found key outside prefix when using descend=false: %s", key)
 			}
 		}
-		
+
 		// Should not have found "b" even if we skip some subtrees
 		if leafsFound["b"] {
 			t.Error("Iterator leaked outside prefix boundary with descend=false")

trie/trie.go

@@ -136,10 +136,10 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) {
 
 // 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 
+// 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 
+// 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.