From eaac2aa2cfb6ea5a61eea2d9a36a8b66b7854ce4 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 1 Nov 2024 15:17:31 +0100
Subject: [PATCH] chore: simplify cached db usage (#12242)

---
 crates/ethereum/payload/src/lib.rs     |  5 ++---
 crates/optimism/payload/src/builder.rs |  5 ++---
 crates/revm/src/cached.rs              | 24 +++++++++++++++++++-----
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/crates/ethereum/payload/src/lib.rs b/crates/ethereum/payload/src/lib.rs
index 8e92d0aa8..ed33292ef 100644
--- a/crates/ethereum/payload/src/lib.rs
+++ b/crates/ethereum/payload/src/lib.rs
@@ -161,7 +161,7 @@ where
     let state_provider = client.state_by_block_hash(config.parent_header.hash())?;
     let state = StateProviderDatabase::new(state_provider);
     let mut db =
-        State::builder().with_database_ref(cached_reads.as_db(state)).with_bundle_update().build();
+        State::builder().with_database(cached_reads.as_db_mut(state)).with_bundle_update().build();
     let PayloadConfig { parent_header, extra_data, attributes } = config;
 
     debug!(target: "payload_builder", id=%attributes.id, parent_header = ?parent_header.hash(), parent_number = parent_header.number, "building new payload");
@@ -372,8 +372,7 @@ where
     // calculate the state root
     let hashed_state = HashedPostState::from_bundle_state(&execution_outcome.state().state);
     let (state_root, trie_output) = {
-        let state_provider = db.database.0.inner.borrow_mut();
-        state_provider.db.state_root_with_updates(hashed_state.clone()).inspect_err(|err| {
+        db.database.inner().state_root_with_updates(hashed_state.clone()).inspect_err(|err| {
             warn!(target: "payload_builder",
                 parent_hash=%parent_header.hash(),
                 %err,
diff --git a/crates/optimism/payload/src/builder.rs b/crates/optimism/payload/src/builder.rs
index 09a443d0f..cae2d34bd 100644
--- a/crates/optimism/payload/src/builder.rs
+++ b/crates/optimism/payload/src/builder.rs
@@ -170,7 +170,7 @@ where
     let state_provider = client.state_by_block_hash(config.parent_header.hash())?;
     let state = StateProviderDatabase::new(state_provider);
     let mut db =
-        State::builder().with_database_ref(cached_reads.as_db(state)).with_bundle_update().build();
+        State::builder().with_database(cached_reads.as_db_mut(state)).with_bundle_update().build();
     let PayloadConfig { parent_header, attributes, mut extra_data } = config;
 
     debug!(target: "payload_builder", id=%attributes.payload_attributes.payload_id(), parent_header = ?parent_header.hash(), parent_number = parent_header.number, "building new payload");
@@ -445,8 +445,7 @@ where
     // calculate the state root
     let hashed_state = HashedPostState::from_bundle_state(&execution_outcome.state().state);
     let (state_root, trie_output) = {
-        let state_provider = db.database.0.inner.borrow_mut();
-        state_provider.db.state_root_with_updates(hashed_state.clone()).inspect_err(|err| {
+        db.database.inner().state_root_with_updates(hashed_state.clone()).inspect_err(|err| {
             warn!(target: "payload_builder",
             parent_header=%parent_header.hash(),
                 %err,
diff --git a/crates/revm/src/cached.rs b/crates/revm/src/cached.rs
index b0eac39c4..2152ca5bd 100644
--- a/crates/revm/src/cached.rs
+++ b/crates/revm/src/cached.rs
@@ -22,10 +22,10 @@ use reth_primitives::revm_primitives::{
 ///
 /// fn build_payload<DB: DatabaseRef>(db: DB) {
 ///     let mut cached_reads = CachedReads::default();
-///     let db_ref = cached_reads.as_db(db);
-///     // this is `Database` and can be used to build a payload, it never writes to `CachedReads` or the underlying database, but all reads from the underlying database are cached in `CachedReads`.
+///     let db = cached_reads.as_db_mut(db);
+///     // this is `Database` and can be used to build a payload, it never commits to `CachedReads` or the underlying database, but all reads from the underlying database are cached in `CachedReads`.
 ///     // Subsequent payload build attempts can use cached reads and avoid hitting the underlying database.
-///     let db = State::builder().with_database_ref(db_ref).build();
+///     let state = State::builder().with_database(db).build();
 /// }
 /// ```
 #[derive(Debug, Clone, Default)]
@@ -40,10 +40,11 @@ pub struct CachedReads {
 impl CachedReads {
     /// Gets a [`DatabaseRef`] that will cache reads from the given database.
     pub fn as_db<DB>(&mut self, db: DB) -> CachedReadsDBRef<'_, DB> {
-        CachedReadsDBRef { inner: RefCell::new(self.as_db_mut(db)) }
+        self.as_db_mut(db).into_db()
     }
 
-    fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
+    /// Gets a mutable [`Database`] that will cache reads from the underlying database.
+    pub fn as_db_mut<DB>(&mut self, db: DB) -> CachedReadsDbMut<'_, DB> {
         CachedReadsDbMut { cached: self, db }
     }
 
@@ -67,6 +68,19 @@ pub struct CachedReadsDbMut<'a, DB> {
     pub db: DB,
 }
 
+impl<'a, DB> CachedReadsDbMut<'a, DB> {
+    /// Converts this [`Database`] implementation into a [`DatabaseRef`] that will still cache
+    /// reads.
+    pub const fn into_db(self) -> CachedReadsDBRef<'a, DB> {
+        CachedReadsDBRef { inner: RefCell::new(self) }
+    }
+
+    /// Returns access to wrapped [`DatabaseRef`].
+    pub const fn inner(&self) -> &DB {
+        &self.db
+    }
+}
+
 impl<DB: DatabaseRef> Database for CachedReadsDbMut<'_, DB> {
     type Error = <DB as DatabaseRef>::Error;