Add insert storage

Cargo.toml

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

pyproject.toml

@@ -9,7 +9,7 @@ classifiers = [
 
 [tool.poetry]
 name = "pyrevm"
-version = "0.3.3"  # Cargo.toml needs to be updated for releases
+version = "0.3.6"  # Cargo.toml needs to be updated for releases
 description = "Python bindings to rust evm (revm)"
 authors = ["Georgios Konstantopoulos <[email protected]>", "Dani Popes", "Daniel Schiavini", "Charles Cooper"]
 

pyrevm.pyi

@@ -201,6 +201,14 @@ class EVM:
         :param info: The account info.
         """
 
+    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.
+        :param index: slot in storage.
+        :param value: value for slot.
+        """
+
     def message_call(
         self: "EVM",
         caller: str,

src/database.rs

@@ -47,6 +47,22 @@ impl DB {
         }
     }
 
+    pub(crate) fn insert_insert_account_storage(
+        &mut self,
+        address: Address,
+        slot: U256,
+        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),
+        }
+    }
+
     pub(crate) fn get_accounts(&self) -> &HashMap<Address, DbAccount> {
         match self {
             DB::Memory(db) => &db.accounts,

src/evm.rs

@@ -52,7 +52,8 @@ pub struct EVM {
 impl EVM {
     /// Create a new EVM instance.
     #[new]
-    #[pyo3(signature = (env = None, fork_url = None, fork_block = None, gas_limit = 18446744073709551615, tracing = false, spec_id = "LATEST"))]
+    #[pyo3(signature = (env = None, fork_url = None, fork_block = None, gas_limit = 18446744073709551615, tracing = false, spec_id = "LATEST")
+    )]
     fn new(
         env: Option<Env>,
         fork_url: Option<&str>,
@@ -151,6 +152,33 @@ impl EVM {
         Ok(())
     }
 
+    /// 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<U256> {
+        let target = addr(address)?;
+
+        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.
     fn set_balance(&mut self, address: &str, balance: U256) -> PyResult<()> {
         let address = addr(address)?;
@@ -166,7 +194,8 @@ impl EVM {
         Ok(balance)
     }
 
-    #[pyo3(signature = (caller, to, calldata = None, value = None, gas = None, gas_price = None, is_static = false))]
+    #[pyo3(signature = (caller, to, calldata = None, value = None, gas = None, gas_price = None, is_static = false)
+    )]
     pub fn message_call(
         &mut self,
         caller: &str,
@@ -193,7 +222,8 @@ impl EVM {
     }
 
     /// Deploy a contract with the given code.
-    #[pyo3(signature = (deployer, code, value = None, gas = None, gas_price = None, is_static = false, _abi = None))]
+    #[pyo3(signature = (deployer, code, value = None, gas = None, gas_price = None, is_static = false, _abi = None)
+    )]
     fn deploy(
         &mut self,
         deployer: &str,

tests/test_evm.py

@@ -11,8 +11,8 @@
 
 # use your own key during development to avoid rate limiting the CI job
 fork_url = (
-    os.getenv("FORK_URL")
-    or "https://mainnet.infura.io/v3/c60b0bb42f8a4c6481ecd229eddaca27"
+        os.getenv("FORK_URL")
+        or "https://mainnet.infura.io/v3/c60b0bb42f8a4c6481ecd229eddaca27"
 )
 
 KWARG_CASES = [
@@ -60,6 +60,33 @@ def test_fork_storage():
     assert value > 0
 
 
+def test_set_into_storage():
+    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
+
+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()
 
@@ -265,7 +292,7 @@ def test_tx_setters():
 
 @pytest.mark.parametrize(
     "excess_blob_gas,expected_fee",
-    [(0, 1), (10**3, 1), (2**24, 152), (2**26, 537070730)],
+    [(0, 1), (10 ** 3, 1), (2 ** 24, 152), (2 ** 26, 537070730)],
 )
 def test_get_blobbasefee(excess_blob_gas, expected_fee):
     evm = EVM()