feat: unchecked ext trait and with_parts_unchecked (#22)

Author: prestwich

src/evm.rs

@@ -1,16 +1,15 @@
 use crate::{
     states::EvmBlockComplete, BasicContext, Block, BlockComplete, BlockContext, Cfg, ErroredState,
-    EvmErrored, EvmNeedsCfg, EvmNeedsFirstBlock, EvmNeedsNextBlock, EvmNeedsTx, EvmReady,
-    EvmTransacted, HasCfg, HasOutputs, NeedsBlock, NeedsCfg, NeedsNextBlock, NeedsTx, Ready,
-    TransactedState, Tx,
+    EvmErrored, EvmExtUnchecked, EvmNeedsCfg, EvmNeedsFirstBlock, EvmNeedsNextBlock, EvmNeedsTx,
+    EvmReady, EvmTransacted, HasCfg, HasContext, HasOutputs, NeedsBlock, NeedsCfg, NeedsNextBlock,
+    NeedsTx, Ready, TransactedState, Tx,
 };
-use alloy_consensus::constants::KECCAK_EMPTY;
 use alloy_primitives::{Address, Bytes, U256};
 use revm::{
     db::{states::bundle_state::BundleRetention, BundleState},
     primitives::{
-        Account, AccountInfo, AccountStatus, Bytecode, EVMError, EvmState, EvmStorageSlot,
-        ExecutionResult, InvalidTransaction, ResultAndState, SpecId,
+        AccountInfo, Bytecode, EVMError, EvmState, ExecutionResult, InvalidTransaction,
+        ResultAndState, SpecId,
     },
     Database, DatabaseCommit, DatabaseRef, Evm, State,
 };
@@ -252,21 +251,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         address: Address,
         f: F,
     ) -> Result<AccountInfo, Db::Error> {
-        let db = self.inner_mut_unchecked().db_mut();
-
-        let mut info = db.basic(address)?.unwrap_or_default();
-        let old = info.clone();
-        f(&mut info);
-
-        // Make a new account with the modified info
-        let mut acct = Account { info, status: AccountStatus::Touched, ..Default::default() };
-        acct.mark_touch();
-
-        // Create a state object with the modified account.
-        let state = [(address, acct)].iter().cloned().collect();
-        self.commit_unchecked(state);
-
-        Ok(old)
+        self.inner.modify_account(address, f)
     }
 
     /// Set the nonce of an account, returning the previous nonce. This is a
@@ -276,7 +261,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         address: Address,
         nonce: u64,
     ) -> Result<u64, Db::Error> {
-        self.try_modify_account_unchecked(address, |info| info.nonce = nonce).map(|info| info.nonce)
+        self.inner.set_nonce(address, nonce)
     }
 
     /// Increment the nonce of an account, returning the previous nonce. This is
@@ -285,17 +270,15 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
     /// If the nonce is already at the maximum value, it will not be
     /// incremented.
     pub fn try_increment_nonce_unchecked(&mut self, address: Address) -> Result<u64, Db::Error> {
-        self.try_modify_account_unchecked(address, |info| info.nonce = info.nonce.saturating_add(1))
-            .map(|info| info.nonce)
+        self.inner.increment_nonce(address)
     }
 
     /// Decrement the nonce of an account, returning the previous nonce. This is
     /// a low-level API, and is not intended for general use.
     ///
     /// If the nonce is already 0, it will not be decremented.
     pub fn try_decrement_nonce_unchecked(&mut self, address: Address) -> Result<u64, Db::Error> {
-        self.try_modify_account_unchecked(address, |info| info.nonce = info.nonce.saturating_sub(1))
-            .map(|info| info.nonce)
+        self.inner.decrement_nonce(address)
     }
 
     /// Set the EVM storage at a slot. This is a low-level API, and is not
@@ -306,22 +289,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         slot: U256,
         value: U256,
     ) -> Result<U256, Db::Error> {
-        let db = self.inner_mut_unchecked().db_mut();
-        let info = db.basic(address)?.unwrap_or_default();
-        let old = db.storage(address, slot)?;
-
-        let change = EvmStorageSlot::new_changed(old, value);
-
-        // Make a new account with the modified storage
-        let storage = [(slot, change)].iter().cloned().collect();
-        let mut acct = Account { storage, info, ..Default::default() };
-        acct.mark_touch();
-
-        // Create a state object with the modified account.
-        let state = [(address, acct)].iter().cloned().collect();
-        self.commit_unchecked(state);
-
-        Ok(old)
+        self.inner.set_storage(address, slot, value)
     }
 
     /// Set the bytecode at a specific address, returning the previous bytecode
