feat: traits for things that behave like State

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "trevm"
-version = "0.19.8"
+version = "0.19.9"
 rust-version = "1.83.0"
 edition = "2021"
 authors = ["init4"]
@@ -35,6 +35,7 @@ zenith-types = { version = "0.15" }
 
 dashmap = { version = "6.1.0", optional = true }
 tracing = { version = "0.1.41", optional = true}
+thiserror = "2.0.11"
 
 [dev-dependencies]
 revm = { version = "19.5.0", features = [

src/db/mod.rs

@@ -1,78 +1,7 @@
-mod builder;
+/// Concurrent version of [`revm::db::State`]
+#[cfg(feature = "concurrent-db")]
+pub mod sync;
 
-use alloy::rpc::types::BlockOverrides;
-pub use builder::ConcurrentStateBuilder;
-
-mod cache_state;
-pub use cache_state::ConcurrentCacheState;
-
-mod sync_state;
-pub use sync_state::{ConcurrentState, ConcurrentStateInfo};
-
-use crate::{Block, EvmNeedsBlock, EvmNeedsTx, Trevm};
-use revm::{
-    db::{states::bundle_state::BundleRetention, BundleState},
-    DatabaseRef,
-};
-
-impl<Ext, Db: DatabaseRef + Sync, TrevmState> Trevm<'_, Ext, ConcurrentState<Db>, TrevmState> {
-    /// Set the [EIP-161] state clear flag, activated in the Spurious Dragon
-    /// hardfork.
-    ///
-    /// This function changes the behavior of the inner [`ConcurrentState`].
-    pub fn set_state_clear_flag(&mut self, flag: bool) {
-        self.inner.db_mut().set_state_clear_flag(flag)
-    }
-}
-
-impl<Ext, Db: DatabaseRef + Sync> EvmNeedsBlock<'_, Ext, ConcurrentState<Db>> {
-    /// Finish execution and return the outputs.
-    ///
-    /// If the State has not been built with
-    /// [revm::StateBuilder::with_bundle_update] then the returned
-    /// [`BundleState`] will be meaningless.
-    ///
-    /// See [`ConcurrentState::merge_transitions`] and
-    /// [`ConcurrentState::take_bundle`].
-    pub fn finish(self) -> BundleState {
-        let Self { inner: mut evm, .. } = self;
-        evm.db_mut().merge_transitions(BundleRetention::Reverts);
-        let bundle = evm.db_mut().take_bundle();
-
-        bundle
-    }
-}
-
-impl<Ext, Db: DatabaseRef + Sync> EvmNeedsTx<'_, Ext, ConcurrentState<Db>> {
-    /// Apply block overrides to the current block.
-    ///
-    /// Note that this is NOT reversible. The overrides are applied directly to
-    /// the underlying state and these changes cannot be removed. If it is
-    /// important that you have access to the pre-change state, you should wrap
-    /// the existing DB in a new [`ConcurrentState`] and apply the overrides to
-    /// that.
-    pub fn apply_block_overrides(mut self, overrides: &BlockOverrides) -> Self {
-        overrides.fill_block(&mut self.inner);
-
-        if let Some(hashes) = &overrides.block_hash {
-            self.inner.db_mut().info.block_hashes.write().unwrap().extend(hashes)
-        }
-
-        self
-    }
-
-    /// Apply block overrides to the current block, if they are provided.
-    ///
-    /// Note that this is NOT reversible. The overrides are applied directly to
-    /// the underlying state and these changes cannot be removed. If it is
-    /// important that you have access to the pre-change state, you should wrap
-    /// the existing DB in a new [`ConcurrentState`] and apply the overrides to
-    /// that.
-    pub fn maybe_apply_block_overrides(self, overrides: Option<&BlockOverrides>) -> Self {
-        if let Some(overrides) = overrides {
-            self.apply_block_overrides(overrides)
-        } else {
-            self
-        }
-    }
-}
+/// Database abstraction traits.
+mod traits;
+pub use traits::{ArcUpgradeError, StateAcc, TryDatabaseCommit, TryStateAcc};

