evm: use clean context for GetCommittedState (#383)

* keep the original context for GetCommittedState api

* fix method mutation

* keep estimateGas consistant

* added test after the original context is recovered

* add integration test for the gas consumption of sstore

* test the committed case

* move methods to keeper module

Co-authored-by: Federico Kunze Küllmer <31522760+fedekunze@users.noreply.github.com>

tests/solidity/suites/basic/contracts/Storage.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity ^0.5.11;
+
+/**
+ * @title Storage
+ * @dev Store & retrieve value in a variable
+ */
+contract Storage {
+
+    uint256 number;
+
+    /**
+     * @dev Store value in variable
+     * @param num value to store
+     */
+    function store(uint256 num) public {
+        number = num + 1;
+        number = num;
+    }
+
+    /**
+     * @dev Return value 
+     * @return value of 'number'
+     */
+    function retrieve() public view returns (uint256){
+        return number;
+    }
+}

tests/solidity/suites/basic/test/estimateGas.js

@@ -0,0 +1,22 @@
+const Storage = artifacts.require("Storage")
+
+contract('Storage', (accounts) => {
+
+  let storage
+  beforeEach(async () => {
+    storage = await Storage.new()
+  })
+
+  it('estimated gas should match', async () => {
+      // set new value
+      let gasUsage = await storage.store.estimateGas(10);
+      expect(gasUsage.toString()).to.equal('43754');
+
+      await storage.store(10);
+
+      // set existing value
+      gasUsage = await storage.store.estimateGas(10);
+      expect(gasUsage.toString()).to.equal('28754');
+  })
+
+})

x/evm/keeper/grpc_query.go

@@ -444,15 +444,14 @@ func (k Keeper) EstimateGas(c context.Context, req *types.EthCallRequest) (*type
 		args.Gas = (*hexutil.Uint64)(&gas)
 
 		// Execute the call in an isolated context
-		sandboxCtx, _ := ctx.CacheContext()
-		k.WithContext(sandboxCtx)
+		k.BeginCachedContext()
 
 		msg := args.ToMessage(req.GasCap)
 		evm := k.NewEVM(msg, ethCfg, params, coinbase)
 		// pass true means execute in query mode, which don't do actual gas refund.
 		rsp, err := k.ApplyMessage(evm, msg, ethCfg, true)
 
-		k.WithContext(ctx)
+		k.EndCachedContext()
 
 		if err != nil {
 			if errors.Is(stacktrace.RootCause(err), core.ErrIntrinsicGas) {

x/evm/keeper/keeper.go

@@ -38,6 +38,8 @@ type Keeper struct {
 	stakingKeeper types.StakingKeeper
 
 	ctx sdk.Context
+	// set in `BeginCachedContext`, used by `GetCommittedState` api.
+	committedCtx sdk.Context
 
 	// chain ID number obtained from the context's chain id
 	eip155ChainID *big.Int
@@ -75,6 +77,11 @@ func NewKeeper(
 	}
 }
 
+// CommittedCtx returns the committed context
+func (k Keeper) CommittedCtx() sdk.Context {
+	return k.committedCtx
+}
+
 // Logger returns a module-specific logger.
 func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 	return ctx.Logger().With("module", types.ModuleName)
@@ -83,6 +90,7 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 // WithContext sets an updated SDK context to the keeper
 func (k *Keeper) WithContext(ctx sdk.Context) {
 	k.ctx = ctx
+	k.committedCtx = ctx
 }
 
 // WithChainID sets the chain id to the local variable in the keeper
@@ -341,3 +349,15 @@ func (k Keeper) ResetAccount(addr common.Address) {
 	k.DeleteCode(addr)
 	k.DeleteAccountStorage(addr)
 }
+
+// BeginCachedContext create the cached context
+func (k *Keeper) BeginCachedContext() (commit func()) {
+	k.committedCtx = k.ctx
+	k.ctx, commit = k.ctx.CacheContext()
+	return
+}
+
+// EndCachedContext recover the committed context
+func (k *Keeper) EndCachedContext() {
+	k.ctx = k.committedCtx
+}

x/evm/keeper/state_transition.go

@@ -137,11 +137,8 @@ func (k *Keeper) ApplyTransaction(tx *ethtypes.Transaction) (*types.MsgEthereumT
 		return nil, stacktrace.Propagate(err, "failed to return ethereum transaction as core message")
 	}
 
-	// create an ethereum StateTransition instance and run TransitionDb
-	// we use a ctx context to avoid modifying to state in case EVM msg is reverted
-	originalCtx := k.ctx
-	cacheCtx, commit := k.ctx.CacheContext()
-	k.ctx = cacheCtx
+	// we use a cached context to avoid modifying to state in case EVM msg is reverted
+	commit := k.BeginCachedContext()
 
 	// get the coinbase address from the block proposer
 	coinbase, err := k.GetCoinbaseAddress()
@@ -149,6 +146,7 @@ func (k *Keeper) ApplyTransaction(tx *ethtypes.Transaction) (*types.MsgEthereumT
 		return nil, stacktrace.Propagate(err, "failed to obtain coinbase address")
 	}
 
+	// create an ethereum EVM instance and run the message
 	evm := k.NewEVM(msg, ethCfg, params, coinbase)
 
 	k.SetTxHashTransient(tx.Hash())
@@ -163,11 +161,12 @@ func (k *Keeper) ApplyTransaction(tx *ethtypes.Transaction) (*types.MsgEthereumT
 	res.Hash = txHash.Hex()
 	logs := k.GetTxLogs(txHash)
 
-	// Commit and switch to original context
+	// Commit and switch to committed context
 	if !res.Failed() {
 		commit()
 	}
-	k.ctx = originalCtx
+
+	k.EndCachedContext()
 
 	// Logs needs to be ignored when tx is reverted
 	// Set the log and bloom filter only when the tx is NOT REVERTED

x/evm/keeper/statedb.go

@@ -352,10 +352,8 @@ func (k *Keeper) GetRefund() uint64 {
 // State
 // ----------------------------------------------------------------------------
 
-// GetCommittedState returns the value set in store for the given key hash. If the key is not registered
-// this function returns the empty hash.
-func (k *Keeper) GetCommittedState(addr common.Address, hash common.Hash) common.Hash {
-	store := prefix.NewStore(k.ctx.KVStore(k.storeKey), types.AddressStoragePrefix(addr))
+func doGetState(ctx sdk.Context, storeKey sdk.StoreKey, addr common.Address, hash common.Hash) common.Hash {
+	store := prefix.NewStore(ctx.KVStore(storeKey), types.AddressStoragePrefix(addr))
 
 	key := types.KeyAddressStorage(addr, hash)
 	value := store.Get(key.Bytes())
@@ -366,10 +364,16 @@ func (k *Keeper) GetCommittedState(addr common.Address, hash common.Hash) common
 	return common.BytesToHash(value)
 }
 
+// GetCommittedState returns the value set in store for the given key hash. If the key is not registered
+// this function returns the empty hash.
+func (k *Keeper) GetCommittedState(addr common.Address, hash common.Hash) common.Hash {
+	return doGetState(k.committedCtx, k.storeKey, addr, hash)
+}
+
 // GetState returns the committed state for the given key hash, as all changes are committed directly
 // to the KVStore.
 func (k *Keeper) GetState(addr common.Address, hash common.Hash) common.Hash {
-	return k.GetCommittedState(addr, hash)
+	return doGetState(k.ctx, k.storeKey, addr, hash)
 }
 
 // SetState sets the given hashes (key, value) to the KVStore. If the value hash is empty, this

x/evm/keeper/statedb_test.go

@@ -383,6 +383,30 @@ func (suite *KeeperTestSuite) TestState() {
 	}
 }
 
+func (suite *KeeperTestSuite) TestCommittedState() {
+	suite.SetupTest()
+
+	var key = common.BytesToHash([]byte("key"))
+	var value1 = common.BytesToHash([]byte("value1"))
+	var value2 = common.BytesToHash([]byte("value2"))
+
+	suite.app.EvmKeeper.SetState(suite.address, key, value1)
+
+	commit := suite.app.EvmKeeper.BeginCachedContext()
+
+	suite.app.EvmKeeper.SetState(suite.address, key, value2)
+	tmp := suite.app.EvmKeeper.GetState(suite.address, key)
+	suite.Require().Equal(value2, tmp)
+	tmp = suite.app.EvmKeeper.GetCommittedState(suite.address, key)
+	suite.Require().Equal(value1, tmp)
+
+	commit()
+	suite.app.EvmKeeper.EndCachedContext()
+
+	tmp = suite.app.EvmKeeper.GetCommittedState(suite.address, key)
+	suite.Require().Equal(value2, tmp)
+}
+
 func (suite *KeeperTestSuite) TestSuicide() {
 	testCases := []struct {
 		name     string