perf(engine): return sorted data from compute_trie_input

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/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/payload_processor/multiproof.rs

@@ -15,7 +15,7 @@ use reth_revm::state::EvmState;
 use reth_trie::{
     added_removed_keys::MultiAddedRemovedKeys, prefix_set::TriePrefixSetsMut,
     updates::TrieUpdatesSorted, DecodedMultiProof, HashedPostState, HashedPostStateSorted,
-    HashedStorage, MultiProofTargets, TrieInput,
+    HashedStorage, MultiProofTargets, TrieInputSorted,
 };
 use reth_trie_parallel::{
     proof::ParallelProof,
@@ -73,12 +73,11 @@ pub(crate) struct MultiProofConfig {
 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) {
+    /// This returns a cleared [`TrieInputSorted`] so that we can reuse any allocated space.
+    pub(crate) fn from_input(mut input: TrieInputSorted) -> (TrieInputSorted, Self) {
         let config = Self {
-            nodes_sorted: Arc::new(input.nodes.drain_into_sorted()),
-            state_sorted: Arc::new(input.state.drain_into_sorted()),
+            nodes_sorted: Arc::new(core::mem::take(&mut input.nodes)),
+            state_sorted: Arc::new(core::mem::take(&mut input.state)),
             prefix_sets: Arc::new(input.prefix_sets.clone()),
         };
         (input.cleared(), config)

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

@@ -39,7 +39,7 @@ use reth_provider::{
     StateRootProvider, TrieReader,
 };
 use reth_revm::db::State;
-use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInput};
+use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInput, 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};
@@ -682,7 +682,7 @@ where
     ) -> Result<(B256, TrieUpdates), ParallelStateRootError> {
         let provider = self.provider.database_provider_ro()?;
 
-        let (mut input, block_number) =
+        let (mut input, _reusable_input, block_number) =
             self.compute_trie_input(provider, parent_hash, state, None)?;
 
         // Extend with block we are validating root for.
@@ -802,15 +802,15 @@ where
         match strategy {
             StateRootStrategy::StateRootTask => {
                 // get allocated trie input if it exists
-                let allocated_trie_input = self.trie_input.take();
+                let allocated = self.trie_input.take();
 
                 // Compute trie input
                 let trie_input_start = Instant::now();
-                let (trie_input, block_number) = self.compute_trie_input(
+                let (trie_input, mut reusable_input, block_number) = self.compute_trie_input(
                     self.provider.database_provider_ro()?,
                     parent_hash,
                     state,
-                    allocated_trie_input,
+                    allocated,
                 )?;
 
                 self.metrics
@@ -820,8 +820,13 @@ where
 
                 // 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);
+                let (mut cleared_sorted_input, multiproof_config) =
+                    MultiProofConfig::from_input(trie_input);
+
+                // Rescue the prefix_sets from the cleared sorted input and attach to reusable
+                // builder
+                reusable_input.prefix_sets = core::mem::take(&mut cleared_sorted_input.prefix_sets);
+                self.trie_input.replace(reusable_input);
 
                 // Create OverlayStateProviderFactory with the multiproof config, for use with
                 // multiproofs.
@@ -982,9 +987,10 @@ where
         parent_hash: B256,
         state: &EngineApiTreeState<N>,
         allocated_trie_input: Option<TrieInput>,
-    ) -> ProviderResult<(TrieInput, BlockNumber)> {
-        // get allocated trie input or use a default trie input
+    ) -> ProviderResult<(TrieInputSorted, TrieInput, BlockNumber)> {
+        // Build with unsorted structures for fast append operations, reusing allocated capacity.
         let mut input = allocated_trie_input.unwrap_or_default();
+        input.clear(); // Keep HashMap capacity, drop contents
 
         let (historical, blocks) = state
             .tree_state
@@ -1002,12 +1008,19 @@ where
             .convert_hash_or_number(historical)?
             .ok_or_else(|| ProviderError::BlockHashNotFound(historical.as_hash().unwrap()))?;
 
-        // Extend with contents of parent in-memory blocks.
-        input.extend_with_blocks(
-            blocks.iter().rev().map(|block| (block.hashed_state(), block.trie_updates())),
-        );
+        // Rebuild the overlay directly in sorted form, reusing the prefix-set allocation we already
+        // paid for on previous iterations.
+        let mut sorted_input = TrieInputSorted {
+            prefix_sets: core::mem::take(&mut input.prefix_sets),
+            ..Default::default()
+        };
+
+        for block in blocks.iter().rev() {
+            sorted_input.state.extend_ref(block.hashed_state());
+            sorted_input.nodes.extend_ref(block.trie_updates());
+        }
 
-        Ok((input, block_number))
+        Ok((sorted_input, input, block_number))
     }
 }
 

crates/trie/common/src/hashed_state.rs

@@ -507,6 +507,13 @@ impl HashedPostStateSorted {
                 .or_insert_with(|| other_storage.clone());
         }
     }
+
+    /// Clears all accounts and storage data.
+    pub fn clear(&mut self) {
+        self.accounts.accounts.clear();
+        self.accounts.destroyed_accounts.clear();
+        self.storages.clear();
+    }
 }
 
 impl AsRef<Self> for HashedPostStateSorted {

crates/trie/common/src/input.rs

@@ -1,4 +1,8 @@
-use crate::{prefix_set::TriePrefixSetsMut, updates::TrieUpdates, HashedPostState};
+use crate::{
+    prefix_set::TriePrefixSetsMut,
+    updates::{TrieUpdates, TrieUpdatesSorted},
+    HashedPostState, HashedPostStateSorted,
+};
 
 /// Inputs for trie-related computations.
 #[derive(Default, Debug, Clone)]
@@ -118,4 +122,75 @@ impl TrieInput {
         self.clear();
         self
     }
+
+    /// Converts trie input into [`TrieInputSorted`], draining the internal maps while keeping
+    /// their allocated capacity.
+    ///
+    /// This effectively clears all the fields in the [`TrieInput`] while preserving the
+    /// `HashMap` capacity for reuse. This allows us to reuse the allocated space for subsequent
+    /// block validations, avoiding repeated allocations and rehashing in hot paths.
+    ///
+    /// The sorted output allocates new `Vec` space, but the original `HashMap` capacity is
+    /// retained for the next cycle.
+    pub fn drain_into_sorted(&mut self) -> TrieInputSorted {
+        TrieInputSorted {
+            nodes: self.nodes.drain_into_sorted(),
+            state: self.state.drain_into_sorted(),
+            prefix_sets: core::mem::take(&mut self.prefix_sets),
+        }
+    }
+}
+
+/// Sorted variant of [`TrieInput`] for efficient proof generation.
+///
+/// This type holds sorted versions of trie data structures, which eliminates the need
+/// for expensive sorting operations during multiproof generation.
+#[derive(Default, Debug, Clone)]
+pub struct TrieInputSorted {
+    /// Sorted cached in-memory intermediate trie nodes.
+    pub nodes: TrieUpdatesSorted,
+    /// Sorted in-memory overlay hashed state.
+    pub state: HashedPostStateSorted,
+    /// Prefix sets for computation.
+    pub prefix_sets: TriePrefixSetsMut,
+}
+
+impl TrieInputSorted {
+    /// Create new sorted trie input.
+    pub const fn new(
+        nodes: TrieUpdatesSorted,
+        state: HashedPostStateSorted,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> Self {
+        Self { nodes, state, prefix_sets }
+    }
+
+    /// Create from unsorted [`TrieInput`] by sorting.
+    pub fn from_unsorted(input: TrieInput) -> Self {
+        Self {
+            nodes: input.nodes.into_sorted(),
+            state: input.state.into_sorted(),
+            prefix_sets: input.prefix_sets,
+        }
+    }
+
+    /// Append state to the input by reference and extend the prefix sets.
+    pub fn append_ref(&mut self, state: &HashedPostState) {
+        self.prefix_sets.extend(state.construct_prefix_sets());
+        let sorted_state = state.clone().into_sorted();
+        self.state.extend_ref(&sorted_state);
+    }
+
+    /// Clear all data.
+    pub fn clear(&mut self) {
+        self.nodes.clear();
+        self.state.clear();
+        self.prefix_sets.clear();
+    }
+
+    /// Return a cleared version of this sorted trie input.
+    pub fn cleared(mut self) -> Self {
+        self.clear();
+        self
+    }
 }

crates/trie/common/src/lib.rs

@@ -17,7 +17,7 @@ pub use hashed_state::*;
 
 /// Input for trie computation.
 mod input;
-pub use input::TrieInput;
+pub use input::{TrieInput, TrieInputSorted};
 
 /// The implementation of hash builder.
 pub mod hash_builder;

crates/trie/common/src/updates.rs

@@ -499,6 +499,12 @@ impl TrieUpdatesSorted {
                 .or_insert_with(|| storage_trie.clone());
         }
     }
+
+    /// Clears all account nodes and storage tries.
+    pub fn clear(&mut self) {
+        self.account_nodes.clear();
+        self.storage_tries.clear();
+    }
 }
 
 impl AsRef<Self> for TrieUpdatesSorted {