perf(engine): return sorted data from compute_trie_input (#19340)

Author: yongkangc

crates/chain-state/src/in_memory.rs

@@ -17,7 +17,7 @@ use reth_primitives_traits::{
     SignedTransaction,
 };
 use reth_storage_api::StateProviderBox;
-use reth_trie::{updates::TrieUpdates, HashedPostState};
+use reth_trie::{updates::TrieUpdatesSorted, HashedPostStateSorted};
 use std::{collections::BTreeMap, sync::Arc, time::Instant};
 use tokio::sync::{broadcast, watch};
 
@@ -725,10 +725,10 @@ pub struct ExecutedBlock<N: NodePrimitives = EthPrimitives> {
     pub recovered_block: Arc<RecoveredBlock<N::Block>>,
     /// Block's execution outcome.
     pub execution_output: Arc<ExecutionOutcome<N::Receipt>>,
-    /// Block's hashed state.
-    pub hashed_state: Arc<HashedPostState>,
-    /// Trie updates that result from calculating the state root for the block.
-    pub trie_updates: Arc<TrieUpdates>,
+    /// Block's sorted hashed state.
+    pub hashed_state: Arc<HashedPostStateSorted>,
+    /// Sorted trie updates that result from calculating the state root for the block.
+    pub trie_updates: Arc<TrieUpdatesSorted>,
 }
 
 impl<N: NodePrimitives> Default for ExecutedBlock<N> {
@@ -763,13 +763,13 @@ impl<N: NodePrimitives> ExecutedBlock<N> {
 
     /// Returns a reference to the hashed state result of the execution outcome
     #[inline]
-    pub fn hashed_state(&self) -> &HashedPostState {
+    pub fn hashed_state(&self) -> &HashedPostStateSorted {
         &self.hashed_state
     }
 
     /// Returns a reference to the trie updates resulting from the execution outcome
     #[inline]
-    pub fn trie_updates(&self) -> &TrieUpdates {
+    pub fn trie_updates(&self) -> &TrieUpdatesSorted {
         &self.trie_updates
     }
 
@@ -875,8 +875,8 @@ mod tests {
         StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
     };
     use reth_trie::{
-        AccountProof, HashedStorage, MultiProof, MultiProofTargets, StorageMultiProof,
-        StorageProof, TrieInput,
+        updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
+        MultiProofTargets, StorageMultiProof, StorageProof, TrieInput,
     };
 
     fn create_mock_state(

crates/chain-state/src/memory_overlay.rs

@@ -53,7 +53,7 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
     /// Return lazy-loaded trie state aggregated from in-memory blocks.
     fn trie_input(&self) -> &TrieInput {
         self.trie_input.get_or_init(|| {
-            TrieInput::from_blocks(
+            TrieInput::from_blocks_sorted(
                 self.in_memory
                     .iter()
                     .rev()

crates/chain-state/src/test_utils.rs

@@ -23,7 +23,7 @@ use reth_primitives_traits::{
     SignedTransaction,
 };
 use reth_storage_api::NodePrimitivesProvider;
-use reth_trie::{root::state_root_unhashed, updates::TrieUpdates, HashedPostState};
+use reth_trie::{root::state_root_unhashed, updates::TrieUpdatesSorted, HashedPostStateSorted};
 use revm_database::BundleState;
 use revm_state::AccountInfo;
 use std::{
@@ -216,8 +216,8 @@ impl<N: NodePrimitives> TestBlockBuilder<N> {
                 block_number,
                 vec![Requests::default()],
             )),
-            hashed_state: Arc::new(HashedPostState::default()),
-            trie_updates: Arc::new(TrieUpdates::default()),
+            hashed_state: Arc::new(HashedPostStateSorted::default()),
+            trie_updates: Arc::new(TrieUpdatesSorted::default()),
         }
     }
 

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

@@ -1794,8 +1794,8 @@ where
         Ok(Some(ExecutedBlock {
             recovered_block: Arc::new(RecoveredBlock::new_sealed(block, senders)),
             execution_output: Arc::new(execution_output),
-            hashed_state: Arc::new(hashed_state),
-            trie_updates: Arc::new(trie_updates.into()),
+            hashed_state: Arc::new(hashed_state.into_sorted()),
+            trie_updates: Arc::new(trie_updates),
         }))
     }
 

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

@@ -13,9 +13,8 @@ use metrics::{Gauge, Histogram};
 use reth_metrics::Metrics;
 use reth_revm::state::EvmState;
 use reth_trie::{
-    added_removed_keys::MultiAddedRemovedKeys, prefix_set::TriePrefixSetsMut,
-    updates::TrieUpdatesSorted, DecodedMultiProof, HashedPostState, HashedPostStateSorted,
-    HashedStorage, MultiProofTargets, TrieInput,
+    added_removed_keys::MultiAddedRemovedKeys, DecodedMultiProof, HashedPostState, HashedStorage,
+    MultiProofTargets,
 };
 use reth_trie_parallel::{
     proof::ParallelProof,
@@ -60,35 +59,6 @@ impl SparseTrieUpdate {
     }
 }
 
-/// Common configuration for multi proof tasks
-#[derive(Debug, Clone, Default)]
-pub(crate) struct MultiProofConfig {
-    /// The sorted collection of cached in-memory intermediate trie nodes that
-    /// can be reused for computation.
-    pub nodes_sorted: Arc<TrieUpdatesSorted>,
-    /// The sorted in-memory overlay hashed state.
-    pub state_sorted: Arc<HashedPostStateSorted>,
-    /// The collection of prefix sets for the computation. Since the prefix sets _always_
-    /// invalidate the in-memory nodes, not all keys from `state_sorted` might be present here,
-    /// if we have cached nodes for them.
-    pub prefix_sets: Arc<TriePrefixSetsMut>,
-}
-
-impl MultiProofConfig {
-    /// Creates a new state root config from the trie input.
-    ///
-    /// This returns a cleared [`TrieInput`] so that we can reuse any allocated space in the
-    /// [`TrieInput`].
-    pub(crate) fn from_input(mut input: TrieInput) -> (TrieInput, Self) {
-        let config = Self {
-            nodes_sorted: Arc::new(input.nodes.drain_into_sorted()),
-            state_sorted: Arc::new(input.state.drain_into_sorted()),
-            prefix_sets: Arc::new(input.prefix_sets.clone()),
-        };
-        (input.cleared(), config)
-    }
-}
-
 /// Messages used internally by the multi proof task.
 #[derive(Debug)]
 pub(super) enum MultiProofMessage {

crates/engine/tree/src/tree/payload_validator.rs

@@ -5,7 +5,7 @@ use crate::tree::{
     error::{InsertBlockError, InsertBlockErrorKind, InsertPayloadError},
     executor::WorkloadExecutor,
     instrumented_state::InstrumentedStateProvider,
-    payload_processor::{multiproof::MultiProofConfig, PayloadProcessor},
+    payload_processor::PayloadProcessor,
     precompile_cache::{CachedPrecompile, CachedPrecompileMetrics, PrecompileCacheMap},
     sparse_trie::StateRootComputeOutcome,
     EngineApiMetrics, EngineApiTreeState, ExecutionEnv, PayloadHandle, StateProviderBuilder,
@@ -38,7 +38,7 @@ use reth_provider::{
     StateRootProvider, TrieReader,
 };
 use reth_revm::db::State;
-use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInput};
+use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInputSorted};
 use reth_trie_parallel::root::{ParallelStateRoot, ParallelStateRootError};
 use std::{collections::HashMap, sync::Arc, time::Instant};
 use tracing::{debug, debug_span, error, info, instrument, trace, warn};
@@ -121,8 +121,6 @@ where
     metrics: EngineApiMetrics,
     /// Validator for the payload.
     validator: V,
-    /// A cleared trie input, kept around to be reused so allocations can be minimized.
-    trie_input: Option<TrieInput>,
 }
 
 impl<N, P, Evm, V> BasicEngineValidator<P, Evm, V>
@@ -166,7 +164,6 @@ where
             invalid_block_hook,
             metrics: EngineApiMetrics::default(),
             validator,
-            trie_input: Default::default(),
         }
     }
 
@@ -530,8 +527,8 @@ where
         Ok(ExecutedBlock {
             recovered_block: Arc::new(block),
             execution_output: Arc::new(ExecutionOutcome::from((output, block_num_hash.number))),
-            hashed_state: Arc::new(hashed_state),
-            trie_updates: Arc::new(trie_output),
+            hashed_state: Arc::new(hashed_state.into_sorted()),
+            trie_updates: Arc::new(trie_output.into_sorted()),
         })
     }
 
@@ -641,26 +638,24 @@ where
         hashed_state: &HashedPostState,
         state: &EngineApiTreeState<N>,
     ) -> Result<(B256, TrieUpdates), ParallelStateRootError> {
-        let (mut input, block_hash) = self.compute_trie_input(parent_hash, state, None)?;
+        let (mut input, block_hash) = self.compute_trie_input(parent_hash, state)?;
 
-        // Extend with block we are validating root for.
-        input.append_ref(hashed_state);
+        // Extend state overlay with current block's sorted state.
+        input.prefix_sets.extend(hashed_state.construct_prefix_sets());
+        let sorted_hashed_state = hashed_state.clone().into_sorted();
+        Arc::make_mut(&mut input.state).extend_ref(&sorted_hashed_state);
 
-        // Convert the TrieInput into a MultProofConfig, since everything uses the sorted
-        // forms of the state/trie fields.
-        let (_, multiproof_config) = MultiProofConfig::from_input(input);
+        let TrieInputSorted { nodes, state, prefix_sets: prefix_sets_mut } = input;
 
         let factory = OverlayStateProviderFactory::new(self.provider.clone())
             .with_block_hash(Some(block_hash))
-            .with_trie_overlay(Some(multiproof_config.nodes_sorted))
-            .with_hashed_state_overlay(Some(multiproof_config.state_sorted));
+            .with_trie_overlay(Some(nodes))
+            .with_hashed_state_overlay(Some(state));
 
         // The `hashed_state` argument is already taken into account as part of the overlay, but we
         // need to use the prefix sets which were generated from it to indicate to the
         // ParallelStateRoot which parts of the trie need to be recomputed.
-        let prefix_sets = Arc::into_inner(multiproof_config.prefix_sets)
-            .expect("MultiProofConfig was never cloned")
-            .freeze();
+        let prefix_sets = prefix_sets_mut.freeze();
 
         ParallelStateRoot::new(factory, prefix_sets).incremental_root_with_updates()
     }
@@ -756,26 +751,23 @@ where
     > {
         match strategy {
             StateRootStrategy::StateRootTask => {
-                // get allocated trie input if it exists
-                let allocated_trie_input = self.trie_input.take();
-
                 // Compute trie input
                 let trie_input_start = Instant::now();
-                let (trie_input, block_hash) =
-                    self.compute_trie_input(parent_hash, state, allocated_trie_input)?;
+                let (trie_input, block_hash) = self.compute_trie_input(parent_hash, state)?;
+
+                self.metrics
+                    .block_validation
+                    .trie_input_duration
+                    .record(trie_input_start.elapsed().as_secs_f64());
 
-                // Convert the TrieInput into a MultProofConfig, since everything uses the sorted
-                // forms of the state/trie fields.
-                let (trie_input, multiproof_config) = MultiProofConfig::from_input(trie_input);
-                self.trie_input.replace(trie_input);
+                // Create OverlayStateProviderFactory with sorted trie data for multiproofs
+                let TrieInputSorted { nodes, state, .. } = trie_input;
 
-                // Create OverlayStateProviderFactory with the multiproof config, for use with
-                // multiproofs.
                 let multiproof_provider_factory =
                     OverlayStateProviderFactory::new(self.provider.clone())
                         .with_block_hash(Some(block_hash))
-                        .with_trie_overlay(Some(multiproof_config.nodes_sorted))
-                        .with_hashed_state_overlay(Some(multiproof_config.state_sorted));
+                        .with_trie_overlay(Some(nodes))
+                        .with_hashed_state_overlay(Some(state));
 
                 // Use state root task only if prefix sets are empty, otherwise proof generation is
                 // too expensive because it requires walking all paths in every proof.
@@ -887,14 +879,14 @@ where
     /// Computes the trie input at the provided parent hash, as well as the block number of the
     /// highest persisted ancestor.
     ///
-    /// The goal of this function is to take in-memory blocks and generate a [`TrieInput`] that
-    /// serves as an overlay to the database blocks.
+    /// The goal of this function is to take in-memory blocks and generate a [`TrieInputSorted`]
+    /// that serves as an overlay to the database blocks.
     ///
     /// It works as follows:
     /// 1. Collect in-memory blocks that are descendants of the provided parent hash using
     ///    [`crate::tree::TreeState::blocks_by_hash`]. This returns the highest persisted ancestor
     ///    hash (`block_hash`) and the list of in-memory descendant blocks.
-    /// 2. Extend the `TrieInput` with the contents of these in-memory blocks (from oldest to
+    /// 2. Extend the `TrieInputSorted` with the contents of these in-memory blocks (from oldest to
     ///    newest) to build the overlay state and trie updates that sit on top of the database view
     ///    anchored at `block_hash`.
     #[instrument(
@@ -907,11 +899,7 @@ where
         &self,
         parent_hash: B256,
         state: &EngineApiTreeState<N>,
-        allocated_trie_input: Option<TrieInput>,
-    ) -> ProviderResult<(TrieInput, B256)> {
-        // get allocated trie input or use a default trie input
-        let mut input = allocated_trie_input.unwrap_or_default();
-
+    ) -> ProviderResult<(TrieInputSorted, B256)> {
         let (block_hash, blocks) =
             state.tree_state.blocks_by_hash(parent_hash).unwrap_or_else(|| (parent_hash, vec![]));
 
@@ -921,10 +909,24 @@ where
             debug!(target: "engine::tree::payload_validator", historical = ?block_hash, blocks = blocks.len(), "Parent found in memory");
         }
 
-        // Extend with contents of parent in-memory blocks.
-        input.extend_with_blocks(
-            blocks.iter().rev().map(|block| (block.hashed_state(), block.trie_updates())),
-        );
+        // Extend with contents of parent in-memory blocks directly in sorted form.
+        let mut input = TrieInputSorted::default();
+        let mut blocks_iter = blocks.iter().rev().peekable();
+
+        if let Some(first) = blocks_iter.next() {
+            input.state = Arc::clone(&first.hashed_state);
+            input.nodes = Arc::clone(&first.trie_updates);
+
+            // Only clone and mutate if there are more in-memory blocks.
+            if blocks_iter.peek().is_some() {
+                let state_mut = Arc::make_mut(&mut input.state);
+                let nodes_mut = Arc::make_mut(&mut input.nodes);
+                for block in blocks_iter {
+                    state_mut.extend_ref(block.hashed_state());
+                    nodes_mut.extend_ref(block.trie_updates());
+                }
+            }
+        }
 
         Ok((input, block_hash))
     }

crates/engine/tree/src/tree/tests.rs

@@ -826,8 +826,8 @@ fn test_tree_state_on_new_head_deep_fork() {
         test_harness.tree.state.tree_state.insert_executed(ExecutedBlock {
             recovered_block: Arc::new(block.clone()),
             execution_output: Arc::new(ExecutionOutcome::default()),
-            hashed_state: Arc::new(HashedPostState::default()),
-            trie_updates: Arc::new(TrieUpdates::default()),
+            hashed_state: Arc::new(HashedPostState::default().into_sorted()),
+            trie_updates: Arc::new(TrieUpdates::default().into_sorted()),
         });
     }
     test_harness.tree.state.tree_state.set_canonical_head(chain_a.last().unwrap().num_hash());
@@ -836,8 +836,8 @@ fn test_tree_state_on_new_head_deep_fork() {
         test_harness.tree.state.tree_state.insert_executed(ExecutedBlock {
             recovered_block: Arc::new(block.clone()),
             execution_output: Arc::new(ExecutionOutcome::default()),
-            hashed_state: Arc::new(HashedPostState::default()),
-            trie_updates: Arc::new(TrieUpdates::default()),
+            hashed_state: Arc::new(HashedPostState::default().into_sorted()),
+            trie_updates: Arc::new(TrieUpdates::default().into_sorted()),
         });
     }
 

crates/optimism/flashblocks/src/worker.rs

@@ -125,7 +125,7 @@ where
             ExecutedBlock {
                 recovered_block: block.into(),
                 execution_output: Arc::new(execution_outcome),
-                hashed_state: Arc::new(hashed_state),
+                hashed_state: Arc::new(hashed_state.into_sorted()),
                 trie_updates: Arc::default(),
             },
         );

crates/optimism/payload/src/builder.rs

@@ -386,6 +386,7 @@ impl<Txs> OpBuilder<'_, Txs> {
         let executed: BuiltPayloadExecutedBlock<N> = BuiltPayloadExecutedBlock {
             recovered_block: Arc::new(block),
             execution_output: Arc::new(execution_outcome),
+            // Keep unsorted; conversion to sorted happens when needed downstream
             hashed_state: either::Either::Left(Arc::new(hashed_state)),
             trie_updates: either::Either::Left(Arc::new(trie_updates)),
         };

crates/payload/primitives/src/traits.rs

@@ -42,17 +42,21 @@ pub struct BuiltPayloadExecutedBlock<N: NodePrimitives> {
 impl<N: NodePrimitives> BuiltPayloadExecutedBlock<N> {
     /// Converts this into an [`reth_chain_state::ExecutedBlock`].
     ///
-    /// If the hashed state or trie updates are in sorted form, they will be converted
-    /// back to their unsorted representations.
+    /// Ensures hashed state and trie updates are in their sorted representations
+    /// as required by `reth_chain_state::ExecutedBlock`.
     pub fn into_executed_payload(self) -> reth_chain_state::ExecutedBlock<N> {
         let hashed_state = match self.hashed_state {
-            Either::Left(unsorted) => unsorted,
-            Either::Right(sorted) => Arc::new(Arc::unwrap_or_clone(sorted).into()),
+            // Convert unsorted to sorted
+            Either::Left(unsorted) => Arc::new(Arc::unwrap_or_clone(unsorted).into_sorted()),
+            // Already sorted
+            Either::Right(sorted) => sorted,
         };
 
         let trie_updates = match self.trie_updates {
-            Either::Left(unsorted) => unsorted,
-            Either::Right(sorted) => Arc::new(Arc::unwrap_or_clone(sorted).into()),
+            // Convert unsorted to sorted
+            Either::Left(unsorted) => Arc::new(Arc::unwrap_or_clone(unsorted).into_sorted()),
+            // Already sorted
+            Either::Right(sorted) => sorted,
         };
 
         reth_chain_state::ExecutedBlock {