perf(trie): use ParallelProof::multiproof in StateRootTask

crates/engine/tree/Cargo.toml

@@ -32,7 +32,6 @@ reth-prune.workspace = true
 reth-revm.workspace = true
 reth-stages-api.workspace = true
 reth-tasks.workspace = true
-reth-trie-db.workspace = true
 reth-trie-parallel.workspace = true
 reth-trie-sparse.workspace = true
 reth-trie.workspace = true
@@ -82,6 +81,7 @@ reth-stages = { workspace = true, features = ["test-utils"] }
 reth-static-file.workspace = true
 reth-testing-utils.workspace = true
 reth-tracing.workspace = true
+reth-trie-db.workspace = true
 
 # alloy
 alloy-rlp.workspace = true
@@ -120,6 +120,6 @@ test-utils = [
     "reth-static-file",
     "reth-tracing",
     "reth-trie/test-utils",
-    "reth-prune-types?/test-utils",
     "reth-trie-db/test-utils",
+    "reth-prune-types?/test-utils",
 ]

crates/engine/tree/src/tree/root.rs

@@ -3,23 +3,18 @@
 use alloy_primitives::{map::HashSet, Address};
 use derive_more::derive::Deref;
 use rayon::iter::{ParallelBridge, ParallelIterator};
-use reth_errors::{ProviderError, ProviderResult};
+use reth_errors::ProviderError;
 use reth_evm::system_calls::OnStateHook;
 use reth_provider::{
-    providers::ConsistentDbView, BlockReader, DBProvider, DatabaseProviderFactory,
-    StateCommitmentProvider,
+    providers::ConsistentDbView, BlockReader, DatabaseProviderFactory, StateCommitmentProvider,
 };
 use reth_trie::{
-    hashed_cursor::HashedPostStateCursorFactory,
     prefix_set::TriePrefixSetsMut,
-    proof::Proof,
-    trie_cursor::InMemoryTrieCursorFactory,
     updates::{TrieUpdates, TrieUpdatesSorted},
     HashedPostState, HashedPostStateSorted, HashedStorage, MultiProof, MultiProofTargets, Nibbles,
     TrieInput,
 };
-use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseProof, DatabaseTrieCursorFactory};
-use reth_trie_parallel::root::ParallelStateRootError;
+use reth_trie_parallel::{proof::ParallelProof, root::ParallelStateRootError};
 use reth_trie_sparse::{
     blinded::{BlindedProvider, BlindedProviderFactory},
     errors::{SparseStateTrieError, SparseStateTrieResult, SparseTrieError, SparseTrieErrorKind},
@@ -400,20 +395,31 @@ where
         state_root_message_sender: Sender<StateRootMessage<BPF>>,
     ) {
         // Dispatch proof gathering for this state update
-        scope.spawn(move |_| match calculate_multiproof(config, proof_targets.clone()) {
-            Ok(proof) => {
-                let _ = state_root_message_sender.send(StateRootMessage::ProofCalculated(
-                    Box::new(ProofCalculated {
-                        state_update: hashed_state_update,
-                        targets: proof_targets,
-                        proof,
-                        sequence_number: proof_sequence_number,
-                    }),
-                ));
-            }
-            Err(error) => {
-                let _ =
-                    state_root_message_sender.send(StateRootMessage::ProofCalculationError(error));
+        scope.spawn(move |_| {
+            let result = ParallelProof::new(
+                config.consistent_view.clone(),
+                config.nodes_sorted.clone(),
+                config.state_sorted.clone(),
+                config.prefix_sets.clone(),
+            )
+            .with_branch_node_hash_masks(true)
+            .multiproof(proof_targets.clone());
+
+            match result {
+                Ok(proof) => {
+                    let _ = state_root_message_sender.send(StateRootMessage::ProofCalculated(
+                        Box::new(ProofCalculated {
+                            state_update: hashed_state_update,
+                            targets: proof_targets,
+                            proof,
+                            sequence_number: proof_sequence_number,
+                        }),
+                    ));
+                }
+                Err(error) => {
+                    let _ = state_root_message_sender
+                        .send(StateRootMessage::ProofCalculationError(error.into()));
+                }
             }
         });
     }
@@ -714,31 +720,6 @@ fn get_proof_targets(
     targets
 }
 
-/// Calculate multiproof for the targets.
-#[inline]
-fn calculate_multiproof<Factory>(
-    config: StateRootConfig<Factory>,
-    proof_targets: MultiProofTargets,
-) -> ProviderResult<MultiProof>
-where
-    Factory: DatabaseProviderFactory<Provider: BlockReader> + StateCommitmentProvider,
-{
-    let provider = config.consistent_view.provider_ro()?;
-
-    Ok(Proof::from_tx(provider.tx_ref())
-        .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
-            DatabaseTrieCursorFactory::new(provider.tx_ref()),
-            &config.nodes_sorted,
-        ))
-        .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
-            DatabaseHashedCursorFactory::new(provider.tx_ref()),
-            &config.state_sorted,
-        ))
-        .with_prefix_sets_mut(config.prefix_sets.as_ref().clone())
-        .with_branch_node_hash_masks(true)
-        .multiproof(proof_targets)?)
-}
-
 /// Updates the sparse trie with the given proofs and state, and returns the updated trie and the
 /// time it took.
 fn update_sparse_trie<