Merkle db iterator (ava-labs#1533)

Signed-off-by: David Boehm <91908103+dboehm-avalabs@users.noreply.github.com>
Co-authored-by: Qian (Tony) Zhang <tonyqzhang@yahoo.com>
Co-authored-by: Dan Laine <daniel.laine@avalabs.org>
Co-authored-by: Darioush Jalali <darioush.jalali@avalabs.org>

Author: 4 people

x/merkledb/db.go

@@ -1162,7 +1162,7 @@ func (db *merkleDB) getHistoricalViewForRange(
 	return newTrieViewWithChanges(db, db, changeHistory, len(changeHistory.nodes))
 }
 
-// Returns all of the keys in range [start, end] that aren't in [keySet].
+// Returns all keys in range [start, end] that aren't in [keySet].
 // If [start] is nil, then the range has no lower bound.
 // If [end] is nil, then the range has no upper bound.
 func (db *merkleDB) getKeysNotInSet(start, end []byte, keySet set.Set[string]) ([][]byte, error) {
@@ -1242,55 +1242,6 @@ func (db *merkleDB) getNode(key path) (*node, error) {
 	return node, err
 }
 
-// If [lock], grabs [db.lock]'s read lock.
-// Otherwise assumes [db.lock] is already read locked.
-func (db *merkleDB) getKeyValues(
-	start []byte,
-	end []byte,
-	maxLength int,
-	keysToIgnore set.Set[string],
-	lock bool,
-) ([]KeyValue, error) {
-	if lock {
-		db.lock.RLock()
-		defer db.lock.RUnlock()
-	}
-
-	if db.closed {
-		return nil, database.ErrClosed
-	}
-
-	if maxLength <= 0 {
-		return nil, fmt.Errorf("%w but was %d", ErrInvalidMaxLength, maxLength)
-	}
-
-	it := db.NewIteratorWithStart(start)
-	defer it.Release()
-
-	remainingLength := maxLength
-	result := make([]KeyValue, 0, maxLength)
-	// Keep adding key/value pairs until one of the following:
-	// * We hit a key that is lexicographically larger than the end key.
-	// * [maxLength] elements are in [result].
-	// * There are no more values to add.
-	for remainingLength > 0 && it.Next() {
-		key := it.Key()
-		if len(end) != 0 && bytes.Compare(it.Key(), end) > 0 {
-			break
-		}
-		if keysToIgnore.Contains(string(key)) {
-			continue
-		}
-		result = append(result, KeyValue{
-			Key:   key,
-			Value: it.Value(),
-		})
-		remainingLength--
-	}
-
-	return result, it.Error()
-}
-
 // Returns a new view atop [db] with the changes in [ops] applied to it.
 // Assumes [db.commitLock] is read locked.
 func (db *merkleDB) prepareBatchView(ops []database.BatchOp) (*trieView, error) {

x/merkledb/mock_db.go

@@ -7,8 +7,8 @@ import (
 	"context"
 	"errors"
 
+	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/ids"
-	"github.com/ava-labs/avalanchego/utils/set"
 )
 
 var errNoNewRoot = errors.New("there was no updated root in change list")
@@ -46,13 +46,7 @@ type ReadOnlyTrie interface {
 	// GetRangeProof generates a proof of up to maxLength smallest key/values with keys between start and end
 	GetRangeProof(ctx context.Context, start, end []byte, maxLength int) (*RangeProof, error)
 
-	getKeyValues(
-		start []byte,
-		end []byte,
-		maxLength int,
-		keysToIgnore set.Set[string],
-		lock bool,
-	) ([]KeyValue, error)
+	database.Iteratee
 }
 
 type Trie interface {

x/merkledb/trie.go

@@ -20,10 +20,12 @@ import (
 
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/ids"
-	"github.com/ava-labs/avalanchego/utils/set"
 )
 
-const defaultPreallocationSize = 100
+const (
+	initKeyValuesSize        = 256
+	defaultPreallocationSize = 100
+)
 
 var (
 	_ TrieView = (*trieView)(nil)
@@ -418,25 +420,20 @@ func (t *trieView) GetRangeProof(
 		t.lock.RLock()
 	}
 
-	var (
-		result RangeProof
-		err    error
-	)
+	var result RangeProof
 
-	result.KeyValues, err = t.getKeyValues(
-		start,
-		end,
-		maxLength,
-		set.Set[string]{},
-		false, /*lock*/
-	)
-	if err != nil {
-		return nil, err
+	result.KeyValues = make([]KeyValue, 0, initKeyValuesSize)
+	it := t.NewIteratorWithStart(start)
+	for it.Next() && len(result.KeyValues) < maxLength && (len(end) == 0 || bytes.Compare(it.Key(), end) <= 0) {
+		// clone the value to prevent editing of the values stored within the trie
+		result.KeyValues = append(result.KeyValues, KeyValue{
+			Key:   it.Key(),
+			Value: slices.Clone(it.Value()),
+		})
 	}
-
-	// copy values, so edits won't affect the underlying arrays
-	for i, kv := range result.KeyValues {
-		result.KeyValues[i] = KeyValue{Key: kv.Key, Value: slices.Clone(kv.Value)}
+	it.Release()
+	if err := it.Error(); err != nil {
+		return nil, err
 	}
 
 	// This proof may not contain all key-value pairs in [start, end] due to size limitations.
@@ -719,140 +716,6 @@ func (t *trieView) getMerkleRoot(ctx context.Context) (ids.ID, error) {
 	return t.root.id, nil
 }
 
-// Returns up to [maxLength] key/values from keys in closed range [start, end].
-// Acts similarly to the merge step of a merge sort to combine state from the view
-// with state from the parent trie.
-// If [lock], grabs [t.lock]'s read lock.
-// Otherwise assumes [t.lock]'s read lock is held.
-func (t *trieView) getKeyValues(
-	start []byte,
-	end []byte,
-	maxLength int,
-	keysToIgnore set.Set[string],
-	lock bool,
-) ([]KeyValue, error) {
-	if lock {
-		t.lock.RLock()
-		defer t.lock.RUnlock()
-	}
-
-	if maxLength <= 0 {
-		return nil, fmt.Errorf("%w but was %d", ErrInvalidMaxLength, maxLength)
-	}
-
-	if t.isInvalid() {
-		return nil, ErrInvalid
-	}
-
-	// collect all values that have changed or been deleted
-	changes := make([]KeyValue, 0, len(t.changes.values))
-	for key, change := range t.changes.values {
-		if change.after.IsNothing() {
-			// This was deleted
-			keysToIgnore.Add(string(key.Serialize().Value))
-		} else {
-			changes = append(changes, KeyValue{
-				Key:   key.Serialize().Value,
-				Value: change.after.value,
-			})
-		}
-	}
-	// sort [changes] so they can be merged with the parent trie's state
-	slices.SortFunc(changes, func(a, b KeyValue) bool {
-		return bytes.Compare(a.Key, b.Key) == -1
-	})
-
-	baseKeyValues, err := t.getParentTrie().getKeyValues(
-		start,
-		end,
-		maxLength,
-		keysToIgnore,
-		true, /*lock*/
-	)
-	if err != nil {
-		return nil, err
-	}
-
-	var (
-		// True if there are more key/value pairs from [baseKeyValues] to add to result
-		baseKeyValuesFinished = false
-		// True if there are more key/value pairs from [changes] to add to result
-		changesFinished = false
-		// The index of the next key/value pair to add from [baseKeyValues].
-		baseKeyValuesIndex = 0
-		// The index of the next key/value pair to add from [changes].
-		changesIndex    = 0
-		remainingLength = maxLength
-		hasUpperBound   = len(end) > 0
-		result          = make([]KeyValue, 0, len(baseKeyValues))
-	)
-
-	// keep adding key/value pairs until one of the following:
-	// * a key that is lexicographically larger than the end key is hit
-	// * the maxLength is hit
-	// * no more values are available to add
-	for remainingLength > 0 {
-		// the baseKeyValues iterator is finished when we have run out of keys or hit a key greater than the end key
-		baseKeyValuesFinished = baseKeyValuesFinished ||
-			(baseKeyValuesIndex >= len(baseKeyValues) || (hasUpperBound && bytes.Compare(baseKeyValues[baseKeyValuesIndex].Key, end) == 1))
-
-		// the changes iterator is finished when we have run out of keys or hit a key greater than the end key
-		changesFinished = changesFinished ||
-			(changesIndex >= len(changes) || (hasUpperBound && bytes.Compare(changes[changesIndex].Key, end) == 1))
-
-		// if both the base state and changes are finished, return the result of the merge
-		if baseKeyValuesFinished && changesFinished {
-			return result, nil
-		}
-
-		// one or both iterators still have values, so one will be added to the result
-		remainingLength--
-
-		// both still have key/values available, so add the smallest key
-		if !changesFinished && !baseKeyValuesFinished {
-			currentChangeState := changes[changesIndex]
-			currentKeyValues := baseKeyValues[baseKeyValuesIndex]
-
-			switch bytes.Compare(currentChangeState.Key, currentKeyValues.Key) {
-			case -1:
-				result = append(result, currentChangeState)
-				changesIndex++
-			case 0:
-				// the keys are the same, so override the base value with the changed value
-				result = append(result, currentChangeState)
-				changesIndex++
-				baseKeyValuesIndex++
-			case 1:
-				result = append(result, currentKeyValues)
-				baseKeyValuesIndex++
-			}
-			continue
-		}
-
-		// the base state is not finished, but the changes is finished.
-		// add the next base state value.
-		if !baseKeyValuesFinished {
-			currentBaseState := baseKeyValues[baseKeyValuesIndex]
-			result = append(result, currentBaseState)
-			baseKeyValuesIndex++
-			continue
-		}
-
-		// the base state is finished, but the changes is not finished.
-		// add the next changes value.
-		currentChangeState := changes[changesIndex]
-		result = append(result, currentChangeState)
-		changesIndex++
-	}
-
-	// ensure no ancestor changes occurred during execution
-	if t.isInvalid() {
-		return nil, ErrInvalid
-	}
-
-	return result, nil
-}
-
 func (t *trieView) GetValues(_ context.Context, keys [][]byte) ([][]byte, []error) {
 	t.lock.RLock()
 	defer t.lock.RUnlock()