forks/go-ethereum
1
0
Fork 1
mirror of https://github.com/ethereum/go-ethereum.git synced 2026-07-23 13:16:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

trie: rewrite the subtreeIterator

Browse source
This commit is contained in:
Gary Rong 2025-09-15 10:19:40 +08:00
parent 2db2a18e83
commit 45ee2fbe19
3 changed files with 231 additions and 139 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

143
trie/iterator.go
View file

@ -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)
// 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
}
// 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)
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]
}
// 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))
}
if trie.Hash() == types.EmptyRootHash {
return &nodeIterator{
trie: trie,
err: errIteratorEnd,
startKey: startKey,
stopKey: stopKey,
startPath: startPath,
stopPath: stopPath,
prefixMode: prefixMode,
// 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
}
it := &nodeIterator{
trie: trie,
startKey: startKey,
stopKey: stopKey,
startPath: startPath,
stopPath: stopPath,
prefixMode: prefixMode,
if it.stopPath != nil && reachedPath(it.NodeIterator.Path(), it.stopPath) {
it.exhausted = true
return false
}
// 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
return true
}

173
trie/iterator_test.go
View file

@ -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,12 +652,103 @@ 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")
allNodes := make(map[string][]byte)
tr, _ = New(TrieID(root), db)
it, err := tr.NodeIterator(nil)
if err != nil {
t.Fatal(err)
}
for it.Next(true) {
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 := newPrefixIterator(tr, prefix)
it, err := newSubtreeIterator(tr, suite.start, suite.end)
if err != nil {
t.Fatal(err)
}
found := make(map[string]string)
for it.Next(true) {
@ -657,20 +756,52 @@ func TestSubtreeIterator(t *testing.T) {
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(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))
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
}
for k, v := range expected {
if found[k] != v {
t.Errorf("wrong value for %s: got %s, want %s", k, found[k], v)
}
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() {

26
trie/trie.go
View file

@ -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