@@ -332,24 +300,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         address: Address,
         bytecode: Bytecode,
     ) -> Result<Option<Bytecode>, Db::Error> {
-        let db = self.inner_mut_unchecked().db_mut();
-        let mut info = db.basic(address)?.unwrap_or_default();
-
-        let old = if info.code_hash != KECCAK_EMPTY {
-            Some(db.code_by_hash(info.code_hash)?)
-        } else {
-            None
-        };
-
-        info.code_hash = if bytecode.is_empty() { KECCAK_EMPTY } else { bytecode.hash_slow() };
-        info.code = Some(bytecode);
-
-        let mut acct = Account { info, ..Default::default() };
-        acct.mark_touch();
-        let state = [(address, acct)].iter().cloned().collect();
-        self.commit_unchecked(state);
-
-        Ok(old)
+        self.inner.set_bytecode(address, bytecode)
     }
 
     /// Increase the balance of an account. Returns the previous balance. This
@@ -362,10 +313,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         address: Address,
         amount: U256,
     ) -> Result<U256, Db::Error> {
-        self.try_modify_account_unchecked(address, |info| {
-            info.balance = info.balance.saturating_add(amount)
-        })
-        .map(|info| info.balance)
+        self.inner.increase_balance(address, amount)
     }
 
     /// Decrease the balance of an account. Returns the previous balance. This
@@ -377,10 +325,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         address: Address,
         amount: U256,
     ) -> Result<U256, Db::Error> {
-        self.try_modify_account_unchecked(address, |info| {
-            info.balance = info.balance.saturating_sub(amount)
-        })
-        .map(|info| info.balance)
+        self.inner.decrease_balance(address, amount)
     }
 
     /// Set the balance of an account. Returns the previous balance. This is a
@@ -390,8 +335,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState> Trevm<'a, Ext, Db, Trev
         address: Address,
         amount: U256,
     ) -> Result<U256, Db::Error> {
-        self.try_modify_account_unchecked(address, |info| info.balance = amount)
-            .map(|info| info.balance)
+        self.inner.set_balance(address, amount)
     }
 }
 
@@ -798,6 +742,8 @@ impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState: NeedsBlock>
     }
 }
 
+// --- HAS OUTPUTS
+
 impl<'a, Ext, Db: Database + DatabaseCommit, Missing: HasOutputs>
     Trevm<'a, Ext, State<Db>, Missing>
 {
@@ -820,6 +766,7 @@ impl<'a, Ext, Db: Database + DatabaseCommit, Missing: HasOutputs>
         bundle
     }
 }
+// --- BLOCK COMPLETE
 
 impl<'a, Ext, Db: Database + DatabaseCommit, C> EvmBlockComplete<'a, Ext, Db, C> {
     /// Destructure the EVM and return the block context and the EVM ready for
@@ -834,6 +781,36 @@ impl<'a, Ext, Db: Database + DatabaseCommit, C> EvmBlockComplete<'a, Ext, Db, C>
     }
 }
 
