Problem: type wrapping needs more allocations

Solution: use sdk types directly, avoid some wrappers in integration

Author: yihuang

memdb.go

@@ -2,7 +2,10 @@ package block_stm
 
 import (
 	"bytes"
+	"io"
 
+	"cosmossdk.io/store/cachekv"
+	"cosmossdk.io/store/tracekv"
 	storetypes "cosmossdk.io/store/types"
 	"github.com/tidwall/btree"
 )
@@ -22,7 +25,7 @@ type MemDB struct {
 	btree.BTreeG[memdbItem]
 }
 
-var _ KVStore = (*MemDB)(nil)
+var _ storetypes.KVStore = (*MemDB)(nil)
 
 func NewMemDB() *MemDB {
 	return &MemDB{*btree.NewBTreeG[memdbItem](KeyItemLess)}
@@ -114,14 +117,28 @@ func (db *MemDB) Equal(other *MemDB) bool {
 	}
 }
 
+func (db *MemDB) GetStoreType() storetypes.StoreType {
+	return storetypes.StoreTypeIAVL
+}
+
+// CacheWrap implements types.KVStore.
+func (db *MemDB) CacheWrap() storetypes.CacheWrap {
+	return cachekv.NewStore(storetypes.KVStore(db))
+}
+
+// CacheWrapWithTrace implements types.KVStore.
+func (db *MemDB) CacheWrapWithTrace(w io.Writer, tc storetypes.TraceContext) storetypes.CacheWrap {
+	return cachekv.NewStore(tracekv.NewStore(db, w, tc))
+}
+
 type MemDBIterator struct {
 	BTreeIteratorG[memdbItem]
 }
 
 var _ storetypes.Iterator = (*MemDBIterator)(nil)
 
 func NewMemDBIterator(start, end Key, iter btree.IterG[memdbItem], ascending bool) *MemDBIterator {
-	return &MemDBIterator{*NewBTreeIteratorG[memdbItem](
+	return &MemDBIterator{*NewBTreeIteratorG(
 		memdbItem{key: start},
 		memdbItem{key: end},
 		iter,
@@ -151,6 +168,6 @@ func (mmdb *MultiMemDB) GetDB(store storetypes.StoreKey) *MemDB {
 	return mmdb.dbs[store]
 }
 
-func (mmdb *MultiMemDB) GetKVStore(store storetypes.StoreKey) KVStore {
+func (mmdb *MultiMemDB) GetKVStore(store storetypes.StoreKey) storetypes.KVStore {
 	return mmdb.GetDB(store)
 }

mock_block.go

@@ -116,7 +116,7 @@ func genRandomSignature() func() {
 	}
 }
 
-func increaseNonce(i int, sender string, store KVStore) error {
+func increaseNonce(i int, sender string, store storetypes.KVStore) error {
 	nonceKey := []byte("nonce" + sender)
 	var nonce uint64
 	v := store.Get(nonceKey)
@@ -136,7 +136,7 @@ func increaseNonce(i int, sender string, store KVStore) error {
 	return nil
 }
 
-func bankTransfer(i int, sender, receiver string, amount uint64, store KVStore) error {
+func bankTransfer(i int, sender, receiver string, amount uint64, store storetypes.KVStore) error {
 	senderKey := []byte("balance" + sender)
 	receiverKey := []byte("balance" + receiver)
 

multimvview.go

@@ -25,7 +25,7 @@ func NewMultiMVMemoryView(
 	}
 }
 
-func (mv *MultiMVMemoryView) GetKVStore(name storetypes.StoreKey) KVStore {
+func (mv *MultiMVMemoryView) GetKVStore(name storetypes.StoreKey) storetypes.KVStore {
 	return mv.views[name]
 }
 

mvdata.go

@@ -22,7 +22,7 @@ func (d *MVData) getTree(key Key) *BTree[secondaryDataItem] {
 func (d *MVData) getTreeOrDefault(key Key) *BTree[secondaryDataItem] {
 	return d.GetOrDefault(dataItem{Key: key}, func(item *dataItem) {
 		if item.Tree == nil {
-			item.Tree = NewBTree[secondaryDataItem](secondaryDataItemLess)
+			item.Tree = NewBTree(secondaryDataItemLess)
 		}
 	}).Tree
 }

mvmemory.go

@@ -148,7 +148,7 @@ func (mv *MVMemory) WriteSnapshot(storage MultiStore) {
 	}
 }
 
-func WriteSnapshot(storage KVStore, snapshot []KVPair) {
+func WriteSnapshot(storage storetypes.KVStore, snapshot []KVPair) {
 	for _, pair := range snapshot {
 		if pair.Value == nil {
 			storage.Delete(pair.Key)

mvview.go

@@ -1,10 +1,16 @@
 package block_stm
 
-import storetypes "cosmossdk.io/store/types"
+import (
+	"io"
+
+	"cosmossdk.io/store/cachekv"
+	"cosmossdk.io/store/tracekv"
+	storetypes "cosmossdk.io/store/types"
+)
 
 // MVMemoryView wraps `MVMemory` for execution of a single transaction.
 type MVMemoryView struct {
-	storage   KVStore
+	storage   storetypes.KVStore
 	mvMemory  *MVMemory
 	scheduler *Scheduler
 	store     int
@@ -14,9 +20,9 @@ type MVMemoryView struct {
 	writeSet WriteSet
 }
 
-var _ KVStore = (*MVMemoryView)(nil)
+var _ storetypes.KVStore = (*MVMemoryView)(nil)
 
-func NewMVMemoryView(store int, storage KVStore, mvMemory *MVMemory, schedule *Scheduler, txn TxnIndex) *MVMemoryView {
+func NewMVMemoryView(store int, storage storetypes.KVStore, mvMemory *MVMemory, schedule *Scheduler, txn TxnIndex) *MVMemoryView {
 	return &MVMemoryView{
 		store:     store,
 		storage:   storage,
@@ -126,3 +132,18 @@ func (s *MVMemoryView) iterator(opts IteratorOptions) storetypes.Iterator {
 		onClose,
 	)
 }
+
+// CacheWrap implements types.KVStore.
+func (s *MVMemoryView) CacheWrap() storetypes.CacheWrap {
+	return cachekv.NewStore(storetypes.KVStore(s))
+}
+
+// CacheWrapWithTrace implements types.KVStore.
+func (s *MVMemoryView) CacheWrapWithTrace(w io.Writer, tc storetypes.TraceContext) storetypes.CacheWrap {
+	return cachekv.NewStore(tracekv.NewStore(s, w, tc))
+}
+
+// GetStoreType implements types.KVStore.
+func (s *MVMemoryView) GetStoreType() storetypes.StoreType {
+	return storetypes.StoreTypeIAVL
+}

types.go

@@ -77,12 +77,6 @@ func KeyItemLess[T KeyItem](a, b T) bool {
 // TxExecutor executes transactions on top of a multi-version memory view.
 type TxExecutor func(TxnIndex, MultiStore)
 
-type KVStore interface {
-	storetypes.BasicKVStore
-	Iterator(start, end []byte) storetypes.Iterator
-	ReverseIterator(start, end []byte) storetypes.Iterator
-}
-
 type MultiStore interface {
-	GetKVStore(storetypes.StoreKey) KVStore
+	GetKVStore(storetypes.StoreKey) storetypes.KVStore
 }