Problem: CacheWrapWithTrace is not supported

runtime/store.go

@@ -2,6 +2,7 @@ package runtime
 
 import (
 	"context"
+	"io"
 
 	"cosmossdk.io/core/store"
 	storetypes "cosmossdk.io/store/types"
@@ -109,6 +110,10 @@ func (kvStoreAdapter) CacheWrap() storetypes.CacheWrap {
 	panic("unimplemented")
 }
 
+func (kvStoreAdapter) CacheWrapWithTrace(w io.Writer, tc storetypes.TraceContext) storetypes.CacheWrap {
+	panic("unimplemented")
+}
+
 func (kvStoreAdapter) GetStoreType() storetypes.StoreType {
 	panic("unimplemented")
 }

server/mock/store.go

@@ -30,6 +30,10 @@ func (ms multiStore) CacheWrap() storetypes.CacheWrap {
 	panic("not implemented")
 }
 
+func (ms multiStore) CacheWrapWithTrace(w io.Writer, tc storetypes.TraceContext) storetypes.CacheWrap {
+	panic("unimplemented")
+}
+
 func (ms multiStore) TracingEnabled() bool {
 	panic("not implemented")
 }
@@ -178,6 +182,10 @@ func (kv kvStore) CacheWrap() storetypes.CacheWrap {
 	panic("not implemented")
 }
 
+func (kv kvStore) CacheWrapWithTrace(w io.Writer, tc storetypes.TraceContext) storetypes.CacheWrap {
+	panic("unimplemented")
+}
+
 func (kv kvStore) GetStoreType() storetypes.StoreType {
 	panic("not implemented")
 }

store/CHANGELOG.md

@@ -35,6 +35,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#243](https://github.com/crypto-org-chain/cosmos-sdk/pull/243) Support `RunAtomic` API to use new CoW cache store.
 * [#244](https://github.com/crypto-org-chain/cosmos-sdk/pull/244) Add `Discard` method to CacheWrap to discard the write buffer.
 * [#258](https://github.com/crypto-org-chain/cosmos-sdk/pull/258) Add `NewFromParent` API to cachemulti store to create a new store from block-stm multiversion data structure.
+* [#1043](https://github.com/crypto-org-chain/cosmos-sdk/pull/1043) Add back CacheWrapWithTrace api.
 
 ## [Unreleased]
 

store/cachekv/store.go

@@ -1,8 +1,11 @@
 package cachekv
 
 import (
+	"io"
+
 	"cosmossdk.io/store/cachekv/internal"
 	"cosmossdk.io/store/internal/btree"
+	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
 
@@ -128,6 +131,14 @@ func (store *GStore[V]) CacheWrap() types.CacheWrap {
 	return NewGStore(store, store.isZero, store.valueLen)
 }
 
+// CacheWrapWithTrace implements the CacheWrapper interface.
+func (store *GStore[V]) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	if store, ok := any(store).(*GStore[[]byte]); ok {
+		return NewStore(tracekv.NewStore(store, w, tc))
+	}
+	return store.CacheWrap()
+}
+
 //----------------------------------------
 // Iteration
 

store/cachemulti/store.go

@@ -145,6 +145,11 @@ func (cms Store) CacheWrap() types.CacheWrap {
 	return cms.CacheMultiStore().(types.CacheWrap)
 }
 
+// CacheWrapWithTrace implements the CacheWrapper interface.
+func (cms Store) CacheWrapWithTrace(_ io.Writer, _ types.TraceContext) types.CacheWrap {
+	return cms.CacheWrap()
+}
+
 // Implements MultiStore.
 func (cms Store) CacheMultiStore() types.CacheMultiStore {
 	return NewFromParent(cms.getCacheWrapper, cms.traceWriter, cms.traceContext)

store/dbadapter/store.go

@@ -1,9 +1,12 @@
 package dbadapter
 
 import (
+	"io"
+
 	dbm "github.com/cosmos/cosmos-db"
 
 	"cosmossdk.io/store/cachekv"
+	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
 
@@ -78,5 +81,10 @@ func (dsa Store) CacheWrap() types.CacheWrap {
 	return cachekv.NewStore(dsa)
 }
 
+// CacheWrapWithTrace implements KVStore.
+func (dsa Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	return cachekv.NewStore(tracekv.NewStore(dsa, w, tc))
+}
+
 // dbm.DB implements KVStore so we can CacheKVStore it.
 var _ types.KVStore = Store{}

store/dbadapter/store_test.go

@@ -80,4 +80,7 @@ func TestCacheWraps(t *testing.T) {
 
 	cacheWrapper := store.CacheWrap()
 	require.IsType(t, &cachekv.Store{}, cacheWrapper)
+
+	cacheWrappedWithTrace := store.CacheWrapWithTrace(nil, nil)
+	require.IsType(t, &cachekv.Store{}, cacheWrappedWithTrace)
 }

store/gaskv/store.go

@@ -1,6 +1,10 @@
 package gaskv
 
-import "cosmossdk.io/store/types"
+import (
+	"io"
+
+	"cosmossdk.io/store/types"
+)
 
 // ObjectValueLength is the emulated number of bytes for storing transient objects in gas accounting.
 const ObjectValueLength = 16
@@ -115,6 +119,11 @@ func (gs *GStore[V]) CacheWrap() types.CacheWrap {
 	panic("cannot CacheWrap a GasKVStore")
 }
 
+// CacheWrapWithTrace implements the KVStore interface.
+func (gs *GStore[V]) CacheWrapWithTrace(_ io.Writer, _ types.TraceContext) types.CacheWrap {
+	panic("cannot CacheWrapWithTrace a GasKVStore")
+}
+
 func (gs *GStore[V]) iterator(start, end []byte, ascending bool) types.GIterator[V] {
 	var parent types.GIterator[V]
 	if ascending {

store/gaskv/store_test.go

@@ -24,6 +24,7 @@ func TestGasKVStoreBasic(t *testing.T) {
 
 	require.Equal(t, types.StoreTypeDB, st.GetStoreType())
 	require.Panics(t, func() { st.CacheWrap() })
+	require.Panics(t, func() { st.CacheWrapWithTrace(nil, nil) })
 
 	require.Panics(t, func() { st.Set(nil, []byte("value")) }, "setting a nil key should panic")
 	require.Panics(t, func() { st.Set([]byte(""), []byte("value")) }, "setting an empty key should panic")

store/iavl/store.go

@@ -3,6 +3,7 @@ package iavl
 import (
 	"errors"
 	"fmt"
+	"io"
 
 	cmtprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 	dbm "github.com/cosmos/cosmos-db"
@@ -15,6 +16,7 @@ import (
 	"cosmossdk.io/store/internal/kv"
 	"cosmossdk.io/store/metrics"
 	pruningtypes "cosmossdk.io/store/pruning/types"
+	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 	"cosmossdk.io/store/wrapper"
 )
@@ -191,6 +193,11 @@ func (st *Store) CacheWrap() types.CacheWrap {
 	return cachekv.NewStore(st)
 }
 
+// CacheWrapWithTrace implements the Store interface.
+func (st *Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	return cachekv.NewStore(tracekv.NewStore(st, w, tc))
+}
+
 // Implements types.KVStore.
 func (st *Store) Set(key, value []byte) {
 	types.AssertValidKey(key)

store/iavl/store_test.go

@@ -663,6 +663,9 @@ func TestCacheWraps(t *testing.T) {
 
 	cacheWrapper := store.CacheWrap()
 	require.IsType(t, &cachekv.Store{}, cacheWrapper)
+
+	cacheWrappedWithTrace := store.CacheWrapWithTrace(nil, nil)
+	require.IsType(t, &cachekv.Store{}, cacheWrappedWithTrace)
 }
 
 func TestChangeSets(t *testing.T) {

store/internal/btreeadaptor.go

@@ -1,6 +1,8 @@
 package internal
 
 import (
+	"io"
+
 	"cosmossdk.io/store/cachekv"
 	"cosmossdk.io/store/internal/btree"
 	"cosmossdk.io/store/types"
@@ -59,3 +61,8 @@ func (ts *BTreeStore[V]) GetStoreType() types.StoreType {
 func (ts *BTreeStore[V]) CacheWrap() types.CacheWrap {
 	return cachekv.NewGStore(ts, ts.isZero, ts.valueLen)
 }
+
+// CacheWrapWithTrace branches the underlying store.
+func (ts *BTreeStore[V]) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	return cachekv.NewGStore(ts, ts.isZero, ts.valueLen)
+}

store/listenkv/store.go

@@ -1,6 +1,8 @@
 package listenkv
 
 import (
+	"io"
+
 	"cosmossdk.io/store/cachekv"
 	"cosmossdk.io/store/types"
 )
@@ -133,3 +135,9 @@ func (s *Store) GetStoreType() types.StoreType {
 func (s *Store) CacheWrap() types.CacheWrap {
 	return cachekv.NewStore(s)
 }
+
+// CacheWrapWithTrace implements the KVStore interface. It panics as a
+// Store cannot be cache wrapped.
+func (s *Store) CacheWrapWithTrace(_ io.Writer, _ types.TraceContext) types.CacheWrap {
+	panic("cannot CacheWrapWithTrace a ListenKVStore")
+}

store/listenkv/store_test.go

@@ -274,3 +274,8 @@ func TestListenKVStoreCacheWrap(t *testing.T) {
 	store := newEmptyListenKVStore(nil)
 	store.CacheWrap()
 }
+
+func TestListenKVStoreCacheWrapWithTrace(t *testing.T) {
+	store := newEmptyListenKVStore(nil)
+	require.Panics(t, func() { store.CacheWrapWithTrace(nil, nil) })
+}

store/mem/mem_test.go

@@ -30,6 +30,9 @@ func TestStore(t *testing.T) {
 
 	cacheWrapper := db.CacheWrap()
 	require.IsType(t, &cachekv.Store{}, cacheWrapper)
+
+	cacheWrappedWithTrace := db.CacheWrapWithTrace(nil, nil)
+	require.IsType(t, &cachekv.Store{}, cacheWrappedWithTrace)
 }
 
 func TestCommit(t *testing.T) {

store/mem/store.go

@@ -1,11 +1,14 @@
 package mem
 
 import (
+	"io"
+
 	dbm "github.com/cosmos/cosmos-db"
 
 	"cosmossdk.io/store/cachekv"
 	"cosmossdk.io/store/dbadapter"
 	pruningtypes "cosmossdk.io/store/pruning/types"
+	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
 
@@ -38,6 +41,11 @@ func (s Store) CacheWrap() types.CacheWrap {
 	return cachekv.NewStore(s)
 }
 
+// CacheWrapWithTrace implements KVStore.
+func (s Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	return cachekv.NewStore(tracekv.NewStore(s, w, tc))
+}
+
 // Commit performs a no-op as entries are persistent between commitments.
 func (s *Store) Commit() (id types.CommitID) { return }
 

store/prefix/store.go

@@ -3,8 +3,10 @@ package prefix
 import (
 	"bytes"
 	"errors"
+	"io"
 
 	"cosmossdk.io/store/cachekv"
+	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
 
@@ -83,6 +85,14 @@ func (s GStore[V]) CacheWrap() types.CacheWrap {
 	return cachekv.NewGStore(s, s.isZero, s.valueLen)
 }
 
+// CacheWrapWithTrace implements the KVStore interface.
+func (s GStore[V]) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	if store, ok := any(s).(*GStore[[]byte]); ok {
+		return cachekv.NewGStore(tracekv.NewStore(store, w, tc), store.isZero, store.valueLen)
+	}
+	return s.CacheWrap()
+}
+
 // Implements KVStore
 func (s GStore[V]) Get(key []byte) V {
 	res := s.parent.Get(s.key(key))

store/prefix/store_test.go

@@ -445,4 +445,7 @@ func TestCacheWraps(t *testing.T) {
 
 	cacheWrapper := store.CacheWrap()
 	require.IsType(t, &cachekv.Store{}, cacheWrapper)
+
+	cacheWrappedWithTrace := store.CacheWrapWithTrace(nil, nil)
+	require.IsType(t, &cachekv.Store{}, cacheWrappedWithTrace)
 }

store/rootmulti/store.go

@@ -572,6 +572,11 @@ func (rs *Store) CacheWrap() types.CacheWrap {
 	return rs.CacheMultiStore().(types.CacheWrap)
 }
 
+// CacheWrapWithTrace implements the CacheWrapper interface.
+func (rs *Store) CacheWrapWithTrace(_ io.Writer, _ types.TraceContext) types.CacheWrap {
+	return rs.CacheWrap()
+}
+
 // CacheMultiStore creates ephemeral branch of the multi-store and returns a CacheMultiStore.
 // It implements the MultiStore interface.
 func (rs *Store) CacheMultiStore() types.CacheMultiStore {

store/rootmulti/store_test.go

@@ -17,7 +17,6 @@ import (
 	sdkmaps "cosmossdk.io/store/internal/maps"
 	"cosmossdk.io/store/metrics"
 	pruningtypes "cosmossdk.io/store/pruning/types"
-	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
 
@@ -731,6 +730,9 @@ func TestCacheWraps(t *testing.T) {
 
 	cacheWrapper := multi.CacheWrap()
 	require.IsType(t, cachemulti.Store{}, cacheWrapper)
+
+	cacheWrappedWithTrace := multi.CacheWrapWithTrace(nil, nil)
+	require.IsType(t, cachemulti.Store{}, cacheWrappedWithTrace)
 }
 
 func TestTraceConcurrency(t *testing.T) {
@@ -748,7 +750,7 @@ func TestTraceConcurrency(t *testing.T) {
 
 	cms := multi.CacheMultiStore()
 	store1 := cms.GetKVStore(key)
-	cw := tracekv.NewStore(store1.CacheWrap().(types.KVStore), b, tc)
+	cw := store1.CacheWrapWithTrace(b, tc)
 	_ = cw
 	require.NotNil(t, store1)
 

store/tracekv/store.go

@@ -6,7 +6,6 @@ import (
 	"io"
 
 	"cosmossdk.io/errors"
-	"cosmossdk.io/store/cachekv"
 	"cosmossdk.io/store/types"
 )
 
@@ -165,7 +164,13 @@ func (tkv *Store) GetStoreType() types.StoreType {
 // CacheWrap implements the KVStore interface. It panics because a Store
 // cannot be branched.
 func (tkv *Store) CacheWrap() types.CacheWrap {
-	return cachekv.NewStore(tkv)
+	return tkv.parent.CacheWrap()
+}
+
+// CacheWrapWithTrace implements the KVStore interface. It panics as a
+// Store cannot be branched.
+func (tkv *Store) CacheWrapWithTrace(_ io.Writer, _ types.TraceContext) types.CacheWrap {
+	panic("cannot CacheWrapWithTrace a TraceKVStore")
 }
 
 // writeOperation writes a KVStore operation to the underlying io.Writer as

store/tracekv/store_test.go

@@ -285,3 +285,8 @@ func TestTraceKVStoreCacheWrap(t *testing.T) {
 	store := newEmptyTraceKVStore(nil)
 	store.CacheWrap()
 }
+
+func TestTraceKVStoreCacheWrapWithTrace(t *testing.T) {
+	store := newEmptyTraceKVStore(nil)
+	require.Panics(t, func() { store.CacheWrapWithTrace(nil, nil) })
+}

store/types/store.go

@@ -353,11 +353,17 @@ type CacheWrap interface {
 
 	// Discard the write set
 	Discard()
+
+	// CacheWrapWithTrace branches a store with tracing enabled.
+	CacheWrapWithTrace(w io.Writer, tc TraceContext) CacheWrap
 }
 
 type CacheWrapper interface {
 	// CacheWrap branches a store.
 	CacheWrap() CacheWrap
+
+	// CacheWrapWithTrace branches a store with tracing enabled.
+	CacheWrapWithTrace(w io.Writer, tc TraceContext) CacheWrap
 }
 
 func (cid CommitID) IsZero() bool {