feat: cachingdb traits (#102)

* feat: cachingdb traits

* chore: bump version

* lint: clippy

Author: prestwich

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "trevm"
-version = "0.20.7"
+version = "0.20.8"
 rust-version = "1.83.0"
 edition = "2021"
 authors = ["init4"]

src/db/cow/mod.rs

@@ -1,3 +1,4 @@
+use crate::db::CachingDb;
 use alloy::{
     consensus::constants::KECCAK_EMPTY,
     primitives::{Address, B256, U256},
@@ -10,6 +11,8 @@ use revm::{
     Database, DatabaseCommit, DatabaseRef,
 };
 
+use super::TryCachingDb;
+
 /// A version of [`CacheDB`] that caches only on write, not on read.
 ///
 /// This saves memory when wrapping some other caching database, like [`State`]
@@ -49,11 +52,21 @@ impl<Db> CacheOnWrite<Db> {
         &self.inner
     }
 
+    /// Get a mutable reference to the inner database.
+    pub const fn inner_mut(&mut self) -> &mut Db {
+        &mut self.inner
+    }
+
     /// Get a refernce to the [`Cache`].
     pub const fn cache(&self) -> &Cache {
         &self.cache
     }
 
+    /// Get a mutable reference to the [`Cache`].
+    pub const fn cache_mut(&mut self) -> &mut Cache {
+        &mut self.cache
+    }
+
     /// Deconstruct the `CacheOnWrite` into its parts.
     pub fn into_parts(self) -> (Db, Cache) {
         (self.inner, self.cache)
@@ -91,29 +104,67 @@ impl<Db> CacheOnWrite<Db> {
     }
 }
 
-impl<Db> CacheOnWrite<CacheOnWrite<Db>> {
-    /// Discard the outer cache, returning the inner.
-    pub fn discard_outer(self) -> CacheOnWrite<Db> {
-        self.inner
+impl<Db> CachingDb for CacheOnWrite<Db> {
+    fn cache(&self) -> &Cache {
+        &self.cache
+    }
+
+    fn cache_mut(&mut self) -> &mut Cache {
+        &mut self.cache
     }
 
+    fn into_cache(self) -> Cache {
+        self.cache
+    }
+}
+
+impl<Db> CacheOnWrite<Db>
+where
+    Db: CachingDb,
+{
     /// Flattens a nested cache by applying the outer cache to the inner cache.
     ///
     /// The behavior is as follows:
     /// - Accounts are overridden with outer accounts
     /// - Contracts are overridden with outer contracts
     /// - Block hashes are overridden with outer block hashes
-    pub fn flatten(self) -> CacheOnWrite<Db> {
+    pub fn flatten(self) -> Db {
         let Self { cache: Cache { accounts, contracts, logs, block_hashes }, mut inner } = self;
 
-        inner.cache.accounts.extend(accounts);
-        inner.cache.contracts.extend(contracts);
-        inner.cache.logs.extend(logs);
-        inner.cache.block_hashes.extend(block_hashes);
+        let inner_cache = inner.cache_mut();
+
+        inner_cache.accounts.extend(accounts);
+        inner_cache.contracts.extend(contracts);
+        inner_cache.logs.extend(logs);
+        inner_cache.block_hashes.extend(block_hashes);
         inner
     }
 }
 
+impl<Db> CacheOnWrite<Db>
+where
+    Db: TryCachingDb,
+{
+    /// Attempts to flatten a nested cache by applying the outer cache to the
+    /// inner cache. This is a fallible version of [`CacheOnWrite::flatten`].
+    ///
+    /// The behavior is as follows:
+    /// - Accounts are overridden with outer accounts
+    /// - Contracts are overridden with outer contracts
+    /// - Block hashes are overridden with outer block hashes
+    pub fn try_flatten(self) -> Result<Db, Db::Error> {
+        let Self { cache: Cache { accounts, contracts, logs, block_hashes }, mut inner } = self;
+
+        let inner_cache = inner.try_cache_mut()?;
+
+        inner_cache.accounts.extend(accounts);
+        inner_cache.contracts.extend(contracts);
+        inner_cache.logs.extend(logs);
+        inner_cache.block_hashes.extend(block_hashes);
+        Ok(inner)
+    }
+}
+
 impl<Db: DatabaseRef> Database for CacheOnWrite<Db> {
     type Error = Db::Error;
 

src/db/mod.rs

@@ -4,7 +4,7 @@ pub mod sync;
 
 /// Database abstraction traits.
 mod traits;
-pub use traits::{ArcUpgradeError, StateAcc, TryDatabaseCommit, TryStateAcc};
+pub use traits::{ArcUpgradeError, CachingDb, StateAcc, TryCachingDb, TryStateAcc};
 
 /// Cache-on-write database. A memory cache that caches only on write, not on
 /// read. Intended to wrap some other caching database.

src/db/traits.rs

@@ -1,8 +1,7 @@
 use revm::{
-    database::{states::bundle_state::BundleRetention, BundleState, State},
-    primitives::{Address, B256},
-    state::Account,
-    Database, DatabaseCommit,
+    database::{states::bundle_state::BundleRetention, BundleState, Cache, CacheDB, State},
+    primitives::B256,
+    Database,
 };
 use std::{collections::BTreeMap, convert::Infallible, sync::Arc};
 
@@ -145,29 +144,63 @@ where
     }
 }
 
-/// A fallible version of [`DatabaseCommit`].
-pub trait TryDatabaseCommit {
-    /// Error type to be thrown when committing changes fails.
+/// Trait for Databases that have a [`Cache`].
+pub trait CachingDb {
+    /// Get the cache.
+    fn cache(&self) -> &Cache;
+
+    /// Get the cache mutably.
+    fn cache_mut(&mut self) -> &mut Cache;
+
+    /// Deconstruct into the cache
+    fn into_cache(self) -> Cache;
+}
+
+impl<Db> CachingDb for CacheDB<Db> {
+    fn cache(&self) -> &Cache {
+        &self.cache
+    }
+
+    fn cache_mut(&mut self) -> &mut Cache {
+        &mut self.cache
+    }
+
+    fn into_cache(self) -> Cache {
+        self.cache
+    }
+}
+
+/// Trait for Databases that have a [`Cache`] and can fail to mutably access
+/// it. E.g. `Arc<CacheDB<Db>>`
+pub trait TryCachingDb {
+    /// Error type to be thrown when cache access 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>;
+    /// Attempt to get the cache.
+    fn cache(&self) -> &Cache;
+
+    /// Attempt to get the cache mutably.
+    fn try_cache_mut(&mut self) -> Result<&mut Cache, Self::Error>;
+
+    /// Attempt to deconstruct into the cache
+    fn try_into_cache(self) -> Result<Cache, Self::Error>;
 }
 
-impl<Db> TryDatabaseCommit for Arc<Db>
+impl<Db> TryCachingDb for Arc<Db>
 where
-    Db: DatabaseCommit,
+    Db: CachingDb,
 {
     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(())
+    fn cache(&self) -> &Cache {
+        self.as_ref().cache()
+    }
+
+    fn try_cache_mut(&mut self) -> Result<&mut Cache, Self::Error> {
+        Self::get_mut(self).ok_or(ArcUpgradeError::NotUnique).map(|db| db.cache_mut())
+    }
+
+    fn try_into_cache(self) -> Result<Cache, Self::Error> {
+        Self::into_inner(self).ok_or(ArcUpgradeError::NotUnique).map(|db| db.into_cache())
     }
 }

src/evm.rs

@@ -1,5 +1,5 @@
 use crate::{
-    db::{StateAcc, TryDatabaseCommit, TryStateAcc},
+    db::{StateAcc, TryStateAcc},
     driver::DriveBlockResult,
     helpers::{Ctx, Evm},
     Block, BlockDriver, BundleDriver, Cfg, ChainDriver, DriveBundleResult, DriveChainResult,
@@ -16,7 +16,7 @@ use revm::{
         result::{EVMError, ExecutionResult, InvalidTransaction, ResultAndState},
         Block as _, BlockEnv, Cfg as _, ContextSetters, ContextTr, Transaction as _, TxEnv,
     },
-    database::{states::bundle_state::BundleRetention, BundleState},
+    database::{states::bundle_state::BundleRetention, BundleState, TryDatabaseCommit},
     inspector::NoOpInspector,
     interpreter::gas::calculate_initial_tx_gas_for_tx,
     primitives::{hardfork::SpecId, TxKind},