feat: Clean up StateDB methods, bump revm

Cargo.toml

@@ -68,8 +68,9 @@ serde_json = "1"
 test-case = "3"
 
 [patch.crates-io]
-# revm = { git = "https://github.com/bluealloy/revm", rev = "11b16259" }
-# op-revm = { git = "https://github.com/bluealloy/revm", rev = "11b16259" }
+# TODO: Remove once new revm version is published
+revm = { git = "https://github.com/0xforerunner/revm", rev = "195c8e9" }
+op-revm = { git = "https://github.com/0xforerunner/revm", rev = "195c8e9" }
 
 #alloy-hardforks = { git = "https://github.com/alloy-rs/hardforks", rev = "0fd230f5aa24c4f6a8c593918b7449df0c20f60a" }
 #alloy-op-hardforks = { git = "https://github.com/alloy-rs/hardforks", rev = "0fd230f5aa24c4f6a8c593918b7449df0c20f60a" }

crates/evm/src/block/mod.rs

@@ -5,9 +5,8 @@ use alloc::{boxed::Box, vec::Vec};
 use alloy_eips::eip7685::Requests;
 use revm::{
     context::result::{ExecutionResult, ResultAndState},
-    database::State,
     inspector::NoOpInspector,
-    Inspector,
+    DatabaseCommit, Inspector,
 };
 
 mod error;
@@ -328,22 +327,22 @@ pub trait BlockExecutor {
 pub trait BlockExecutorFor<'a, F: BlockExecutorFactory + ?Sized, DB, I = NoOpInspector>
 where
     Self: BlockExecutor<
-        Evm = <F::EvmFactory as EvmFactory>::Evm<&'a mut State<DB>, I>,
+        Evm = <F::EvmFactory as EvmFactory>::Evm<DB, I>,
         Transaction = F::Transaction,
         Receipt = F::Receipt,
     >,
-    DB: Database + 'a,
-    I: Inspector<<F::EvmFactory as EvmFactory>::Context<&'a mut State<DB>>> + 'a,
+    DB: StateDB + Database + 'a,
+    I: Inspector<<F::EvmFactory as EvmFactory>::Context<DB>> + 'a,
 {
 }
 
 impl<'a, F, DB, I, T> BlockExecutorFor<'a, F, DB, I> for T
 where
     F: BlockExecutorFactory,
-    DB: Database + 'a,
-    I: Inspector<<F::EvmFactory as EvmFactory>::Context<&'a mut State<DB>>> + 'a,
+    DB: StateDB + Database + 'a,
+    I: Inspector<<F::EvmFactory as EvmFactory>::Context<DB>> + 'a,
     T: BlockExecutor<
-        Evm = <F::EvmFactory as EvmFactory>::Evm<&'a mut State<DB>, I>,
+        Evm = <F::EvmFactory as EvmFactory>::Evm<DB, I>,
         Transaction = F::Transaction,
         Receipt = F::Receipt,
     >,
@@ -467,10 +466,10 @@ pub trait BlockExecutorFactory: 'static {
     /// ```
     fn create_executor<'a, DB, I>(
         &'a self,
-        evm: <Self::EvmFactory as EvmFactory>::Evm<&'a mut State<DB>, I>,
+        evm: <Self::EvmFactory as EvmFactory>::Evm<DB, I>,
         ctx: Self::ExecutionCtx<'a>,
     ) -> impl BlockExecutorFor<'a, Self, DB, I>
     where
-        DB: Database + 'a,
-        I: Inspector<<Self::EvmFactory as EvmFactory>::Context<&'a mut State<DB>>> + 'a;
+        DB: StateDB + DatabaseCommit + Database + 'a,
+        I: Inspector<<Self::EvmFactory as EvmFactory>::Context<DB>> + 'a;
 }

crates/evm/src/block/state.rs

@@ -1,7 +1,6 @@
 //! State database abstraction.
 
-use alloy_primitives::Address;
-use revm::database::State;
+use revm::database::{states::bundle_state::BundleRetention, BundleState, State};
 
 /// A type which has the state of the blockchain.
 ///
