feat(engine): share sparse trie pipeline with payload builder

crates/engine/primitives/src/config.rs

@@ -153,6 +153,8 @@ pub struct TreeConfig {
     state_root_task_timeout: Option<Duration>,
     /// Whether to share execution cache with the payload builder.
     share_execution_cache_with_payload_builder: bool,
+    /// Whether to share sparse trie with the payload builder.
+    share_sparse_trie_with_payload_builder: bool,
     /// Maximum random jitter applied before each proof computation (trie-debug only).
     /// When set, each proof worker sleeps for a random duration up to this value
     /// before starting a proof calculation.
@@ -189,6 +191,7 @@ impl Default for TreeConfig {
             disable_sparse_trie_cache_pruning: false,
             state_root_task_timeout: Some(DEFAULT_STATE_ROOT_TASK_TIMEOUT),
             share_execution_cache_with_payload_builder: false,
+            share_sparse_trie_with_payload_builder: false,
             #[cfg(feature = "trie-debug")]
             proof_jitter: None,
         }
@@ -224,6 +227,7 @@ impl TreeConfig {
         slow_block_threshold: Option<Duration>,
         state_root_task_timeout: Option<Duration>,
         share_execution_cache_with_payload_builder: bool,
+        share_sparse_trie_with_payload_builder: bool,
     ) -> Self {
         Self {
             persistence_threshold,
@@ -252,6 +256,7 @@ impl TreeConfig {
             disable_sparse_trie_cache_pruning: false,
             state_root_task_timeout,
             share_execution_cache_with_payload_builder,
+            share_sparse_trie_with_payload_builder,
             #[cfg(feature = "trie-debug")]
             proof_jitter: None,
         }
@@ -569,6 +574,11 @@ impl TreeConfig {
         self.share_execution_cache_with_payload_builder
     }
 
+    /// Returns whether to share sparse trie with the payload builder.
+    pub const fn share_sparse_trie_with_payload_builder(&self) -> bool {
+        self.share_sparse_trie_with_payload_builder
+    }
+
     /// Setter for whether to share execution cache with the payload builder.
     pub const fn with_share_execution_cache_with_payload_builder(
         mut self,
@@ -579,6 +589,15 @@ impl TreeConfig {
         self
     }
 
+    /// Setter for whether to share sparse trie with the payload builder.
+    pub const fn with_share_sparse_trie_with_payload_builder(
+        mut self,
+        share_sparse_trie_with_payload_builder: bool,
+    ) -> Self {
+        self.share_sparse_trie_with_payload_builder = share_sparse_trie_with_payload_builder;
+        self
+    }
+
     /// Returns the proof jitter duration, if configured (trie-debug only).
     #[cfg(feature = "trie-debug")]
     pub const fn proof_jitter(&self) -> Option<Duration> {

crates/engine/tree/src/tree/mod.rs

@@ -3101,12 +3101,23 @@ where
             None
         };
 
+        let trie_handle = if self.config.share_sparse_trie_with_payload_builder() {
+            self.payload_validator.sparse_trie_handle_for(
+                state.head_block_hash,
+                head.state_root(),
+                &self.state,
+            )
+        } else {
+            None
+        };
+
         // send the payload to the builder and return the receiver for the pending payload
         // id, initiating payload job is handled asynchronously
         let pending_payload_id = self.payload_builder.send_new_payload(BuildNewPayload {
             parent_hash: state.head_block_hash,
             attributes,
             cache,
+            trie_handle,
         });
 
         // Client software MUST respond to this method call in the following way:

crates/engine/tree/src/tree/payload_processor/mod.rs

@@ -2,17 +2,13 @@
 
 use super::precompile_cache::PrecompileCacheMap;
 use crate::tree::{
-    payload_processor::{
-        prewarm::{PrewarmCacheTask, PrewarmContext, PrewarmMode, PrewarmTaskEvent},
-        sparse_trie::StateRootComputeOutcome,
-    },
+    payload_processor::prewarm::{PrewarmCacheTask, PrewarmContext, PrewarmMode, PrewarmTaskEvent},
     sparse_trie::SparseTrieCacheTask,
     CacheWaitDurations, CachedStateMetrics, ExecutionCache, PayloadExecutionCache, SavedCache,
     StateProviderBuilder, TreeConfig, WaitForCaches,
 };
 use alloy_eip7928::BlockAccessList;
 use alloy_eips::{eip1898::BlockWithParent, eip4895::Withdrawal};
-use alloy_evm::block::StateChangeSource;
 use alloy_primitives::B256;
 use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender};
 use multiproof::*;
@@ -28,7 +24,7 @@ use reth_primitives_traits::{FastInstant as Instant, NodePrimitives};
 use reth_provider::{
     BlockExecutionOutput, BlockReader, DatabaseProviderROFactory, StateProviderFactory, StateReader,
 };
-use reth_revm::{db::BundleState, state::EvmState};
+use reth_revm::db::BundleState;
 use reth_tasks::{utils::increase_thread_priority, ForEachOrdered, Runtime};
 use reth_trie::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory};
 use reth_trie_parallel::{
@@ -271,12 +267,18 @@ where
 
         let span = Span::current();
 
-        let state_root_handle = self.spawn_state_root(multiproof_provider_factory, &env, config);
+        let halve_workers = env.transaction_count <= Self::SMALL_BLOCK_PROOF_WORKER_TX_THRESHOLD;
+        let state_root_handle = self.spawn_state_root(
+            multiproof_provider_factory,
+            env.parent_state_root,
+            halve_workers,
+            config,
+        );
         let prewarm_handle = self.spawn_caching_with(
             env,
             prewarm_rx,
             provider_builder,
-            Some(state_root_handle.to_multi_proof.clone()),
+            Some(state_root_handle.updates_tx().clone()),
             bal,
         );
 
@@ -322,11 +324,15 @@ where
     ///   state root.
     ///
     /// The state hook **must** be dropped after execution to signal the end of state updates.
+    ///
+    /// When `halve_workers` is true, the proof worker pool is halved (for small blocks where
+    /// fewer transactions produce fewer state changes and most workers would be idle).
     #[instrument(level = "debug", target = "engine::tree::payload_processor", skip_all)]
     pub fn spawn_state_root<F>(
-        &mut self,
+        &self,
         multiproof_provider_factory: F,
-        env: &ExecutionEnv<Evm>,
+        parent_state_root: B256,
+        halve_workers: bool,
         config: &TreeConfig,
     ) -> StateRootHandle
     where
@@ -336,12 +342,11 @@ where
             + Sync
             + 'static,
     {
-        let (to_multi_proof, from_multi_proof) = crossbeam_channel::unbounded();
+        let (updates_tx, from_multi_proof) = crossbeam_channel::unbounded();
 
         let task_ctx = ProofTaskCtx::new(multiproof_provider_factory);
         #[cfg(feature = "trie-debug")]
         let task_ctx = task_ctx.with_proof_jitter(config.proof_jitter());
-        let halve_workers = env.transaction_count <= Self::SMALL_BLOCK_PROOF_WORKER_TX_THRESHOLD;
         let proof_handle = ProofWorkerHandle::new(&self.executor, task_ctx, halve_workers);
 
         let (state_root_tx, state_root_rx) = channel();
@@ -350,11 +355,11 @@ where
             proof_handle,
             state_root_tx,
             from_multi_proof,
-            env.parent_state_root,
+            parent_state_root,
             config.multiproof_chunk_size(),
         );
 
-        StateRootHandle::new(to_multi_proof, state_root_rx)
+        StateRootHandle::new(parent_state_root, updates_tx, state_root_rx)
     }
 
     /// Transaction count threshold below which proof workers are halved, since fewer transactions
@@ -742,67 +747,6 @@ fn convert_serial<RawTx, Tx, TxEnv, InnerTx, Recovered, Err, C>(
     }
 }
 
-/// Handle to a background state root computation task.
-///
-/// Unlike [`PayloadHandle`], this does not include transaction iteration or cache prewarming.
-/// It only provides access to the state root computation via [`Self::state_hook`] and
-/// [`Self::state_root`].
-///
-/// Created by [`PayloadProcessor::spawn_state_root`].
-#[derive(Debug)]
-pub struct StateRootHandle {
-    /// Channel for evm state updates to the multiproof pipeline.
-    to_multi_proof: CrossbeamSender<MultiProofMessage>,
-    /// Receiver for the computed state root.
-    state_root_rx: Option<mpsc::Receiver<Result<StateRootComputeOutcome, ParallelStateRootError>>>,
-}
-
-impl StateRootHandle {
-    /// Creates a new state root handle.
-    pub const fn new(
-        to_multi_proof: CrossbeamSender<MultiProofMessage>,
-        state_root_rx: mpsc::Receiver<Result<StateRootComputeOutcome, ParallelStateRootError>>,
-    ) -> Self {
-        Self { to_multi_proof, state_root_rx: Some(state_root_rx) }
-    }
-
-    /// Returns a state hook that streams state updates to the background state root task.
-    ///
-    /// The hook must be dropped after execution completes to signal the end of state updates.
-    pub fn state_hook(&self) -> impl OnStateHook {
-        let to_multi_proof = StateHookSender::new(self.to_multi_proof.clone());
-
-        move |source: StateChangeSource, state: &EvmState| {
-            let _ =
-                to_multi_proof.send(MultiProofMessage::StateUpdate(source.into(), state.clone()));
-        }
-    }
-
-    /// Awaits the state root computation result.
-    ///
-    /// # Panics
-    ///
-    /// If called more than once.
-    pub fn state_root(&mut self) -> Result<StateRootComputeOutcome, ParallelStateRootError> {
-        self.state_root_rx
-            .take()
-            .expect("state_root already taken")
-            .recv()
-            .map_err(|_| ParallelStateRootError::Other("sparse trie task dropped".to_string()))?
-    }
-
-    /// Takes the state root receiver for use with custom waiting logic (e.g., timeouts).
-    ///
-    /// # Panics
-    ///
-    /// If called more than once.
-    pub const fn take_state_root_rx(
-        &mut self,
-    ) -> mpsc::Receiver<Result<StateRootComputeOutcome, ParallelStateRootError>> {
-        self.state_root_rx.take().expect("state_root already taken")
-    }
-}
-
 /// Handle to all the spawned tasks.
 ///
 /// Generic over `R` (receipt type) to allow sharing `Arc<ExecutionOutcome<R>>` with the