refactor(trie): move proof database related operations to an extension trait in reth-trie-db crate (#9743)

crates/storage/provider/src/providers/state/historical.rs

@@ -15,8 +15,8 @@ use reth_primitives::{
 };
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, StateRoot};
-use reth_trie_db::DatabaseStateRoot;
+use reth_trie::{proof::Proof, updates::TrieUpdates, AccountProof, HashedPostState, StateRoot};
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot};
 use std::fmt::Debug;
 
 /// State provider for a given block number which takes a tx reference.
@@ -285,8 +285,7 @@ impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> {
     ) -> ProviderResult<AccountProof> {
         let mut revert_state = self.revert_state()?;
         revert_state.extend(hashed_state.clone());
-        revert_state
-            .account_proof(self.tx, address, slots)
+        Proof::overlay_account_proof(self.tx, revert_state, address, slots)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 }

crates/storage/provider/src/providers/state/latest.rs

@@ -12,8 +12,8 @@ use reth_primitives::{
 };
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, StateRoot};
-use reth_trie_db::DatabaseStateRoot;
+use reth_trie::{proof::Proof, updates::TrieUpdates, AccountProof, HashedPostState, StateRoot};
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot};
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -96,8 +96,7 @@ impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        Ok(hashed_state
-            .account_proof(self.tx, address, slots)
+        Ok(Proof::overlay_account_proof(self.tx, hashed_state.clone(), address, slots)
             .map_err(Into::<reth_db::DatabaseError>::into)?)
     }
 }

crates/trie/db/src/lib.rs

@@ -1,7 +1,9 @@
 //! An integration of [`reth-trie`] with [`reth-db`].
 
+mod proof;
 mod state;
 mod storage;
 
+pub use proof::DatabaseProof;
 pub use state::DatabaseStateRoot;
 pub use storage::DatabaseStorageRoot;

crates/trie/db/src/proof.rs

@@ -0,0 +1,46 @@
+use reth_db_api::transaction::DbTx;
+use reth_execution_errors::StateRootError;
+use reth_primitives::{Address, B256};
+use reth_trie::{
+    hashed_cursor::{DatabaseHashedCursorFactory, HashedPostStateCursorFactory},
+    proof::Proof,
+    HashedPostState,
+};
+use reth_trie_common::AccountProof;
+
+/// Extends [`Proof`] with operations specific for working with a database transaction.
+pub trait DatabaseProof<'a, TX> {
+    /// Create a new [Proof] from database transaction.
+    fn from_tx(tx: &'a TX) -> Self;
+
+    /// Generates the state proof for target account and slots on top of this [`HashedPostState`].
+    fn overlay_account_proof(
+        tx: &'a TX,
+        post_state: HashedPostState,
+        address: Address,
+        slots: &[B256],
+    ) -> Result<AccountProof, StateRootError>;
+}
+
+impl<'a, TX: DbTx> DatabaseProof<'a, TX> for Proof<&'a TX, DatabaseHashedCursorFactory<'a, TX>> {
+    /// Create a new [Proof] instance from database transaction.
+    fn from_tx(tx: &'a TX) -> Self {
+        Self::new(tx, DatabaseHashedCursorFactory::new(tx))
+    }
+
+    fn overlay_account_proof(
+        tx: &'a TX,
+        post_state: HashedPostState,
+        address: Address,
+        slots: &[B256],
+    ) -> Result<AccountProof, StateRootError> {
+        let prefix_sets = post_state.construct_prefix_sets();
+        let sorted = post_state.into_sorted();
+        let hashed_cursor_factory =
+            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &sorted);
+        Proof::from_tx(tx)
+            .with_hashed_cursor_factory(hashed_cursor_factory)
+            .with_prefix_sets_mut(prefix_sets)
+            .account_proof(address, slots)
+    }
+}

crates/trie/db/tests/proof.rs