@@ -10,18 +9,29 @@ pub trait StateDB: revm::Database {
     /// State clear EIP-161 is enabled in Spurious Dragon hardfork.
     fn set_state_clear_flag(&mut self, has_state_clear: bool);
 
-    /// Iterates over received balances and increment all account balances.
+    /// Gets a reference to the internal [`BundleState`]
+    fn bundle_state(&self) -> &BundleState;
+
+    /// Gets a mutable reference to the internal [`BundleState`]
+    fn bundle_state_mut(&mut self) -> &mut BundleState;
+
+    /// This will not apply any pending [`revm::database::TransitionState`].
     ///
-    /// **Note**: If account is not found inside cache state it will be loaded from database.
+    /// It is recommended to call [`StateDB::merge_transitions`] before taking the bundle.
     ///
-    /// Update will create transitions for all accounts that are updated.
+    /// If the `State` has been built with the
+    /// [`revm::database::StateBuilder::with_bundle_prestate`] option, the pre-state will be
+    /// taken along with any changes made by [`StateDB::merge_transitions`].
+    fn take_bundle(&mut self) -> BundleState {
+        core::mem::take(self.bundle_state_mut())
+    }
+
+    /// Take all transitions and merge them inside [`BundleState`].
     ///
-    /// If using this to implement withdrawals, zero balances must be filtered out before calling
-    /// this function.
-    fn increment_balances(
-        &mut self,
-        balances: impl IntoIterator<Item = (Address, u128)>,
-    ) -> Result<(), Self::Error>;
+    /// This action will create final post state and all reverts so that
+    /// we at any time revert state of bundle to the state before transition
+    /// is applied.
+    fn merge_transitions(&mut self, retention: BundleRetention);
 }
 
 /// auto_impl unable to reconcile return associated type from supertrait
@@ -30,23 +40,33 @@ impl<T: StateDB> StateDB for &mut T {
         StateDB::set_state_clear_flag(*self, has_state_clear);
     }
 
-    fn increment_balances(
-        &mut self,
-        balances: impl IntoIterator<Item = (Address, u128)>,
-    ) -> Result<(), Self::Error> {
-        StateDB::increment_balances(*self, balances)
+    fn bundle_state(&self) -> &BundleState {
+        StateDB::bundle_state(*self)
+    }
+
+    fn bundle_state_mut(&mut self) -> &mut BundleState {
+        StateDB::bundle_state_mut(*self)
+    }
+
+    fn merge_transitions(&mut self, retention: BundleRetention) {
+        StateDB::merge_transitions(*self, retention);
     }
 }
 
 impl<DB: revm::Database> StateDB for State<DB> {
     fn set_state_clear_flag(&mut self, has_state_clear: bool) {
-        self.cache.set_state_clear_flag(has_state_clear);
+        Self::set_state_clear_flag(self, has_state_clear);
+    }
+
+    fn bundle_state(&self) -> &BundleState {
+        &self.bundle_state
+    }
+
+    fn bundle_state_mut(&mut self) -> &mut BundleState {
+        &mut self.bundle_state
     }
 
-    fn increment_balances(
-        &mut self,
-        balances: impl IntoIterator<Item = (Address, u128)>,
-    ) -> Result<(), Self::Error> {
-        Self::increment_balances(self, balances)
+    fn merge_transitions(&mut self, retention: BundleRetention) {
+        Self::merge_transitions(self, retention);
     }
 }

crates/evm/src/eth/block.rs

@@ -11,7 +11,7 @@ use crate::{
         state_changes::{balance_increment_state, post_block_balance_increments},
         BlockExecutionError, BlockExecutionResult, BlockExecutor, BlockExecutorFactory,
         BlockExecutorFor, BlockValidationError, ExecutableTx, OnStateHook,
-        StateChangePostBlockSource, StateChangeSource, SystemCaller,
+        StateChangePostBlockSource, StateChangeSource, StateDB, SystemCaller,
     },
     Database, Evm, EvmFactory, FromRecoveredTx, FromTxWithEncoded,
 };