src/db/builder.rs → src/db/sync/builder.rs

@@ -1,4 +1,4 @@
-use crate::db::ConcurrentState;
+use crate::db::sync::ConcurrentState;
 use revm::{
     db::{
         states::{BundleState, TransitionState},

src/db/sync/mod.rs

@@ -0,0 +1,35 @@
+mod builder;
+
+pub use builder::ConcurrentStateBuilder;
+
+mod cache;
+pub use cache::ConcurrentCacheState;
+
+mod state;
+pub use state::{ConcurrentState, ConcurrentStateInfo};
+
+use crate::db::StateAcc;
+use revm::{
+    db::{states::bundle_state::BundleRetention, BundleState},
+    primitives::B256,
+    DatabaseRef,
+};
+use std::collections::BTreeMap;
+
+impl<Db: DatabaseRef + Sync> StateAcc for ConcurrentState<Db> {
+    fn set_state_clear_flag(&mut self, flag: bool) {
+        Self::set_state_clear_flag(self, flag)
+    }
+
+    fn merge_transitions(&mut self, retention: BundleRetention) {
+        Self::merge_transitions(self, retention)
+    }
+
+    fn take_bundle(&mut self) -> BundleState {
+        Self::take_bundle(self)
+    }
+
+    fn set_block_hashes(&mut self, block_hashes: &BTreeMap<u64, B256>) {
+        self.info.block_hashes.write().unwrap().extend(block_hashes)
+    }
+}

src/db/sync_state.rs → src/db/sync/state.rs

@@ -1,4 +1,4 @@
-use crate::db::ConcurrentCacheState;
+use crate::db::sync::ConcurrentCacheState;
 use alloy::primitives::{Address, B256, U256};
 use dashmap::mapref::one::RefMut;
 use revm::{

src/db/traits.rs

@@ -0,0 +1,172 @@
+use revm::{
+    db::{states::bundle_state::BundleRetention, BundleState, State},
+    primitives::{Account, Address, B256},
+    Database, DatabaseCommit,
+};
+use std::{collections::BTreeMap, convert::Infallible, sync::Arc};
+
+/// Abstraction trait covering types that accumulate state changes into a
+/// [`BundleState`]. The prime example of this is [`State`]. These types are
+/// use to accumulate state changes during the execution of a sequence of
+/// transactions, and then provide access to the net changes in the form of a
+/// [`BundleState`].
+pub trait StateAcc {
+    /// Set the state clear flag. See [`State::set_state_clear_flag`].
+    fn set_state_clear_flag(&mut self, flag: bool);
+
+    /// Merge transitions into the bundle. See [`State::merge_transitions`].
+    fn merge_transitions(&mut self, retention: BundleRetention);
+
+    /// Take the bundle. See [`State::take_bundle`].
+    fn take_bundle(&mut self) -> BundleState;
+
+    /// Set the block hashes, overriding any existing values, and inserting any
+    /// absent values.
+    fn set_block_hashes(&mut self, block_hashes: &BTreeMap<u64, B256>);
+}
+
+impl<Db: Database> StateAcc for State<Db> {
+    fn set_state_clear_flag(&mut self, flag: bool) {
+        Self::set_state_clear_flag(self, flag)
+    }
+
+    fn merge_transitions(&mut self, retention: BundleRetention) {
+        Self::merge_transitions(self, retention)
+    }
+
+    fn take_bundle(&mut self) -> BundleState {
+        Self::take_bundle(self)
+    }
+
+    fn set_block_hashes(&mut self, block_hashes: &BTreeMap<u64, B256>) {
+        self.block_hashes.extend(block_hashes)
+    }
+}
+
+/// Fallible version of [`StateAcc`].
+///
+/// Abstraction trait covering types that accumulate state changes into a
+/// [`BundleState`]. The prime example of this is [`State`]. These types are
+/// use to accumulate state changes during the execution of a sequence of
+/// transactions, and then provide access to the net changes in the form of a
+/// [`BundleState`].
+///
+/// The primary motivator for this trait is to allow for the implementation of
+/// [`StateAcc`] for [`Arc`]-wrapped DBs, which may fail to mutate if the
+/// reference is not unique.
+pub trait TryStateAcc: Sync {
+    /// Error type to be thrown when state accumulation fails.
+    type Error: core::error::Error;
+
+    /// Attempt to set the state clear flag. See [`State::set_state_clear_flag`].
+    fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), Self::Error>;
+
+    /// Attempt to merge transitions into the bundle. See
+    /// [`State::merge_transitions`].
+    fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), Self::Error>;
+
+    /// Attempt to take the bundle. See [`State::take_bundle`].
+    fn try_take_bundle(&mut self) -> Result<BundleState, Self::Error>;
+
+    /// Attempt to set the block hashes, overriding any existing values, and
+    /// inserting any absent values.
+    fn try_set_block_hashes(
+        &mut self,
+        block_hashes: &BTreeMap<u64, B256>,
+    ) -> Result<(), Self::Error>;
+}
+
+impl<Db> TryStateAcc for Db
+where
+    Db: StateAcc + Sync,
+{
+    type Error = Infallible;
+
+    fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), Infallible> {
+        self.set_state_clear_flag(flag);
+        Ok(())
+    }
+
+    fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), Infallible> {
+        self.merge_transitions(retention);
+        Ok(())
+    }
+
+    fn try_take_bundle(&mut self) -> Result<BundleState, Infallible> {
+        Ok(self.take_bundle())
+    }
+
+    fn try_set_block_hashes(
+        &mut self,
+        block_hashes: &BTreeMap<u64, B256>,
+    ) -> Result<(), Infallible> {
+        self.set_block_hashes(block_hashes);
+        Ok(())
+    }
+}
+
+/// Error type for implementation of [`TryStateAcc`] for [`Arc`]-wrapped
+/// DBs.
+#[derive(thiserror::Error, Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ArcUpgradeError {
+    /// Arc reference is not unique. Ensure that all other references are
+    /// dropped before attempting to mutate the state.
+    #[error("Arc reference is not unique, cannot mutate")]
+    NotUnique,
+}
+
+impl<Db> TryStateAcc for Arc<Db>
+where
+    Db: StateAcc + Sync + Send,
+{
+    type Error = ArcUpgradeError;
+
+    fn try_set_state_clear_flag(&mut self, flag: bool) -> Result<(), ArcUpgradeError> {
+        Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.set_state_clear_flag(flag);
+        Ok(())
+    }
+
+    fn try_merge_transitions(&mut self, retention: BundleRetention) -> Result<(), ArcUpgradeError> {
+        Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.merge_transitions(retention);
+        Ok(())
+    }
+
+    fn try_take_bundle(&mut self) -> Result<BundleState, ArcUpgradeError> {
+        Ok(Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.take_bundle())
+    }
+
+    fn try_set_block_hashes(
+        &mut self,
+        block_hashes: &BTreeMap<u64, B256>,
+    ) -> Result<(), ArcUpgradeError> {
+        Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.set_block_hashes(block_hashes);
+        Ok(())
+    }
+}
+
+/// A fallible version of [`DatabaseCommit`].
+pub trait TryDatabaseCommit {
+    /// Error type to be thrown when committing changes fails.
+    type Error: core::error::Error;
+
+    /// Attempt to commit changes to the database.
+    fn try_commit(
+        &mut self,
+        changes: revm::primitives::HashMap<Address, Account>,
+    ) -> Result<(), Self::Error>;
+}
+
+impl<Db> TryDatabaseCommit for Arc<Db>
+where
+    Db: DatabaseCommit,
+{
+    type Error = ArcUpgradeError;
+
+    fn try_commit(
+        &mut self,
+        changes: revm::primitives::HashMap<Address, Account>,
+    ) -> Result<(), Self::Error> {
+        Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique)?.commit(changes);
+        Ok(())
+    }
+}