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

* feat(x/vm): add cacheStack for CacheMultiStore snapshot

* fix(x/vm): StateDB.snapshotter

* feat(x/vm): add keys field to x/vm keeper and use keys for snapshot

* deps: remove cosmossdk.io/store fork

* refactor(x/vm): refactor snapshot multi store and add comments

* test(x/vm): add benchmark test for snapshotmulti.Store

* test(precompiles): add edge case test for precompile

* refactor(x/vm): modify snapshot stores

dep(evmd): go mod tidy

refactor: change package name of stack based store

refactor(x/vm/store): modify type casting

refactor(x/vm): rename store packages

* fix(x/vm): order store keys for snapshot store

* chore(x/vm/store): modify comments

* chore(x/vm/store): add comments

* test(x/vm/store): add unit tests

* chore(x/vm): fix lint

* chore: fix lint

* chore: fix lint

* chore(x/vm/store): remove unnecessary function call

* test: modify TestCMS

* fix(tests): precompile test case that intermittently fails

* chore(x/vm/store): rename variable cms into snapshotStore in test code

* chore: resolve merge conflict

* chore(x/vm) fix lint

* chore: fix typo

* chore: modify comment

* chore(x/vm/store): add comments

* test(x/vm/store): add test case for overwrite to same key

Author: cloudgray

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"
@@ -481,7 +482,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.json

@@ -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();

precompiles/testutil/contracts/StakingReverter.sol

@@ -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 precompile 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 precompile call,
+				// CacheMultiStore.Write() is not called when tx fails.
+				testutil.ValidateWrites(s.T(), cms, 0)
 			}
 		})
 	}

tests/integration/precompiles/distribution/test_distribution.go

@@ -2274,7 +2274,7 @@ func TestPrecompileIntegrationTestSuite(t *testing.T, create network.CreateEvmAp
 
 				// set gas such that the internal keeper function called by the precompile fails out mid-execution
 				txArgs.GasLimit = 80_000
-				_, _, err = s.factory.CallContractAndCheckLogs(
+				_, txRes, err := s.factory.CallContractAndCheckLogs(
 					s.keyring.GetPrivKey(0),
 					txArgs,
 					callArgs,
@@ -2286,7 +2286,8 @@ func TestPrecompileIntegrationTestSuite(t *testing.T, create network.CreateEvmAp
 				balRes, err := s.grpcHandler.GetBalanceFromBank(s.keyring.GetAccAddr(0), s.bondDenom)
 				Expect(err).To(BeNil())
 				finalBalance := balRes.Balance
-				expectedGasCost := math.NewInt(79_416_000_000_000)
+
+				expectedGasCost := math.NewIntFromUint64(txRes.GasUsed).Mul(math.NewIntFromBigInt(txArgs.GasPrice))
 				Expect(finalBalance.Amount.Equal(initialBalance.Amount.Sub(expectedGasCost))).To(BeTrue(), "expected final balance must be initial balance minus any gas spent")
 
 				res, err = s.grpcHandler.GetDelegationTotalRewards(s.keyring.GetAccAddr(0).String())

tests/integration/precompiles/distribution/test_integration.go

@@ -1834,6 +1834,57 @@ 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)
+
+						// The transaction should succeed with delegations occurring both before and after the intended revert.
+						// The revert itself is not propagated because it occurs within the scope of a try-catch statement,
+						// but is not caught by the catch block.
+						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_integration.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 precompile 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 precompile 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")

tests/integration/precompiles/staking/test_staking.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
+}