@@ -21,8 +21,8 @@ use alloy_eips::{eip4895::Withdrawals, eip7685::Requests, Encodable2718};
 use alloy_hardforks::EthereumHardfork;
 use alloy_primitives::{Log, B256};
 use revm::{
-    context::Block, context_interface::result::ResultAndState, database::State, DatabaseCommit,
-    Inspector,
+    context::Block, context_interface::result::ResultAndState, database::DatabaseCommitExt as _,
+    DatabaseCommit, Inspector,
 };
 
 /// Context for Ethereum block execution.
@@ -83,11 +83,10 @@ where
     }
 }
 
-impl<'db, DB, E, Spec, R> BlockExecutor for EthBlockExecutor<'_, E, Spec, R>
+impl<E, Spec, R> BlockExecutor for EthBlockExecutor<'_, E, Spec, R>
 where
-    DB: Database + 'db,
     E: Evm<
-        DB = &'db mut State<DB>,
+        DB: StateDB + DatabaseCommit,
         Tx: FromRecoveredTx<R::Transaction> + FromTxWithEncoded<R::Transaction>,
     >,
     Spec: EthExecutorSpec,
@@ -314,12 +313,12 @@ where
 
     fn create_executor<'a, DB, I>(
         &'a self,
-        evm: EvmF::Evm<&'a mut State<DB>, I>,
+        evm: EvmF::Evm<DB, I>,
         ctx: Self::ExecutionCtx<'a>,
     ) -> impl BlockExecutorFor<'a, Self, DB, I>
     where
-        DB: Database + 'a,
-        I: Inspector<EvmF::Context<&'a mut State<DB>>> + 'a,
+        DB: StateDB + DatabaseCommit + Database + 'a,
+        I: Inspector<EvmF::Context<DB>> + 'a,
     {
         EthBlockExecutor::new(evm, ctx, &self.spec, &self.receipt_builder)
     }

crates/evm/src/traits.rs

@@ -5,7 +5,10 @@ use alloc::boxed::Box;
 use alloy_primitives::{Address, Log, B256, U256};
 use core::{error::Error, fmt, fmt::Debug};
 use revm::{
-    context::{journaled_state::TransferError, Block, DBErrorMarker, JournalTr},
+    context::{
+        journaled_state::{account::JournaledAccountTr as _, TransferError},
+        Block, DBErrorMarker, JournalTr,
+    },
     interpreter::{SStoreResult, StateLoad},
     primitives::{StorageKey, StorageValue},
     state::{Account, AccountInfo, Bytecode},

crates/op-evm/src/block/mod.rs

@@ -26,7 +26,7 @@ pub use receipt_builder::OpAlloyReceiptBuilder;
 use receipt_builder::OpReceiptBuilder;
 use revm::{
     context::{result::ResultAndState, Block},
-    database::State,
+    database::DatabaseCommitExt as _,
     Database as _, DatabaseCommit, Inspector,
 };
 
@@ -416,12 +416,12 @@ where
 
     fn create_executor<'a, DB, I>(
         &'a self,
-        evm: EvmF::Evm<&'a mut State<DB>, I>,
+        evm: EvmF::Evm<DB, I>,
         ctx: Self::ExecutionCtx<'a>,
     ) -> impl BlockExecutorFor<'a, Self, DB, I>
     where
-        DB: Database + 'a,
-        I: Inspector<EvmF::Context<&'a mut State<DB>>> + 'a,
+        DB: StateDB + DatabaseCommit + Database + 'a,
+        I: Inspector<EvmF::Context<DB>> + 'a,
     {
         OpBlockExecutor::new(evm, ctx, &self.spec, &self.receipt_builder)
     }
@@ -446,7 +446,7 @@ mod tests {
     };
     use revm::{
         context::BlockEnv,
-        database::{CacheDB, EmptyDB, InMemoryDB},
+        database::{CacheDB, EmptyDB, InMemoryDB, State},
         inspector::NoOpInspector,
         primitives::HashMap,
         state::AccountInfo,