Fix insert storage overwrite

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyrevm"
-version = "0.3.4"
+version = "0.3.5"
 edition = "2021"
 
 [lib]

pyrevm.pyi

@@ -201,7 +201,7 @@ class EVM:
         :param info: The account info.
         """
 
-    def insert_account_storage(self: "EVM", address: str, index: int, value: int) -> None:
+    def insert_account_storage(self: "EVM", address: str, index: int, value: int) -> int:
         """
         Inserts the provided value for slot of in the database at the specified address
         :param address: The address of the account.

src/database.rs

@@ -54,10 +54,12 @@ impl DB {
         value: U256,
     ) -> PyResult<()> {
         match self {
-            DB::Memory(db) =>
-                db.insert_account_storage(address, slot, value).map_err(pyerr),
-            DB::Fork(db) =>
-                db.insert_account_storage(address, slot, value).map_err(pyerr),
+            DB::Memory(db) => db
+                .insert_account_storage(address, slot, value)
+                .map_err(pyerr),
+            DB::Fork(db) => db
+                .insert_account_storage(address, slot, value)
+                .map_err(pyerr),
         }
     }
 

src/evm.rs

@@ -153,9 +153,30 @@ impl EVM {
     }
 
     /// Inserts the provided value for slot of in the database at the specified address
-    fn insert_account_storage(&mut self, address: &str, index: U256, value: U256) -> PyResult<()> {
+    fn insert_account_storage(
+        &mut self,
+        address: &str,
+        index: U256,
+        value: U256,
+    ) -> PyResult<U256> {
         let target = addr(address)?;
-        self.context.db.insert_insert_account_storage(target, index, value)
+
+        match self.context.journaled_state.state.get_mut(&target) {
+            // account is cold, just insert into the DB
+            None => {
+                self.context
+                    .db
+                    .insert_insert_account_storage(target, index, value)
+                    .map_err(pyerr)?;
+                self.context.load_account(target).map_err(pyerr)?;
+                Ok(U256::ZERO)
+            }
+            // just replace old value
+            Some(_) => {
+                let store_result = self.context.sstore(target, index, value).map_err(pyerr)?;
+                Ok(store_result.original_value)
+            }
+        }
     }
 
     /// Set the balance of a given address.

tests/test_evm.py

@@ -67,6 +67,25 @@ def test_set_into_storage():
     value = evm.storage(weth, 0)
     assert value == 10
 
+def test_set_into_storage_with_update():
+    weth = "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
+    evm = EVM(fork_url=fork_url, fork_block="latest")
+    evm.insert_account_storage(weth, 0, 10)
+    value = evm.storage(weth, 0)
+    assert value == 10
+    evm.insert_account_storage(weth, 0, 20)
+    value = evm.storage(weth, 0)
+    assert value == 20
+
+def test_set_into_storage_old_value():
+    weth = "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
+    evm = EVM(fork_url=fork_url, fork_block="latest")
+    old_value = evm.insert_account_storage(weth, 0, 10)
+    assert old_value == 0
+    old_value = evm.insert_account_storage(weth, 0, 20)
+    assert old_value == 10
+
+
 
 def test_deploy():
     evm = EVM()