db/state: adopt improvements from #19742 in temporal_mem_batch and domain_stream

Author: anacrolix

db/state/domain_stream.go

@@ -359,7 +359,7 @@ func (hi *DomainLatestIterFile) Next() ([]byte, []byte, error) {
 // debugIteratePrefix iterates over key-value pairs of the storage domain that start with given prefix
 //
 // k and v lifetime is bounded by the lifetime of the iterator
-func (dt *DomainRoTx) debugIteratePrefixLatest(prefix []byte, ramIter *btree2.MapIterator[string, []dataWithTxNum], it func(k []byte, v []byte, step kv.Step) (cont bool, err error), roTx kv.Tx) error {
+func (dt *DomainRoTx) debugIteratePrefixLatest(prefix []byte, storage *btree2.Map[string, []dataWithTxNum], it func(k []byte, v []byte, step kv.Step) (cont bool, err error), roTx kv.Tx) error {
 	// Implementation:
 	//     File endTxNum  = last txNum of file step
 	//     DB endTxNum    = first txNum of step in db
@@ -375,15 +375,16 @@ func (dt *DomainRoTx) debugIteratePrefixLatest(prefix []byte, ramIter *btree2.Ma
 	var k, v []byte
 	var err error
 
-	if ramIter != nil {
-		ramIter.SeekGE(string(prefix))
-		if ramIter.Valid() {
-			k := common.ToBytesZeroCopy(ramIter.Cur())
+	if storage != nil {
+		ramIt := storage.Iterator()
+		ramIt.SeekGE(string(prefix))
+		if ramIt.Valid() {
+			k := common.ToBytesZeroCopy(ramIt.Cur())
 
-			v = ramIter.Value()[len(ramIter.Value())-1].data
+			v = ramIt.Value()[len(ramIt.Value())-1].data
 
 			if len(k) > 0 && bytes.HasPrefix(k, prefix) {
-				heap.Push(cpPtr, &CursorItem{t: RAM_CURSOR, key: common.Copy(k), val: common.Copy(v), step: 0, iter: ramIter, endTxNum: math.MaxUint64, reverse: true})
+				heap.Push(cpPtr, &CursorItem{t: RAM_CURSOR, key: common.Copy(k), val: common.Copy(v), step: 0, iter: &ramIt, endTxNum: math.MaxUint64, reverse: true})
 			}
 		}
 	}

db/state/temporal_mem_batch.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"maps"
 	"sort"
+	"strings"
 	"sync"
 
 	btree2 "github.com/anacrolix/btree"
@@ -60,7 +61,7 @@ type TemporalMemBatch struct {
 
 	latestStateLock sync.RWMutex
 	domains         [kv.DomainLen]map[string][]dataWithTxNum
-	storage         *btree2.Map[string, []dataWithTxNum] // TODO: replace hardcoded domain name to per-config configuration of available Guarantees/AccessMethods (range vs get)
+	storage         btree2.Map[string, []dataWithTxNum] // TODO: replace hardcoded domain name to per-config configuration of available Guarantees/AccessMethods (range vs get)
 
 	domainWriters   [kv.DomainLen]*DomainBufferedWriter
 	iiWriters       []*InvertedIndexBufferedWriter
@@ -79,22 +80,9 @@ type TemporalMemBatch struct {
 	metrics *changeset.DomainMetrics
 }
 
-func newStringDataMap() *btree2.Map[string, []dataWithTxNum] {
-	m := btree2.MakeMap[string, []dataWithTxNum](func(a, b string) int {
-		if a == b {
-			return 0
-		}
-		if a < b {
-			return -1
-		}
-		return 1
-	})
-	return &m
-}
-
 func NewTemporalMemBatch(tx kv.TemporalTx, ioMetrics any) *TemporalMemBatch {
 	sd := &TemporalMemBatch{
-		storage:           newStringDataMap(),
+		storage:           btree2.MakeMap[string, []dataWithTxNum](strings.Compare),
 		metrics:           ioMetrics.(*changeset.DomainMetrics),
 		inMemHistoryReads: true,
 	}
@@ -321,13 +309,12 @@ func (sd *TemporalMemBatch) ClearRam() {
 func (sd *TemporalMemBatch) IteratePrefix(domain kv.Domain, prefix []byte, roTx kv.Tx, it func(k []byte, v []byte, step kv.Step) (cont bool, err error)) error {
 	sd.latestStateLock.RLock()
 	defer sd.latestStateLock.RUnlock()
-	var ramIter *btree2.MapIterator[string, []dataWithTxNum]
+	var storageMap *btree2.Map[string, []dataWithTxNum]
 	if domain == kv.StorageDomain {
-		iter := sd.storage.Iterator()
-		ramIter = &iter
+		storageMap = &sd.storage
 	}
 
-	return AggTx(roTx).d[domain].debugIteratePrefixLatest(prefix, ramIter, it, roTx)
+	return AggTx(roTx).d[domain].debugIteratePrefixLatest(prefix, storageMap, it, roTx)
 }
 
 func (sd *TemporalMemBatch) HasPrefix(domain kv.Domain, prefix []byte, roTx kv.Tx) ([]byte, []byte, bool, error) {