feat: use memstore instead of block wise transient store

app/ibc-hooks/common_test.go

@@ -309,7 +309,6 @@ func _createTestInput(
 		ac,
 		appCodec,
 		runtime.NewKVStoreService(keys[evmtypes.StoreKey]),
-		runtime.NewTransientStoreService(tkeys[evmtypes.TStoreKey]),
 		accountKeeper,
 		bankKeeper,
 		communityPoolKeeper,

app/ibc-hooks/receive_test.go

@@ -18,7 +18,6 @@ import (
 	nfttransfertypes "github.com/initia-labs/initia/x/ibc/nft-transfer/types"
 	evm_hooks "github.com/initia-labs/minievm/app/ibc-hooks"
 	"github.com/initia-labs/minievm/x/evm/contracts/counter"
-	"github.com/initia-labs/minievm/x/evm/types"
 	evmtypes "github.com/initia-labs/minievm/x/evm/types"
 )
 
@@ -111,7 +110,7 @@ func Test_onReceiveIcs20Packet_memo(t *testing.T) {
 	require.Equal(t, uint256.NewInt(1).Bytes32(), [32]byte(queryRes))
 
 	// check allowance
-	erc20Addr, err := types.DenomToContractAddr(ctx, input.EVMKeeper, localDenom)
+	erc20Addr, err := evmtypes.DenomToContractAddr(ctx, input.EVMKeeper, localDenom)
 	require.NoError(t, err)
 	queryInputBz, err = input.EVMKeeper.ERC20Keeper().GetERC20ABI().Pack("allowance", common.BytesToAddress(intermediateSender.Bytes()), contractAddr)
 	require.NoError(t, err)

app/keepers/keepers.go

@@ -553,7 +553,6 @@ func NewAppKeeper(
 		ac,
 		appCodec,
 		runtime.NewKVStoreService(appKeepers.keys[evmtypes.StoreKey]),
-		runtime.NewTransientStoreService(appKeepers.tkeys[evmtypes.TStoreKey]),
 		accountKeeper,
 		bankKeeper,
 		communityPoolKeeper,

app/keepers/keys.go

@@ -57,7 +57,7 @@ func (appKeepers *AppKeepers) GenerateKeys() {
 	)
 
 	// Define transient store keys
-	appKeepers.tkeys = storetypes.NewTransientStoreKeys(evmtypes.TStoreKey, forwardingtypes.TransientStoreKey)
+	appKeepers.tkeys = storetypes.NewTransientStoreKeys(forwardingtypes.TransientStoreKey)
 
 	// MemKeys are for information that is stored only in RAM.
 	appKeepers.memKeys = storetypes.NewMemoryStoreKeys(capabilitytypes.MemStoreKey)

x/bank/keeper/common_test.go

@@ -301,7 +301,6 @@ func _createTestInput(
 		ac,
 		appCodec,
 		runtime.NewKVStoreService(keys[evmtypes.StoreKey]),
-		runtime.NewTransientStoreService(tkeys[evmtypes.TStoreKey]),
 		accountKeeper,
 		bankKeeper,
 		communityPoolKeeper,

x/evm/keeper/common_test.go

@@ -262,7 +262,6 @@ func _createTestInput(
 		ac,
 		appCodec,
 		runtime.NewKVStoreService(keys[evmtypes.StoreKey]),
-		runtime.NewTransientStoreService(tkeys[evmtypes.TStoreKey]),
 		accountKeeper,
 		bankKeeper,
 		communityPoolKeeper,

x/evm/keeper/context.go

@@ -27,11 +27,8 @@
 func (k Keeper) NewStateDB(ctx context.Context, evm *vm.EVM, fee types.Fee) (*evmstate.StateDB, error) {
 	return evmstate.NewStateDB(
 		// delegate gas meter to the EVM
-		sdk.UnwrapSDKContext(ctx).WithGasMeter(storetypes.NewInfiniteGasMeter()), k.Logger(ctx),
-		k.accountKeeper, k.VMStore, k.TransientVMStore, k.TransientCreated,
-		k.TransientSelfDestruct, k.TransientLogs, k.TransientLogSize,
-		k.TransientAccessList, k.TransientRefund, k.execIndex,
-		evm, k.ERC20Keeper().GetERC20ABI(), fee.Contract(),
+		sdk.UnwrapSDKContext(ctx).WithGasMeter(storetypes.NewInfiniteGasMeter()), k.cdc, k.Logger(ctx),
+		k.accountKeeper, k.VMStore, evm, k.ERC20Keeper().GetERC20ABI(), fee.Contract(),
 	)
 }
 
@@ -159,48 +156,54 @@
 	}
 
 	// prepare SDK context for EVM execution
-	ctx, err = prepareSDKContext(sdk.UnwrapSDKContext(ctx))
+	sdkCtx, err := prepareSDKContext(sdk.UnwrapSDKContext(ctx))
 	if err != nil {
 		return ctx, nil, err
 	}
 
-	evm := &vm.EVM{}
-	blockContext, err := k.buildBlockContext(ctx, evm, fee)
+	chainConfig := types.DefaultChainConfig(sdkCtx)
+	vmConfig := vm.Config{Tracer: tracer, ExtraEips: extraEIPs, NumRetainBlockHashes: &params.NumRetainBlockHashes}
+
+	// use dummy block context for chain rules in EVM creation
+	dummyBlockContext, err := k.buildBlockContext(sdkCtx, nil, fee)
 	if err != nil {
 		return ctx, nil, err
 	}
-	txContext, err := k.buildTxContext(ctx, caller, fee)
+	txContext, err := k.buildTxContext(sdkCtx, caller, fee)
 	if err != nil {
 		return ctx, nil, err
 	}
-	stateDB, err := k.NewStateDB(ctx, evm, fee)
+
+	// NOTE: need to check if the EVM is correctly initialized with empty context and stateDB
+	evm := vm.NewEVM(
+		dummyBlockContext,
+		txContext,
+		nil,
+		chainConfig,
+		vmConfig,
+	)
+	// customize EVM contexts and stateDB and precompiles
+	evm.Context, err = k.buildBlockContext(sdkCtx, evm, fee)
 	if err != nil {
 		return ctx, nil, err
 	}
-
-	chainConfig := types.DefaultChainConfig(ctx)
-	rules := chainConfig.Rules(blockContext.BlockNumber, blockContext.Random != nil, blockContext.Time)
-	vmConfig := vm.Config{Tracer: tracer, ExtraEips: extraEIPs, NumRetainBlockHashes: &params.NumRetainBlockHashes}
-	precompiles, err := k.precompiles(rules, stateDB)
+	evm.StateDB, err = k.NewStateDB(sdkCtx, evm, fee)
 	if err != nil {
 		return ctx, nil, err
 	}
-
-	*evm = *vm.NewEVMWithPrecompiles(
-		blockContext,
-		txContext,
-		stateDB,
-		chainConfig,
-		vmConfig,
-		precompiles,
-	)
+	rules := chainConfig.Rules(evm.Context.BlockNumber, evm.Context.Random != nil, evm.Context.Time)
+	precompiles, err := k.precompiles(rules, evm.StateDB.(types.StateDB))
+	if err != nil {
+		return ctx, nil, err
+	}
+	evm.SetPrecompiles(precompiles)
 
 	if tracer != nil {
 		// register vm context to tracer
 		tracer.OnTxStart(evm.GetVMContext(), nil, caller)
 	}
 
-	return ctx, evm, nil
+	return sdkCtx, evm, nil
 }
 
 // prepare SDK context for EVM execution

x/evm/keeper/keeper.go

@@ -48,20 +48,6 @@ type Keeper struct {
 	Params  collections.Item[types.Params]
 	VMStore collections.Map[[]byte, []byte]
 
-	// execIndex is unique index for each execution, which is used
-	// unique key for transient stores.
-	execIndex *atomic.Uint64
-
-	// transient store
-	TSchema               collections.Schema
-	TransientVMStore      collections.Map[collections.Pair[uint64, []byte], []byte]
-	TransientLogs         collections.Map[collections.Pair[uint64, uint64], types.Log]
-	TransientLogSize      collections.Map[uint64, uint64]
-	TransientRefund       collections.Map[uint64, uint64]
-	TransientCreated      collections.KeySet[collections.Pair[uint64, []byte]]
-	TransientSelfDestruct collections.KeySet[collections.Pair[uint64, []byte]]
-	TransientAccessList   collections.KeySet[collections.Pair[uint64, []byte]]
-
 	// erc20 stores of users
 	ERC20FactoryAddr          collections.Item[[]byte]
 	ERC20WrapperAddr          collections.Item[[]byte]
@@ -86,7 +72,6 @@ func NewKeeper(
 	ac address.Codec,
 	cdc codec.Codec,
 	storeService corestoretypes.KVStoreService,
-	transientService corestoretypes.TransientStoreService,
 	accountKeeper types.AccountKeeper,
 	bankKeeper types.BankKeeper,
 	communityPoolKeeper types.CommunityPoolKeeper,
@@ -99,7 +84,6 @@ func NewKeeper(
 	queryCosmosWhitelist types.QueryCosmosWhitelist,
 ) *Keeper {
 	sb := collections.NewSchemaBuilder(storeService)
-	tsb := collections.NewSchemaBuilderFromAccessor(transientService.OpenTransientStore)
 
 	if evmConfig.ContractSimulationGasLimit == 0 {
 		evmConfig.ContractSimulationGasLimit = evmconfig.DefaultContractSimulationGasLimit
@@ -132,16 +116,6 @@ func NewKeeper(
 		Params:  collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
 		VMStore: collections.NewMap(sb, types.VMStorePrefix, "vm_store", collections.BytesKey, collections.BytesValue),
 
-		execIndex: execIndex,
-
-		TransientVMStore:      collections.NewMap(tsb, types.TransientVMStorePrefix, "transient_vm_store", collections.PairKeyCodec(collections.Uint64Key, collections.BytesKey), collections.BytesValue),
-		TransientCreated:      collections.NewKeySet(tsb, types.TransientCreatedPrefix, "transient_created", collections.PairKeyCodec(collections.Uint64Key, collections.BytesKey)),
-		TransientSelfDestruct: collections.NewKeySet(tsb, types.TransientSelfDestructPrefix, "transient_self_destruct", collections.PairKeyCodec(collections.Uint64Key, collections.BytesKey)),
-		TransientLogs:         collections.NewMap(tsb, types.TransientLogsPrefix, "transient_logs", collections.PairKeyCodec(collections.Uint64Key, collections.Uint64Key), codec.CollValue[types.Log](cdc)),
-		TransientLogSize:      collections.NewMap(tsb, types.TransientLogSizePrefix, "transient_log_size", collections.Uint64Key, collections.Uint64Value),
-		TransientAccessList:   collections.NewKeySet(tsb, types.TransientAccessListPrefix, "transient_access_list", collections.PairKeyCodec(collections.Uint64Key, collections.BytesKey)),
-		TransientRefund:       collections.NewMap(tsb, types.TransientRefundPrefix, "transient_refund", collections.Uint64Key, collections.Uint64Value),
-
 		ERC20WrapperAddr:          collections.NewItem(sb, types.ERC20WrapperAddrKey, "erc20_wrapper_addr", collections.BytesValue),
 		ERC20FactoryAddr:          collections.NewItem(sb, types.ERC20FactoryAddrKey, "erc20_factory_addr", collections.BytesValue),
 		ERC20s:                    collections.NewKeySet(sb, types.ERC20sPrefix, "erc20s", collections.BytesKey),
@@ -164,13 +138,8 @@ func NewKeeper(
 	if err != nil {
 		panic(err)
 	}
-	tSchema, err := tsb.Build()
-	if err != nil {
-		panic(err)
-	}
 
 	k.Schema = schema
-	k.TSchema = tSchema
 	k.erc20StoresKeeper = NewERC20StoresKeeper(k)
 	k.erc20Keeper, err = NewERC20Keeper(k)
 	if err != nil {

x/evm/precompiles/cosmos/common_test.go

@@ -24,14 +24,15 @@ import (
 var _ evmtypes.StateDB = &MockStateDB{}
 
 type MockStateDB struct {
-	ctx        sdk.Context
-	initialCtx sdk.Context
+	ctx        state.Context
+	initialCtx state.Context
 
 	// Snapshot stack
 	snaps []*state.Snapshot
 }
 
-func NewMockStateDB(ctx sdk.Context) *MockStateDB {
+func NewMockStateDB(sdkCtx sdk.Context) *MockStateDB {
+	ctx := state.NewContext(sdkCtx)
 	return &MockStateDB{
 		ctx:        ctx,
 		initialCtx: ctx,
@@ -73,15 +74,15 @@ func (m *MockStateDB) RevertToSnapshot(i int) {
 // ContextOfSnapshot implements types.StateDB.
 func (m *MockStateDB) ContextOfSnapshot(i int) sdk.Context {
 	if i == -1 {
-		return m.initialCtx
+		return m.initialCtx.Context
 	}
 
-	return m.snaps[i].Context()
+	return m.snaps[i].Context().Context
 }
 
 // Context implements types.StateDB.
 func (m *MockStateDB) Context() sdk.Context {
-	return m.ctx
+	return m.ctx.Context
 }
 
 //////////////////////// MOCKED METHODS ////////////////////////

x/evm/precompiles/erc20_registry/common_test.go

@@ -18,14 +18,15 @@ import (
 var _ evmtypes.StateDB = &MockStateDB{}
 
 type MockStateDB struct {
-	ctx        sdk.Context
-	initialCtx sdk.Context
+	ctx        state.Context
+	initialCtx state.Context
 
 	// Snapshot stack
 	snaps []*state.Snapshot
 }
 
-func NewMockStateDB(ctx sdk.Context) *MockStateDB {
+func NewMockStateDB(sdkCtx sdk.Context) *MockStateDB {
+	ctx := state.NewContext(sdkCtx)
 	return &MockStateDB{
 		ctx:        ctx,
 		initialCtx: ctx,
@@ -67,15 +68,15 @@ func (m *MockStateDB) RevertToSnapshot(i int) {
 // ContextOfSnapshot implements types.StateDB.
 func (m *MockStateDB) ContextOfSnapshot(i int) sdk.Context {
 	if i == -1 {
-		return m.initialCtx
+		return m.initialCtx.Context
 	}
 
-	return m.snaps[i].Context()
+	return m.snaps[i].Context().Context
 }
 
 // Context implements types.StateDB.
 func (m *MockStateDB) Context() sdk.Context {
-	return m.ctx
+	return m.ctx.Context
 }
 
 //////////////////////// MOCKED METHODS ////////////////////////

x/evm/state/common_test.go

@@ -261,7 +261,6 @@ func _createTestInput(
 		ac,
 		appCodec,
 		runtime.NewKVStoreService(keys[evmtypes.StoreKey]),
-		runtime.NewTransientStoreService(tkeys[evmtypes.TStoreKey]),
 		accountKeeper,
 		bankKeeper,
 		communityPoolKeeper,

x/evm/state/context.go

@@ -0,0 +1,45 @@
+package state
+
+import (
+	storetypes "cosmossdk.io/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+type Context struct {
+	sdk.Context
+
+	memStore    storetypes.MultiStore
+	memStoreKey storetypes.StoreKey
+}
+
+func NewContext(ctx sdk.Context) Context {
+	memStore, memStoreKey := newMemStore()
+	return Context{
+		Context:     ctx,
+		memStore:    memStore,
+		memStoreKey: memStoreKey,
+	}
+}
+
+func (c Context) WithSDKContext(sdkCtx sdk.Context) Context {
+	c.Context = sdkCtx
+	return c
+}
+
+func (c Context) WithMemStore(memStore storetypes.MultiStore) Context {
+	c.memStore = memStore
+	return c
+}
+
+func (c Context) CacheContext() (cc Context, writeCache func()) {
+	cacheCtx, commit := c.Context.CacheContext()
+	cacheMemStore := c.memStore.CacheMultiStore()
+
+	cc = c.WithSDKContext(cacheCtx).WithMemStore(cacheMemStore)
+	writeCache = func() {
+		commit()
+		cacheMemStore.Write()
+	}
+
+	return cc, writeCache
+}

x/evm/state/keys.go

@@ -37,3 +37,14 @@ func uint64ToBytes(v uint64) []byte {
 func bytesToUint64(bz []byte) uint64 {
 	return binary.BigEndian.Uint64(bz)
 }
+
+// Keys for transient store
+var (
+	memStoreVMStorePrefix      = []byte{0x01} // prefix for transient vm store
+	memStoreCreatedPrefix      = []byte{0x02} // prefix for transient created accounts
+	memStoreSelfDestructPrefix = []byte{0x03} // prefix for transient self destruct accounts
+	memStoreLogsPrefix         = []byte{0x04} // prefix for transient logs
+	memStoreLogSizePrefix      = []byte{0x05} // prefix for transient log size
+	memStoreAccessListPrefix   = []byte{0x06} // prefix for transient access list
+	memStoreRefundPrefix       = []byte{0x07} // prefix for transient refund
+)