feat(x/vm): apply stack-based StateDB snapshot mechamism for precompile call

contracts/solidity/precompiles/testutil/contracts/StakingReverter.sol

@@ -31,6 +31,20 @@ contract StakingReverter {
         }
     }
 
+    function callPrecompileBeforeAndAfterRevert(uint numTimes, string calldata validatorAddress) external {
+        STAKING_CONTRACT.delegate(address(this), validatorAddress, 10);
+
+        for (uint i = 0; i < numTimes; i++) {
+            try
+            StakingReverter(address(this)).performDelegation(
+                validatorAddress
+            )
+            {} catch {}
+        }
+
+        STAKING_CONTRACT.delegate(address(this), validatorAddress, 10);
+    }
+
     function performDelegation(string calldata validatorAddress) external {
         STAKING_CONTRACT.delegate(address(this), validatorAddress, 10);
         revert();

evmd/app.go

@@ -30,6 +30,7 @@ import (
 	feemarketkeeper "github.com/cosmos/evm/x/feemarket/keeper"
 	feemarkettypes "github.com/cosmos/evm/x/feemarket/types"
 	ibccallbackskeeper "github.com/cosmos/evm/x/ibc/callbacks/keeper"
+
 	// NOTE: override ICS20 keeper to support IBC transfers of ERC20 tokens
 	"github.com/cosmos/evm/x/ibc/transfer"
 	transferkeeper "github.com/cosmos/evm/x/ibc/transfer/keeper"
@@ -483,7 +484,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], tkeys[evmtypes.TransientKey],
+		appCodec, keys[evmtypes.StoreKey], tkeys[evmtypes.TransientKey], keys,
 		authtypes.NewModuleAddress(govtypes.ModuleName),
 		app.AccountKeeper,
 		app.PreciseBankKeeper,

evmd/go.mod

@@ -13,6 +13,7 @@ require (
 	cosmossdk.io/tools/confix v0.1.2
 	cosmossdk.io/x/evidence v0.2.0
 	cosmossdk.io/x/feegrant v0.2.0
+	cosmossdk.io/x/tx v0.14.0
 	cosmossdk.io/x/upgrade v0.2.0
 	github.com/cometbft/cometbft v0.38.17
 	github.com/cosmos/cosmos-db v1.1.3
@@ -43,7 +44,6 @@ require (
 	cosmossdk.io/collections v1.2.1 // indirect
 	cosmossdk.io/depinject v1.2.1 // indirect
 	cosmossdk.io/schema v1.1.0 // indirect
-	cosmossdk.io/x/tx v0.14.0 // indirect
 	filippo.io/edwards25519 v1.1.0 // indirect
 	github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect
 	github.com/99designs/keyring v1.2.2 // indirect

precompiles/testutil/contracts/StakingReverter.sol

@@ -31,6 +31,24 @@ contract StakingReverter {
         }
     }
 
+    /// @dev callPrecompileBeforeAndAfterRevert tests whether precompile calls that occur 
+    /// before and after an intentionally ignored revert correctly modify the state.
+    /// This method assumes that the StakingReverter.sol contract holds a native balance. 
+    /// Therefore, in order to call this method, the contract must be funded with a balance in advance.
+    function callPrecompileBeforeAndAfterRevert(uint numTimes, string calldata validatorAddress) external {
+        STAKING_CONTRACT.delegate(address(this), validatorAddress, 10);
+
+        for (uint i = 0; i < numTimes; i++) {
+            try
+            StakingReverter(address(this)).performDelegation(
+                validatorAddress
+            )
+            {} catch {}
+        }
+
+        STAKING_CONTRACT.delegate(address(this), validatorAddress, 10);
+    }
+
     function performDelegation(string calldata validatorAddress) external {
         STAKING_CONTRACT.delegate(address(this), validatorAddress, 10);
         revert();

tests/integration/precompiles/distribution/test_distribution.go

@@ -497,13 +497,18 @@ func (s *PrecompileTestSuite) TestCMS() {
 			if tc.expPass {
 				s.Require().NoError(err, "expected no error when running the precompile")
 				s.Require().NotNil(resp.Ret, "expected returned bytes not to be nil")
-				testutil.ValidateWrites(s.T(), cms, 2)
+				// NOTES: After stack-based snapshot mechanism is added for precopile call,
+				// CacheMultiStore.Write() is always called once when tx succeeds.
+				// It is because CacheMultiStore() is not called when creating snapshot for MultiStore,
+				// Count of Write() is not accumulated.
+				testutil.ValidateWrites(s.T(), cms, 1)
 			} else {
 				s.Require().Error(err, "expected error to be returned when running the precompile")
 				s.Require().Nil(resp.Ret, "expected returned bytes to be nil")
 				s.Require().ErrorContains(err, tc.errContains)
-				// Writes once because of gas usage
-				testutil.ValidateWrites(s.T(), cms, 1)
+				// NOTES: After stack-based snapshot mechanism is added for precopile call,
+				// CacheMultiStore.Write() is not called when tx fails.
+				testutil.ValidateWrites(s.T(), cms, 0)
 			}
 		})
 	}

tests/integration/precompiles/staking/test_integration.go

@@ -1834,6 +1834,54 @@ func TestPrecompileIntegrationTestSuite(t *testing.T, create network.CreateEvmAp
 						Expect(err).NotTo(BeNil())
 						Expect(err.Error()).To(ContainSubstring("not found"), "expected NO delegation created")
 					})