@@ -10,7 +10,7 @@ use reth_provider::{
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{proof::Proof, Nibbles, StateRoot};
 use reth_trie_common::{AccountProof, StorageProof};
-use reth_trie_db::DatabaseStateRoot;
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot};
 use std::{str::FromStr, sync::Arc};
 
 /*

crates/trie/trie/src/proof.rs

@@ -1,14 +1,12 @@
 use crate::{
-    hashed_cursor::{DatabaseHashedCursorFactory, HashedCursorFactory, HashedStorageCursor},
+    hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
     node_iter::{TrieElement, TrieNodeIter},
     prefix_set::TriePrefixSetsMut,
-    trie_cursor::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor},
+    trie_cursor::TrieCursorFactory,
     walker::TrieWalker,
     HashBuilder, Nibbles,
 };
 use alloy_rlp::{BufMut, Encodable};
-use reth_db::tables;
-use reth_db_api::transaction::DbTx;
 use reth_execution_errors::{StateRootError, StorageRootError};
 use reth_primitives::{constants::EMPTY_ROOT_HASH, keccak256, Address, B256};
 use reth_trie_common::{proof::ProofRetainer, AccountProof, StorageProof, TrieAccount};
@@ -19,24 +17,32 @@ use reth_trie_common::{proof::ProofRetainer, AccountProof, StorageProof, TrieAcc
 /// on the hash builder and follows the same algorithm as the state root calculator.
 /// See `StateRoot::root` for more info.
 #[derive(Debug)]
-pub struct Proof<'a, TX, H> {
-    /// A reference to the database transaction.
-    tx: &'a TX,
+pub struct Proof<T, H> {
     /// The factory for hashed cursors.
     hashed_cursor_factory: H,
+    /// Creates cursor for traversing trie entities.
+    trie_cursor_factory: T,
     /// A set of prefix sets that have changes.
     prefix_sets: TriePrefixSetsMut,
 }
 
-impl<'a, TX, H> Proof<'a, TX, H> {
-    /// Creates a new proof generator.
-    pub fn new(tx: &'a TX, hashed_cursor_factory: H) -> Self {
-        Self { tx, hashed_cursor_factory, prefix_sets: TriePrefixSetsMut::default() }
+impl<T, H> Proof<T, H> {
+    /// Create a new [Proof] instance.
+    pub fn new(t: T, h: H) -> Self {
+        Self {
+            trie_cursor_factory: t,
+            hashed_cursor_factory: h,
+            prefix_sets: TriePrefixSetsMut::default(),
+        }
     }
 
     /// Set the hashed cursor factory.
-    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> Proof<'a, TX, HF> {
-        Proof { tx: self.tx, hashed_cursor_factory, prefix_sets: self.prefix_sets }
+    pub fn with_hashed_cursor_factory<HF>(self, hashed_cursor_factory: HF) -> Proof<T, HF> {
+        Proof {
+            trie_cursor_factory: self.trie_cursor_factory,
+            hashed_cursor_factory,
+            prefix_sets: self.prefix_sets,
+        }
     }
 
     /// Set the prefix sets. They have to be mutable in order to allow extension with proof target.
@@ -46,16 +52,9 @@ impl<'a, TX, H> Proof<'a, TX, H> {
     }
 }
 
-impl<'a, TX> Proof<'a, TX, DatabaseHashedCursorFactory<'a, TX>> {
-    /// Create a new [Proof] instance from database transaction.
-    pub fn from_tx(tx: &'a TX) -> Self {
-        Self::new(tx, DatabaseHashedCursorFactory::new(tx))
-    }
-}
-
-impl<'a, TX, H> Proof<'a, TX, H>
+impl<T, H> Proof<T, H>
 where
-    TX: DbTx,
+    T: TrieCursorFactory,
     H: HashedCursorFactory + Clone,
 {
     /// Generate an account proof from intermediate nodes.
@@ -69,8 +68,7 @@ where
         let mut account_proof = AccountProof::new(address);
 
         let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
-        let trie_cursor =
-            DatabaseAccountTrieCursor::new(self.tx.cursor_read::<tables::AccountsTrie>()?);
+        let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
 
         // Create the walker.
         let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
@@ -141,10 +139,7 @@ where
         let mut prefix_set =
             self.prefix_sets.storage_prefix_sets.get(&hashed_address).cloned().unwrap_or_default();
         prefix_set.extend(target_nibbles.clone());
-        let trie_cursor = DatabaseStorageTrieCursor::new(
-            self.tx.cursor_dup_read::<tables::StoragesTrie>()?,
-            hashed_address,
-        );
+        let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(hashed_address)?;
         let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
 
         let retainer = ProofRetainer::from_iter(target_nibbles);

crates/trie/trie/src/state.rs

@@ -1,7 +1,5 @@
 use crate::{
-    hashed_cursor::{DatabaseHashedCursorFactory, HashedPostStateCursorFactory},
     prefix_set::{PrefixSetMut, TriePrefixSetsMut},
-    proof::Proof,
     Nibbles,
 };
 use itertools::Itertools;
@@ -12,9 +10,7 @@ use reth_db_api::{
     models::{AccountBeforeTx, BlockNumberAddress},
     transaction::DbTx,
 };
-use reth_execution_errors::StateRootError;
 use reth_primitives::{keccak256, Account, Address, BlockNumber, B256, U256};
-use reth_trie_common::AccountProof;
 use revm::db::BundleAccount;
 use std::collections::{hash_map, HashMap, HashSet};
 
@@ -193,23 +189,6 @@ impl HashedPostState {
 
         TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
     }
-
-    /// Generates the state proof for target account and slots on top of this [`HashedPostState`].
-    pub fn account_proof<TX: DbTx>(
-        &self,
-        tx: &TX,
-        address: Address,
-        slots: &[B256],
-    ) -> Result<AccountProof, StateRootError> {
-        let sorted = self.clone().into_sorted();
-        let prefix_sets = self.construct_prefix_sets();
-        let hashed_cursor_factory =
-            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &sorted);
-        Proof::from_tx(tx)
-            .with_hashed_cursor_factory(hashed_cursor_factory)
-            .with_prefix_sets_mut(prefix_sets)
-            .account_proof(address, slots)
-    }
 }
 
 /// Representation of in-memory hashed storage.