+// --- HAS CONTEXT
+
+impl<'a, Ext, Db: Database + DatabaseCommit, TrevmState: HasContext>
+    Trevm<'a, Ext, Db, TrevmState>
+{
+    /// Get a reference to the block context.
+    pub fn block_context(&self) -> &TrevmState::Context {
+        self.state.context()
+    }
+
+    /// Get a mutable reference to the block context.
+    pub fn block_context_mut(&mut self) -> &mut TrevmState::Context {
+        self.state.context_mut()
+    }
+
+    /// Apply a closure to the components. This allows for modifying the block
+    /// context and the EVM in a single operation, by (e.g.) passing the EVM
+    /// to a block context method.
+    ///
+    /// This is a low-level API and is not intended for general use.
+    pub fn with_parts_unchecked<F>(self, f: F) -> Self
+    where
+        F: FnOnce(&mut TrevmState, &mut Evm<'a, Ext, Db>),
+    {
+        let Trevm { mut inner, mut state } = self;
+        f(&mut state, &mut inner);
+        Trevm { inner, state }
+    }
+}
+
 // --- NEEDS FIRST TX
 
 impl<'a, Ext, Db: Database + DatabaseCommit, C: BlockContext<Ext, Db>> EvmNeedsTx<'a, Ext, Db, C> {

src/ext.rs

@@ -0,0 +1,177 @@
+use std::collections::HashMap;
+
+use alloy_primitives::{Address, B256, U256};
+use revm::{
+    primitives::{Account, AccountInfo, Bytecode, EvmStorageSlot},
+    Database, DatabaseCommit,
+};
+
+/// Extension trait for [`revm::Evm`] with convenience functions for reading
+/// and modifying state.
+pub trait EvmExtUnchecked<Db: Database> {
+    /// Get a mutable reference to the database.
+    fn db_mut_ext(&mut self) -> &mut Db;
+
+    /// Load an account from the database.
+    fn account(&mut self, address: Address) -> Result<Account, Db::Error> {
+        let info = self.db_mut_ext().basic(address)?;
+        let created = info.is_none();
+        let mut acct = Account { info: info.unwrap_or_default(), ..Default::default() };
+        if created {
+            acct.mark_created();
+        }
+        Ok(acct)
+    }
+
+    /// Load a storage slot from an account.
+    fn storage(&mut self, address: Address, index: U256) -> Result<U256, Db::Error> {
+        self.db_mut_ext().storage(address, index)
+    }
+
+    /// Load the bytecode for a contract. Returns empty bytecode if the account
+    /// does not currently have contract code.
+    fn bytecode(&mut self, code_hash: B256) -> Result<Bytecode, Db::Error> {
+        self.db_mut_ext().code_by_hash(code_hash)
+    }
+
+    /// Modify multiple accounts with a closure and commit the modified
+    /// accounts.
+    fn modify_accounts<I, F>(
+        &mut self,
+        changes: I,
+        f: F,
+    ) -> Result<HashMap<Address, AccountInfo>, Db::Error>
+    where
+        I: IntoIterator<Item = Address>,
+        F: Fn(&mut AccountInfo),
+        Db: DatabaseCommit,
+    {
+        let mut state = HashMap::default();
+        let mut old = HashMap::default();
+        for account in changes.into_iter() {
+            let mut acct = self.account(account)?;
+            old.insert(account, acct.info.clone());
+            f(&mut acct.info);
+            acct.mark_touch();
+            state.insert(account, acct);
+        }
+        self.db_mut_ext().commit(state);
+        Ok(old)
+    }
+
+    /// Modify an account with a closure and commit the modified account.
+    fn modify_account<F>(&mut self, address: Address, f: F) -> Result<AccountInfo, Db::Error>
+    where
+        F: FnOnce(&mut AccountInfo),
+        Db: DatabaseCommit,
+    {
+        let mut acct = self.account(address)?;
+        let old = acct.info.clone();
+        f(&mut acct.info);
+        acct.mark_touch();
+        let changes = [(address, acct)].into_iter().collect();
+        self.db_mut_ext().commit(changes);
+        Ok(old)
+    }
+
+    /// Set the storage value of an account, returning the previous value.
+    fn set_storage(&mut self, address: Address, index: U256, value: U256) -> Result<U256, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        let mut acct = self.account(address)?;
+        let old = self.storage(address, index)?;
+
+        let change = EvmStorageSlot::new_changed(old, value);
+        acct.storage.insert(index, change);
+        acct.mark_touch();
+
+        let changes = [(address, acct)].into_iter().collect();
+        self.db_mut_ext().commit(changes);
+        Ok(old)
+    }
+
+    /// Set the bytecode for an account, returns the previous bytecode.
+    fn set_bytecode(
+        &mut self,
+        address: Address,
+        bytecode: Bytecode,
+    ) -> Result<Option<Bytecode>, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        let mut acct = self.account(address)?;
+        let old = self.db_mut_ext().code_by_hash(acct.info.code_hash).map(|old| {
+            if old.is_empty() {
+                None
+            } else {
+                Some(old)
+            }
+        })?;
+        acct.info.code_hash = bytecode.hash_slow();
+        acct.info.code = Some(bytecode);
+        acct.mark_touch();
+
+        let changes = [(address, acct)].into_iter().collect();
+        self.db_mut_ext().commit(changes);
+        Ok(old)
+    }
+
+    /// Increase the balance of an account, returning the previous balance.
+    fn increase_balance(&mut self, address: Address, amount: U256) -> Result<U256, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        self.modify_account(address, |info| info.balance = info.balance.saturating_add(amount))
+            .map(|info| info.balance)
+    }
+
+    /// Decrease the balance of an account, returning the previous balance.
+    fn decrease_balance(&mut self, address: Address, amount: U256) -> Result<U256, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        self.modify_account(address, |info| info.balance = info.balance.saturating_sub(amount))
+            .map(|info| info.balance)
+    }
+
+    /// Set the balance of an account, returning the previous balance.
+    fn set_balance(&mut self, address: Address, amount: U256) -> Result<U256, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        self.modify_account(address, |info| info.balance = amount).map(|info| info.balance)
+    }
+
+    /// Set the nonce of an account, returning the previous nonce.
+    fn set_nonce(&mut self, address: Address, nonce: u64) -> Result<u64, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        self.modify_account(address, |info| info.nonce = nonce).map(|info| info.nonce)
+    }
+
+    /// Increment the nonce of an account, returning the new nonce.
+    fn increment_nonce(&mut self, address: Address) -> Result<u64, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        self.modify_account(address, |info| info.nonce = info.nonce.saturating_add(1))
+            .map(|info| info.nonce)
+    }
+
+    /// Decrement the nonce of an account, returning the new nonce.
+    fn decrement_nonce(&mut self, address: Address) -> Result<u64, Db::Error>
+    where
+        Db: DatabaseCommit,
+    {
+        self.modify_account(address, |info| info.nonce = info.nonce.saturating_sub(1))
+            .map(|info| info.nonce)
+    }
+}
+
+impl<Ext, Db: Database> EvmExtUnchecked<Db> for revm::Evm<'_, Ext, Db> {
+    fn db_mut_ext(&mut self) -> &mut Db {
+        self.db_mut()
+    }
+}

src/lib.rs

@@ -402,6 +402,9 @@
 mod evm;
 pub use evm::Trevm;
 
+mod ext;
+pub use ext::EvmExtUnchecked;
+
 mod fill;
 pub use fill::{Block, Cfg, NoopBlock, NoopCfg, Tx};