Problem: state override not replacing the whole state

tests/integration_tests/test_call.py

@@ -4,10 +4,10 @@
 from web3 import Web3
 from web3._utils.contracts import encode_transaction_data
 
-from .utils import CONTRACTS
+from .utils import CONTRACTS, deploy_contract
 
 
-def test_state_override(ethermint):
+def test_temporary_contract_code(ethermint):
     state = 100
     w3: Web3 = ethermint.w3
     info = json.loads(CONTRACTS["Greeter"].read_text())
@@ -33,3 +33,65 @@ def test_state_override(ethermint):
         overrides,
     )
     assert (state,) == w3.codec.decode(("uint256",), result)
+
+
+def test_override_state(ethermint):
+    w3: Web3 = ethermint.w3
+    contract, _ = deploy_contract(w3, CONTRACTS["Greeter"])
+
+    assert "Hello" == contract.functions.greet().call()
+    assert 0 == contract.functions.intValue().call()
+
+    info = json.loads(CONTRACTS["Greeter"].read_text())
+    int_value = 100
+    state = {
+        ("0x" + "0" * 64): HexBytes(
+            w3.codec.encode(("uint256",), (int_value,))
+        ).hex(),
+    }
+    result = w3.eth.call(
+        {
+            "to": contract.address,
+            "data": encode_transaction_data(w3, "intValue", info["abi"]),
+        },
+        "latest",
+        {
+            contract.address: {
+                "code": info["deployedBytecode"],
+                "stateDiff": state,
+            },
+        },
+    )
+    assert (int_value,) == w3.codec.decode(("uint256",), result)
+
+    # stateDiff don't affect the other state slots
+    result = w3.eth.call(
+        {
+            "to": contract.address,
+            "data": encode_transaction_data(w3, "greet", info["abi"]),
+        },
+        "latest",
+        {
+            contract.address: {
+                "code": info["deployedBytecode"],
+                "stateDiff": state,
+            },
+        },
+    )
+    assert ("Hello",) == w3.codec.decode(("string",), result)
+
+    # state will overrides the whole state
+    result = w3.eth.call(
+        {
+            "to": contract.address,
+            "data": encode_transaction_data(w3, "greet", info["abi"]),
+        },
+        "latest",
+        {
+            contract.address: {
+                "code": info["deployedBytecode"],
+                "state": state,
+            },
+        },
+    )
+    assert ("",) == w3.codec.decode(("string",), result)

x/evm/statedb/state_object.go

@@ -71,6 +71,9 @@ type stateObject struct {
 	// state storage
 	originStorage Storage
 	dirtyStorage  Storage
+	// overridden state, when not nil, replace the whole committed state,
+	// mainly to support the stateOverrides in eth_call.
+	overrideStorage Storage
 
 	address common.Address
 
@@ -172,6 +175,13 @@ func (s *stateObject) Nonce() uint64 {
 
 // GetCommittedState query the committed state
 func (s *stateObject) GetCommittedState(key common.Hash) common.Hash {
+	if s.overrideStorage != nil {
+		if value, ok := s.overrideStorage[key]; ok {
+			return value
+		}
+		return common.Hash{}
+	}
+
 	if value, cached := s.originStorage[key]; cached {
 		return value
 	}
@@ -205,6 +215,12 @@ func (s *stateObject) SetState(key common.Hash, value common.Hash) {
 	s.setState(key, value)
 }
 
+func (s *stateObject) SetStorage(storage Storage) {
+	s.overrideStorage = storage
+	s.originStorage = make(Storage)
+	s.dirtyStorage = make(Storage)
+}
+
 func (s *stateObject) setState(key, value common.Hash) {
 	s.dirtyStorage[key] = value
 }

x/evm/statedb/statedb.go

@@ -418,11 +418,9 @@ func (s *StateDB) SetState(addr common.Address, key, value common.Hash) {
 // SetStorage replaces the entire storage for the specified account with given
 // storage. This function should only be used for debugging and the mutations
 // must be discarded afterwards.
-func (s *StateDB) SetStorage(addr common.Address, storage map[common.Hash]common.Hash) {
+func (s *StateDB) SetStorage(addr common.Address, storage Storage) {
 	stateObject := s.getOrNewStateObject(addr)
-	for k, v := range storage {
-		stateObject.SetState(k, v)
-	}
+	stateObject.SetStorage(storage)
 }
 
 // Suicide marks the given account as suicided.

x/evm/statedb/statedb_test.go

@@ -747,6 +747,30 @@ func (suite *StateDBTestSuite) TestNativeAction() {
 	suite.Require().Equal(sdk.Events{{Type: "success1"}, {Type: "success3"}}, ctx.EventManager().Events())
 }
 
+func (suite *StateDBTestSuite) TestSetStorage() {
+	contract := common.BigToAddress(big.NewInt(101))
+
+	_, ctx, keeper := setupTestEnv(suite.T())
+	stateDB := statedb.New(ctx, keeper, emptyTxConfig)
+
+	stateDB.SetState(contract, common.BigToHash(big.NewInt(0)), common.BigToHash(big.NewInt(0)))
+	stateDB.SetState(contract, common.BigToHash(big.NewInt(1)), common.BigToHash(big.NewInt(1)))
+	stateDB.SetState(contract, common.BigToHash(big.NewInt(2)), common.BigToHash(big.NewInt(2)))
+	suite.Require().NoError(stateDB.Commit())
+
+	suite.Require().Equal(common.BigToHash(big.NewInt(0)), stateDB.GetState(contract, common.BigToHash(big.NewInt(0))))
+	suite.Require().Equal(common.BigToHash(big.NewInt(1)), stateDB.GetState(contract, common.BigToHash(big.NewInt(1))))
+	suite.Require().Equal(common.BigToHash(big.NewInt(2)), stateDB.GetState(contract, common.BigToHash(big.NewInt(2))))
+
+	stateDB.SetStorage(contract, map[common.Hash]common.Hash{
+		common.BigToHash(big.NewInt(1)): common.BigToHash(big.NewInt(3)),
+	})
+
+	suite.Require().Equal(common.Hash{}, stateDB.GetState(contract, common.BigToHash(big.NewInt(0))))
+	suite.Require().Equal(common.BigToHash(big.NewInt(3)), stateDB.GetState(contract, common.BigToHash(big.NewInt(1))))
+	suite.Require().Equal(common.Hash{}, stateDB.GetState(contract, common.BigToHash(big.NewInt(2))))
+}
+
 func CollectContractStorage(db vm.StateDB, address common.Address) statedb.Storage {
 	storage := make(statedb.Storage)
 	db.ForEachStorage(address, func(k, v common.Hash) bool {