+
+					It("should delegate before and after intentionaly ignored delegation revert - successful tx", func() {
+						delegationAmount := math.NewInt(10)
+						expectedDelegationAmount := delegationAmount.Add(delegationAmount)
+
+						callArgs := testutiltypes.CallArgs{
+							ContractABI: stakingReverterContract.ABI,
+							MethodName:  "callPrecompileBeforeAndAfterRevert",
+							Args: []interface{}{
+								big.NewInt(5), s.network.GetValidators()[0].OperatorAddress,
+							},
+						}
+
+						delegateCheck := passCheck.WithExpEvents(staking.EventTypeDelegate, staking.EventTypeDelegate)
+						// Tx should be successful, but no state changes happened
+						res, _, err := s.factory.CallContractAndCheckLogs(
+							s.keyring.GetPrivKey(0),
+							evmtypes.EvmTxArgs{
+								To:       &stkReverterAddr,
+								GasPrice: gasPrice.BigInt(),
+							},
+							callArgs,
+							delegateCheck,
+						)
+						Expect(err).To(BeNil(), "error while calling the smart contract: %v", err)
+						Expect(s.network.NextBlock()).To(BeNil())
+
+						fees := gasPrice.MulRaw(res.GasUsed)
+
+						// delegation should have been created
+						qRes, err := s.grpcHandler.GetDelegation(sdk.AccAddress(stkReverterAddr.Bytes()).String(), s.network.GetValidators()[0].OperatorAddress)
+						Expect(err).To(BeNil())
+						Expect(qRes.DelegationResponse.Delegation.GetDelegatorAddr()).To(Equal(sdk.AccAddress(stkReverterAddr.Bytes()).String()), "expected delegator address is equal to contract address")
+						Expect(qRes.DelegationResponse.Delegation.GetShares().BigInt()).To(Equal(expectedDelegationAmount.BigInt()), "expected different delegation shares")
+
+						// contract balance should be deducted by delegation amount
+						balRes, err := s.grpcHandler.GetBalanceFromBank(stkReverterAddr.Bytes(), s.bondDenom)
+						Expect(err).To(BeNil())
+						contractFinalBalance := balRes.Balance
+						Expect(contractFinalBalance.Amount).To(Equal(contractInitialBalance.Amount.Sub(expectedDelegationAmount)))
+
+						// fees deducted on tx sender.
+						// delegation amount is deducted on contract balance that is previously funded.
+						balRes, err = s.grpcHandler.GetBalanceFromBank(s.keyring.GetAccAddr(0), s.bondDenom)
+						Expect(err).To(BeNil())
+						txSenderFinalBal := balRes.Balance
+						Expect(txSenderFinalBal.Amount).To(Equal(txSenderInitialBal.Amount.Sub(fees)), "expected tx sender balance to be deducted by fees")
+					})
 				})
 
 				Context("Table-driven tests for Delegate method", func() {

tests/integration/precompiles/staking/test_staking.go

@@ -777,7 +777,11 @@ func (s *PrecompileTestSuite) TestCMS() {
 			if tc.expPass {
 				s.Require().NoError(err, "expected no error when running the precompile")
 				s.Require().NotNil(resp.Ret, "expected returned bytes not to be nil")
-				testutil.ValidateWrites(s.T(), cms, 2)
+				// NOTES: After stack-based snapshot mechanism is added for precopile call,
+				// CacheMultiStore.Write() is always called once when tx succeeds.
+				// It is because CacheMultiStore() is not called when creating snapshot for MultiStore,
+				// Count of Write() is not accumulated.
+				testutil.ValidateWrites(s.T(), cms, 1)
 			} else {
 				if tc.expKeeperPass {
 					s.Require().Contains(resp.VmError, tc.errContains,
@@ -786,8 +790,9 @@ func (s *PrecompileTestSuite) TestCMS() {
 					consumed := ctx.GasMeter().GasConsumed()
 					// LessThanOrEqual because the gas is consumed before the error is returned
 					s.Require().LessOrEqual(tc.gas, consumed, "expected gas consumed to be equal to gas limit")
-					// Writes once because of gas usage
-					testutil.ValidateWrites(s.T(), cms, 1)
+					// NOTES: After stack-based snapshot mechanism is added for precopile call,
+					// CacheMultiStore.Write() is not called when tx fails.
+					testutil.ValidateWrites(s.T(), cms, 0)
 				} else {
 					s.Require().Error(err, "expected error to be returned when running the precompile")
 					s.Require().Nil(resp, "expected returned response to be nil")

x/vm/keeper/keeper.go

@@ -40,6 +40,9 @@ type Keeper struct {
 	// key to access the transient store, which is reset on every block during Commit
 	transientKey storetypes.StoreKey
 
+	// KVStore Keys for modules wired to app
+	storeKeys map[string]*storetypes.KVStoreKey
+
 	// the address capable of executing a MsgUpdateParams message. Typically, this should be the x/gov module account.
 	authority sdk.AccAddress
 
@@ -73,6 +76,7 @@ type Keeper struct {
 func NewKeeper(
 	cdc codec.BinaryCodec,
 	storeKey, transientKey storetypes.StoreKey,
+	keys map[string]*storetypes.KVStoreKey,
 	authority sdk.AccAddress,
 	ak types.AccountKeeper,
 	bankKeeper types.BankKeeper,
@@ -106,6 +110,7 @@ func NewKeeper(
 		transientKey:     transientKey,
 		tracer:           tracer,
 		erc20Keeper:      erc20Keeper,
+		storeKeys:        keys,
 	}
 }
 
@@ -351,3 +356,8 @@ func (k Keeper) AddTransientGasUsed(ctx sdk.Context, gasUsed uint64) (uint64, er
 	k.SetTransientGasUsed(ctx, result)
 	return result, nil
 }
+
+// KVStoreKeys returns KVStore keys injected to keeper
+func (k Keeper) KVStoreKeys() map[string]*storetypes.KVStoreKey {
+	return k.storeKeys
+}

x/vm/statedb/interfaces.go

@@ -4,6 +4,8 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/vm"
 
+	storetypes "cosmossdk.io/store/types"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -33,4 +35,8 @@ type Keeper interface {
 	DeleteCode(ctx sdk.Context, codeHash []byte)
 	SetCode(ctx sdk.Context, codeHash []byte, code []byte)
 	DeleteAccount(ctx sdk.Context, addr common.Address) error
+
+	// Getter for injected KVStore keys
+	// It is used for StateDB.snapshotter creation
+	KVStoreKeys() map[string]*storetypes.KVStoreKey
 }

x/vm/statedb/journal.go

@@ -23,8 +23,6 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/holiman/uint256"
 
-	storetypes "cosmossdk.io/store/types"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -146,8 +144,8 @@ type (
 		slot    *common.Hash
 	}
 	precompileCallChange struct {
-		multiStore storetypes.CacheMultiStore
-		events     sdk.Events
+		snapshot int
+		events   sdk.Events
 	}
 	createContractChange struct {
 		account *common.Address
@@ -182,7 +180,7 @@ func (ch createContractChange) Dirtied() *common.Address {
 func (pc precompileCallChange) Revert(s *StateDB) {
 	// rollback multi store from cache ctx to the previous
 	// state stored in the snapshot
-	s.RevertMultiStore(pc.multiStore, pc.events)
+	s.RevertMultiStore(pc.snapshot, pc.events)
 }
 
 func (pc precompileCallChange) Dirtied() *common.Address {

x/vm/statedb/mock_test.go

@@ -11,6 +11,8 @@ import (
 	"github.com/cosmos/evm/x/vm/statedb"
 	"github.com/cosmos/evm/x/vm/types"
 
+	storetypes "cosmossdk.io/store/types"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -115,3 +117,7 @@ func (k MockKeeper) Clone() *MockKeeper {
 	codes := maps.Clone(k.codes)
 	return &MockKeeper{accounts, codes}
 }
+
+func (k MockKeeper) KVStoreKeys() map[string]*storetypes.KVStoreKey {
+	return make(map[string]*storetypes.KVStoreKey)
+}

x/vm/statedb/state_object.go

@@ -9,8 +9,6 @@ import (
 
 	"github.com/cosmos/evm/x/vm/types"
 
-	storetypes "cosmossdk.io/store/types"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -141,10 +139,10 @@ func (s *stateObject) SetBalance(amount *uint256.Int) uint256.Int {
 // AddPrecompileFn appends to the journal an entry
 // with a snapshot of the multi-store and events
 // previous to the precompile call
-func (s *stateObject) AddPrecompileFn(cms storetypes.CacheMultiStore, events sdk.Events) {
+func (s *stateObject) AddPrecompileFn(snapshot int, events sdk.Events) {
 	s.db.journal.append(precompileCallChange{
-		multiStore: cms,
-		events:     events,
+		snapshot: snapshot,
+		events:   events,
 	})
 }
 

x/vm/statedb/statedb.go

@@ -18,6 +18,8 @@ import (
 	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/holiman/uint256"
 
+	"github.com/cosmos/evm/x/vm/store/snapshotmulti"
+	vmstoretypes "github.com/cosmos/evm/x/vm/store/types"
 	"github.com/cosmos/evm/x/vm/types"
 
 	errorsmod "cosmossdk.io/errors"
@@ -49,6 +51,9 @@ type StateDB struct {
 	cacheCtx sdk.Context
 	// writeCache function contains all the changes related to precompile calls.
 	writeCache func()
+	// snapshotter is used for snapshot creation and revert
+	// this snapshot is used for precompile call
+	snapshotter vmstoretypes.Snapshotter
 
 	// Transient storage
 	transientStorage transientStorage
@@ -163,29 +168,17 @@ func (s *StateDB) GetCacheContext() (sdk.Context, error) {
 	return s.cacheCtx, nil
 }
 
-// MultiStoreSnapshot returns a copy of the stateDB CacheMultiStore.
-func (s *StateDB) MultiStoreSnapshot() storetypes.CacheMultiStore {
-	if s.writeCache == nil {
-		err := s.cache()
-		if err != nil {
-			return s.ctx.MultiStore().CacheMultiStore()
-		}
-	}
-	// the cacheCtx multi store is already a CacheMultiStore
-	// so we need to pass a copy of the current state of it
-	cms := s.cacheCtx.MultiStore().(storetypes.CacheMultiStore)
-	snapshot := cms.CacheMultiStore()
-
-	return snapshot
+// MultiStoreSnapshot snapshots stateDB CacheMultiStore
+// and returns snapshot index
+func (s *StateDB) MultiStoreSnapshot() int {
+	return s.snapshotter.Snapshot()
 }
 
-func (s *StateDB) RevertMultiStore(cms storetypes.CacheMultiStore, events sdk.Events) {
-	s.cacheCtx = s.cacheCtx.WithMultiStore(cms)
+func (s *StateDB) RevertMultiStore(snapshot int, events sdk.Events) {
+	s.snapshotter.RevertToSnapshot(snapshot)
 	s.writeCache = func() {
-		// rollback the events to the ones
-		// on the snapshot
 		s.ctx.EventManager().EmitEvents(events)
-		cms.Write()
+		s.cacheCtx.MultiStore().(storetypes.CacheMultiStore).Write()
 	}
 }
 
@@ -194,7 +187,21 @@ func (s *StateDB) cache() error {
 	if s.ctx.MultiStore() == nil {
 		return errors.New("ctx has no multi store")
 	}
-	s.cacheCtx, s.writeCache = s.ctx.CacheContext()
+	s.cacheCtx, _ = s.ctx.CacheContext()
+
+	// Get KVStores for modules wired to app
+	cms := s.cacheCtx.MultiStore().(storetypes.CacheMultiStore)
+	storeKeys := s.keeper.KVStoreKeys()
+
+	// Create and set snapshot store to stateDB
+	snapshotStore := snapshotmulti.NewStore(cms, storeKeys)
+	s.snapshotter = snapshotStore
+	s.cacheCtx = s.cacheCtx.WithMultiStore(snapshotStore)
+	s.writeCache = func() {
+		s.ctx.EventManager().EmitEvents(s.cacheCtx.EventManager().Events())
+		s.cacheCtx.MultiStore().(storetypes.CacheMultiStore).Write()
+	}
+
 	return nil
 }
 
@@ -409,12 +416,12 @@ func (s *StateDB) setStateObject(object *stateObject) {
 // AddPrecompileFn adds a precompileCall journal entry
 // with a snapshot of the multi-store and events previous
 // to the precompile call.
-func (s *StateDB) AddPrecompileFn(addr common.Address, cms storetypes.CacheMultiStore, events sdk.Events) error {
+func (s *StateDB) AddPrecompileFn(addr common.Address, snapshot int, events sdk.Events) error {
 	stateObject := s.getOrNewStateObject(addr)
 	if stateObject == nil {
 		return fmt.Errorf("could not add precompile call to address %s. State object not found", addr)
 	}
-	stateObject.AddPrecompileFn(cms, events)
+	stateObject.AddPrecompileFn(snapshot, events)
 	s.precompileCallsCounter++
 	if s.precompileCallsCounter > types.MaxPrecompileCalls {
 		return fmt.Errorf("max calls to precompiles (%d) reached", types.MaxPrecompileCalls)