fix: handle transient store keys in snapshot multi store

evmd/app.go

@@ -245,6 +245,7 @@ func NewExampleApp(
 		evmtypes.StoreKey, feemarkettypes.StoreKey, erc20types.StoreKey,
 	)
 	oKeys := storetypes.NewObjectStoreKeys(banktypes.ObjectStoreKey, evmtypes.ObjectKey)
+	tKeys := storetypes.NewTransientStoreKeys(evmtypes.TransientKey)
 
 	var nonTransientKeys []storetypes.StoreKey
 	for _, k := range keys {
@@ -471,7 +472,7 @@ func NewExampleApp(
 	// NOTE: it's required to set up the EVM keeper before the ERC-20 keeper, because it is used in its instantiation.
 	app.EVMKeeper = evmkeeper.NewKeeper(
 		// TODO: check why this is not adjusted to use the runtime module methods like SDK native keepers
-		appCodec, keys[evmtypes.StoreKey], oKeys[evmtypes.ObjectKey], nonTransientKeys,
+		appCodec, keys[evmtypes.StoreKey], oKeys[evmtypes.ObjectKey], nonTransientKeys, tKeys,
 		authtypes.NewModuleAddress(govtypes.ModuleName),
 		app.AccountKeeper,
 		app.BankKeeper,

mempool/interface.go

@@ -40,6 +40,7 @@ type VMKeeperI interface {
 	SetCode(ctx sdk.Context, codeHash []byte, code []byte)
 	DeleteAccount(ctx sdk.Context, addr common.Address) error
 	KVStoreKeys() map[string]storetypes.StoreKey
+	TransientStoreKeys() map[string]*storetypes.TransientStoreKey
 	SetEvmMempool(evmMempool NotifiedMempool)
 }
 

x/erc20/types/interfaces.go

@@ -59,6 +59,7 @@ type EVMKeeper interface {
 	SetState(ctx sdk.Context, addr common.Address, key common.Hash, value []byte)
 	DeleteCode(ctx sdk.Context, codeHash []byte)
 	KVStoreKeys() map[string]storetypes.StoreKey
+	TransientStoreKeys() map[string]*storetypes.TransientStoreKey
 }
 
 type Erc20Keeper interface {

x/ibc/callbacks/types/expected_keepers.go

@@ -40,6 +40,7 @@ type EVMKeeper interface {
 	SetCode(ctx sdk.Context, codeHash []byte, code []byte)
 	DeleteAccount(ctx sdk.Context, addr common.Address) error
 	KVStoreKeys() map[string]storetypes.StoreKey
+	TransientStoreKeys() map[string]*storetypes.TransientStoreKey
 }
 
 type ERC20Keeper interface {

x/vm/keeper/keeper.go

@@ -51,6 +51,9 @@ type Keeper struct {
 	// Store Keys for modules wired to app
 	storeKeys map[string]storetypes.StoreKey
 
+	// Transient Store Keys for modules wired to app
+	transientStoreKeys map[string]*storetypes.TransientStoreKey
+
 	// the address capable of executing a MsgUpdateParams message. Typically, this should be the x/gov module account.
 	authority sdk.AccAddress
 
@@ -102,6 +105,7 @@ func NewKeeper(
 	cdc codec.BinaryCodec,
 	storeKey, objectKey storetypes.StoreKey,
 	keys []storetypes.StoreKey,
+	tkeys map[string]*storetypes.TransientStoreKey,
 	authority sdk.AccAddress,
 	ak types.AccountKeeper,
 	bankKeeper types.BankKeeper,
@@ -137,18 +141,19 @@ func NewKeeper(
 
 	// NOTE: we pass in the parameter space to the CommitStateDB in order to use custom denominations for the EVM operations
 	return &Keeper{
-		cdc:              cdc,
-		authority:        authority,
-		accountKeeper:    ak,
-		bankWrapper:      bankWrapper,
-		stakingKeeper:    sk,
-		feeMarketWrapper: feeMarketWrapper,
-		storeKey:         storeKey,
-		objectKey:        objectKey,
-		tracer:           tracer,
-		consensusKeeper:  consensusKeeper,
-		erc20Keeper:      erc20Keeper,
-		storeKeys:        storeKeys,
+		cdc:                cdc,
+		authority:          authority,
+		accountKeeper:      ak,
+		bankWrapper:        bankWrapper,
+		stakingKeeper:      sk,
+		feeMarketWrapper:   feeMarketWrapper,
+		storeKey:           storeKey,
+		objectKey:          objectKey,
+		tracer:             tracer,
+		consensusKeeper:    consensusKeeper,
+		erc20Keeper:        erc20Keeper,
+		storeKeys:          storeKeys,
+		transientStoreKeys: tkeys,
 	}
 }
 
@@ -440,6 +445,11 @@ func (k Keeper) KVStoreKeys() map[string]storetypes.StoreKey {
 	return k.storeKeys
 }
 
+// TransientStoreKeys returns transient store keys injected to keeper
+func (k Keeper) TransientStoreKeys() map[string]*storetypes.TransientStoreKey {
+	return k.transientStoreKeys
+}
+
 // SetMempool sets the mempool that is notified of new blocks via the
 // EndBlocker.
 func (k *Keeper) SetEvmMempool(evmMempool evmmempool.NotifiedMempool) {

x/vm/keeper/keeper_test.go

@@ -43,8 +43,8 @@ func (suite *KeeperTestSuite) SetupTest() {
 	key := storetypes.NewKVStoreKey(vmtypes.StoreKey)
 	oKey := storetypes.NewObjectStoreKey(vmtypes.ObjectKey)
 	allKeys := []storetypes.StoreKey{key, oKey}
-	testCtx := testutil.DefaultContextWithObjectStore(suite.T(), key,
-		storetypes.NewTransientStoreKey("store_test"), oKey)
+	tkeys := storetypes.NewTransientStoreKeys(vmtypes.TransientKey)
+	testCtx := testutil.DefaultContextWithObjectStore(suite.T(), key, tkeys[vmtypes.TransientKey], oKey)
 	ctx := testCtx.Ctx.WithBlockHeader(cmtproto.Header{Time: cmttime.Now()})
 	encCfg := moduletestutil.MakeTestEncodingConfig()
 
@@ -65,6 +65,7 @@ func (suite *KeeperTestSuite) SetupTest() {
 		key,
 		oKey,
 		allKeys,
+		tkeys,
 		authority,
 		suite.accKeeper,
 		suite.bankKeeper,

x/vm/statedb/interfaces.go

@@ -40,4 +40,5 @@ type Keeper interface {
 	// Getter for injected Store keys
 	// It is used for StateDB.snapshotter creation
 	KVStoreKeys() map[string]storetypes.StoreKey
+	TransientStoreKeys() map[string]*storetypes.TransientStoreKey
 }

x/vm/statedb/mock_test.go

@@ -123,3 +123,7 @@ func (k MockKeeper) Clone() *MockKeeper {
 func (k MockKeeper) KVStoreKeys() map[string]storetypes.StoreKey {
 	return make(map[string]storetypes.StoreKey)
 }
+
+func (k MockKeeper) TransientStoreKeys() map[string]*storetypes.TransientStoreKey {
+	return nil
+}

x/vm/statedb/statedb.go

@@ -3,6 +3,7 @@ package statedb
 import (
 	"errors"
 	"fmt"
+	"maps"
 	"slices"
 	"sort"
 
@@ -203,10 +204,15 @@ func (s *StateDB) cache() error {
 
 	// Get KVStores for modules wired to app
 	cms := s.cacheCtx.MultiStore()
-	storeKeys := s.keeper.KVStoreKeys()
+	mergedStoreKeys := maps.Clone(s.keeper.KVStoreKeys())
+
+	// Merge transient store keys into storeKeys map
+	for name, tkey := range s.keeper.TransientStoreKeys() {
+		mergedStoreKeys[name] = tkey
+	}
 
 	// Create and set snapshot store to stateDB
-	snapshotStore := snapshotmulti.NewStore(cms, storeKeys)
+	snapshotStore := snapshotmulti.NewStore(cms, mergedStoreKeys)
 	s.snapshotter = snapshotStore
 	s.cacheCtx = s.cacheCtx.WithMultiStore(snapshotStore)
 	s.writeCache = func() {

x/vm/store/snapshotmulti/store.go

@@ -29,10 +29,11 @@ func NewStore(cms storetypes.MultiStore, keys map[string]storetypes.StoreKey) *S
 	}
 
 	for _, key := range s.storeKeys {
-		if _, ok := key.(*storetypes.KVStoreKey); ok {
+		switch key.(type) {
+		case *storetypes.KVStoreKey, *storetypes.TransientStoreKey:
 			store := cms.GetKVStore(key)
 			s.stores[key] = snapshotkv.NewStore(store.(storetypes.CacheWrap))
-		} else {
+		default:
 			store := cms.GetObjKVStore(key)
 			s.stores[key] = snapshotkv.NewStore(store.(storetypes.CacheWrap))
 		}

x/vm/types/key.go

@@ -19,6 +19,10 @@ const (
 	// during the Commit phase.
 	ObjectKey = "object:" + ModuleName
 
+	// TransientKey is the key to access the EVM transient store, that is reset
+	// on every block.
+	TransientKey = "transient:" + ModuleName
+
 	// RouterKey uses module name for routing
 	RouterKey = ModuleName
 )

x/vm/types/mocks/EVMKeeper.go

@@ -135,6 +135,10 @@ func (k EVMKeeper) KVStoreKeys() map[string]storetypes.StoreKey {
 	return result
 }
 
+func (k EVMKeeper) TransientStoreKeys() map[string]*storetypes.TransientStoreKey {
+	return nil
+}
+
 func (k EVMKeeper) GetCodeHash(_ sdk.Context, _ common.Address) common.Hash {
 	return common.Hash{}
 }