feat: use memstore instead of block wise transient store (#117)

* use memstore instead of block wise transient store

* add testing

* remove unnecessary sdkCtx

* rollback log

* fix to set key value with account address

* fix comments

* fix error handling

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/contracts/counter/Counter.sol

@@ -14,15 +14,17 @@ contract Counter is IIBCAsyncCallback {
 
     constructor() payable {}
 
-    function increase_for_fuzz(uint64 num) external {
-        for (uint64 i = 0; i < num; i++) {
-            increase();
+    function increase_for_fuzz(uint64 num) public {
+        if (num == 0) {
+            return;
         }
+
+        increase();
+        increase_for_fuzz(num - 1);
     }
 
     function increase() public payable {
         count++;
-
         emit increased(count - 1, count);
     }
 

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 @@ import (
 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(),
 	)
 }
 
@@ -164,36 +161,42 @@ func (k Keeper) CreateEVM(ctx context.Context, caller common.Address, tracer *tr
 		return ctx, nil, err
 	}
 
-	evm := &vm.EVM{}
-	blockContext, err := k.buildBlockContext(ctx, evm, fee)
+	chainConfig := types.DefaultChainConfig(ctx)
+	vmConfig := vm.Config{Tracer: tracer, ExtraEips: extraEIPs, NumRetainBlockHashes: &params.NumRetainBlockHashes}
+
+	// use dummy block context for chain rules in EVM creation
+	dummyBlockContext, err := k.buildBlockContext(ctx, nil, fee)
 	if err != nil {
 		return ctx, nil, err
 	}
 	txContext, err := k.buildTxContext(ctx, 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(ctx, 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(ctx, 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

x/evm/keeper/fuzz_test.go

@@ -1,17 +1,14 @@
 package keeper_test
 
 import (
-	"bytes"
 	"sync"
-	"sync/atomic"
 	"testing"
 
 	storetypes "cosmossdk.io/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
-	coretypes "github.com/ethereum/go-ethereum/core/types"
 
 	"github.com/holiman/uint256"
 	"github.com/stretchr/testify/require"
@@ -48,9 +45,6 @@ func Fuzz_Concurrent_Counter(f *testing.F) {
 		inputBz, err := parsed.Pack("increase_for_fuzz", uint64(numCount))
 		require.NoError(t, err)
 
-		atomicBloomBytes := atomic.Pointer[[]byte]{}
-		atomicBloomBytes.Store(nil)
-
 		var wg sync.WaitGroup
 		cacheCtxes := make([]sdk.Context, numThread)
 		for i := uint8(0); i < numThread; i++ {
@@ -65,21 +59,33 @@ func Fuzz_Concurrent_Counter(f *testing.F) {
 				res, logs, err := input.EVMKeeper.EVMCall(ctx, caller, contractAddr, inputBz, nil, nil)
 				require.NoError(t, err)
 				require.Empty(t, res)
-				bloomBytes := coretypes.LogsBloom(logs.ToEthLogs())
-				prev := atomicBloomBytes.Swap(&bloomBytes)
-				require.True(t, prev == nil || bytes.Equal(*prev, bloomBytes))
+				assertLogs(t, contractAddr, logs)
 			}(cacheCtx)
 		}
 		wg.Wait()
 
 		for i := uint8(0); i < numThread; i++ {
 			count := getCount(t, cacheCtxes[i], input, contractAddr)
 			require.Equal(t, uint256.NewInt(uint64(numCount)), count)
-			require.NotEmpty(t, atomicBloomBytes.Load())
 		}
 	})
 }
 
+func assertLogs(t *testing.T, contractAddr common.Address, logs []types.Log) {
+	for i := range logs {
+		require.Equal(t, contractAddr.Hex(), logs[i].Address)
+		dataBz, err := hexutil.Decode(logs[i].Data)
+		require.NoError(t, err)
+		require.Len(t, dataBz, 64)
+		before := new(uint256.Int).SetBytes(dataBz[:32])
+		after := new(uint256.Int).SetBytes(dataBz[32:])
+
+		require.NoError(t, err)
+		require.Equal(t, uint256.NewInt(uint64(i)), before)
+		require.Equal(t, uint256.NewInt(uint64(i+1)), after)
+	}
+}
+
 func getCount(t *testing.T, ctx sdk.Context, input TestKeepers, contractAddr common.Address) *uint256.Int {
 	parsed, err := counter.CounterMetaData.GetAbi()
 	require.NoError(t, err)

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,