perf(trie): add HashedPostStateSorted::from_reverts (#20047)

Author: yongkangc

crates/chain-state/src/deferred_trie.rs

@@ -155,7 +155,7 @@ impl DeferredTrieData {
         ancestors: &[Self],
     ) -> ComputedTrieData {
         // Sort the current block's hashed state and trie updates
-        let sorted_hashed_state = Arc::new(hashed_state.clone().into_sorted());
+        let sorted_hashed_state = Arc::new(hashed_state.clone_into_sorted());
         let sorted_trie_updates = Arc::new(trie_updates.clone().into_sorted());
 
         // Merge trie data from ancestors (oldest -> newest so later state takes precedence)

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

@@ -647,7 +647,7 @@ where
 
         // 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();
+        let sorted_hashed_state = hashed_state.clone_into_sorted();
         Arc::make_mut(&mut input.state).extend_ref(&sorted_hashed_state);
 
         let TrieInputSorted { nodes, state, prefix_sets: prefix_sets_mut } = input;

crates/stages/stages/src/stages/merkle_changesets.rs

@@ -12,9 +12,11 @@ use reth_stages_api::{
     BlockErrorKind, ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId,
     UnwindInput, UnwindOutput,
 };
-use reth_trie::{updates::TrieUpdates, HashedPostState, KeccakKeyHasher, StateRoot, TrieInput};
+use reth_trie::{
+    updates::TrieUpdates, HashedPostStateSorted, KeccakKeyHasher, StateRoot, TrieInputSorted,
+};
 use reth_trie_db::{DatabaseHashedPostState, DatabaseStateRoot};
-use std::ops::Range;
+use std::{ops::Range, sync::Arc};
 use tracing::{debug, error};
 
 /// The `MerkleChangeSets` stage.
@@ -105,12 +107,12 @@ impl MerkleChangeSets {
         Ok(target_start..target_end)
     }
 
-    /// Calculates the trie updates given a [`TrieInput`], asserting that the resulting state root
-    /// matches the expected one for the block.
+    /// Calculates the trie updates given a [`TrieInputSorted`], asserting that the resulting state
+    /// root matches the expected one for the block.
     fn calculate_block_trie_updates<Provider: DBProvider + HeaderProvider>(
         provider: &Provider,
         block_number: BlockNumber,
-        input: TrieInput,
+        input: TrieInputSorted,
     ) -> Result<TrieUpdates, StageError> {
         let (root, trie_updates) =
             StateRoot::overlay_root_from_nodes_with_updates(provider.tx_ref(), input).map_err(
@@ -192,21 +194,21 @@ impl MerkleChangeSets {
         );
         let mut per_block_state_reverts = Vec::new();
         for block_number in target_range.clone() {
-            per_block_state_reverts.push(HashedPostState::from_reverts::<KeccakKeyHasher>(
+            per_block_state_reverts.push(HashedPostStateSorted::from_reverts::<KeccakKeyHasher>(
                 provider.tx_ref(),
                 block_number..=block_number,
             )?);
         }
 
         // Helper to retrieve state revert data for a specific block from the pre-computed array
-        let get_block_state_revert = |block_number: BlockNumber| -> &HashedPostState {
+        let get_block_state_revert = |block_number: BlockNumber| -> &HashedPostStateSorted {
             let index = (block_number - target_start) as usize;
             &per_block_state_reverts[index]
         };
 
         // Helper to accumulate state reverts from a given block to the target end
-        let compute_cumulative_state_revert = |block_number: BlockNumber| -> HashedPostState {
-            let mut cumulative_revert = HashedPostState::default();
+        let compute_cumulative_state_revert = |block_number: BlockNumber| -> HashedPostStateSorted {
+            let mut cumulative_revert = HashedPostStateSorted::default();
             for n in (block_number..target_end).rev() {
                 cumulative_revert.extend_ref(get_block_state_revert(n))
             }
@@ -216,7 +218,7 @@ impl MerkleChangeSets {
         // To calculate the changeset for a block, we first need the TrieUpdates which are
         // generated as a result of processing the block. To get these we need:
         // 1) The TrieUpdates which revert the db's trie to _prior_ to the block
-        // 2) The HashedPostState to revert the db's state to _after_ the block
+        // 2) The HashedPostStateSorted to revert the db's state to _after_ the block
         //
         // To get (1) for `target_start` we need to do a big state root calculation which takes
         // into account all changes between that block and db tip. For each block after the
@@ -227,12 +229,15 @@ impl MerkleChangeSets {
             ?target_start,
             "Computing trie state at starting block",
         );
-        let mut input = TrieInput::default();
-        input.state = compute_cumulative_state_revert(target_start);
-        input.prefix_sets = input.state.construct_prefix_sets();
+        let initial_state = compute_cumulative_state_revert(target_start);
+        let initial_prefix_sets = initial_state.construct_prefix_sets();
+        let initial_input =
+            TrieInputSorted::new(Arc::default(), Arc::new(initial_state), initial_prefix_sets);
         // target_start will be >= 1, see `determine_target_range`.
-        input.nodes =
-            Self::calculate_block_trie_updates(provider, target_start - 1, input.clone())?;
+        let mut nodes = Arc::new(
+            Self::calculate_block_trie_updates(provider, target_start - 1, initial_input)?
+                .into_sorted(),
+        );
 
         for block_number in target_range {
             debug!(
@@ -242,21 +247,24 @@ impl MerkleChangeSets {
             );
             // Revert the state so that this block has been just processed, meaning we take the
             // cumulative revert of the subsequent block.
-            input.state = compute_cumulative_state_revert(block_number + 1);
+            let state = Arc::new(compute_cumulative_state_revert(block_number + 1));
+
+            // Construct prefix sets from only this block's `HashedPostStateSorted`, because we only
+            // care about trie updates which occurred as a result of this block being processed.
+            let prefix_sets = get_block_state_revert(block_number).construct_prefix_sets();
 
-            // Construct prefix sets from only this block's `HashedPostState`, because we only care
-            // about trie updates which occurred as a result of this block being processed.
-            input.prefix_sets = get_block_state_revert(block_number).construct_prefix_sets();
+            let input = TrieInputSorted::new(Arc::clone(&nodes), state, prefix_sets);
 
             // Calculate the trie updates for this block, then apply those updates to the reverts.
             // We calculate the overlay which will be passed into the next step using the trie
             // reverts prior to them being updated.
             let this_trie_updates =
-                Self::calculate_block_trie_updates(provider, block_number, input.clone())?;
+                Self::calculate_block_trie_updates(provider, block_number, input)?.into_sorted();
 
-            let trie_overlay = input.nodes.clone().into_sorted();
-            input.nodes.extend_ref(&this_trie_updates);
-            let this_trie_updates = this_trie_updates.into_sorted();
+            let trie_overlay = Arc::clone(&nodes);
+            let mut nodes_mut = Arc::unwrap_or_clone(nodes);
+            nodes_mut.extend_ref(&this_trie_updates);
+            nodes = Arc::new(nodes_mut);
 
             // Write the changesets to the DB using the trie updates produced by the block, and the
             // trie reverts as the overlay.

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

@@ -21,8 +21,9 @@ use reth_trie::{
     proof::{Proof, StorageProof},
     updates::TrieUpdates,
     witness::TrieWitness,
-    AccountProof, HashedPostState, HashedStorage, KeccakKeyHasher, MultiProof, MultiProofTargets,
-    StateRoot, StorageMultiProof, StorageRoot, TrieInput,
+    AccountProof, HashedPostState, HashedPostStateSorted, HashedStorage, KeccakKeyHasher,
+    MultiProof, MultiProofTargets, StateRoot, StorageMultiProof, StorageRoot, TrieInput,
+    TrieInputSorted,
 };
 use reth_trie_db::{
     DatabaseHashedPostState, DatabaseHashedStorage, DatabaseProof, DatabaseStateRoot,
@@ -118,7 +119,7 @@ impl<'b, Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'b, P
     }
 
     /// Retrieve revert hashed state for this history provider.
-    fn revert_state(&self) -> ProviderResult<HashedPostState> {
+    fn revert_state(&self) -> ProviderResult<HashedPostStateSorted> {
         if !self.lowest_available_blocks.is_account_history_available(self.block_number) ||
             !self.lowest_available_blocks.is_storage_history_available(self.block_number)
         {
@@ -133,7 +134,8 @@ impl<'b, Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'b, P
             );
         }
 
-        Ok(HashedPostState::from_reverts::<KeccakKeyHasher>(self.tx(), self.block_number..)?)
+        HashedPostStateSorted::from_reverts::<KeccakKeyHasher>(self.tx(), self.block_number..)
+            .map_err(ProviderError::from)
     }
 
     /// Retrieve revert hashed storage for this history provider and target address.
@@ -287,14 +289,15 @@ impl<Provider: DBProvider + BlockNumReader> StateRootProvider
 {
     fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
         let mut revert_state = self.revert_state()?;
-        revert_state.extend(hashed_state);
-        StateRoot::overlay_root(self.tx(), revert_state)
+        let hashed_state_sorted = hashed_state.into_sorted();
+        revert_state.extend_ref(&hashed_state_sorted);
+        StateRoot::overlay_root(self.tx(), &revert_state)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
     fn state_root_from_nodes(&self, mut input: TrieInput) -> ProviderResult<B256> {
-        input.prepend(self.revert_state()?);
-        StateRoot::overlay_root_from_nodes(self.tx(), input)
+        input.prepend(self.revert_state()?.into());
+        StateRoot::overlay_root_from_nodes(self.tx(), TrieInputSorted::from_unsorted(input))
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
@@ -303,18 +306,22 @@ impl<Provider: DBProvider + BlockNumReader> StateRootProvider
         hashed_state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         let mut revert_state = self.revert_state()?;
-        revert_state.extend(hashed_state);
-        StateRoot::overlay_root_with_updates(self.tx(), revert_state)
+        let hashed_state_sorted = hashed_state.into_sorted();
+        revert_state.extend_ref(&hashed_state_sorted);
+        StateRoot::overlay_root_with_updates(self.tx(), &revert_state)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
     fn state_root_from_nodes_with_updates(
         &self,
         mut input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        input.prepend(self.revert_state()?);
-        StateRoot::overlay_root_from_nodes_with_updates(self.tx(), input)
-            .map_err(|err| ProviderError::Database(err.into()))
+        input.prepend(self.revert_state()?.into());
+        StateRoot::overlay_root_from_nodes_with_updates(
+            self.tx(),
+            TrieInputSorted::from_unsorted(input),
+        )
+        .map_err(|err| ProviderError::Database(err.into()))
     }
 }
 
@@ -367,7 +374,7 @@ impl<Provider: DBProvider + BlockNumReader> StateProofProvider
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        input.prepend(self.revert_state()?);
+        input.prepend(self.revert_state()?.into());
         let proof = <Proof<_, _> as DatabaseProof>::from_tx(self.tx());
         proof.overlay_account_proof(input, address, slots).map_err(ProviderError::from)
     }
@@ -377,13 +384,13 @@ impl<Provider: DBProvider + BlockNumReader> StateProofProvider
         mut input: TrieInput,
         targets: MultiProofTargets,
     ) -> ProviderResult<MultiProof> {
-        input.prepend(self.revert_state()?);
+        input.prepend(self.revert_state()?.into());
         let proof = <Proof<_, _> as DatabaseProof>::from_tx(self.tx());
         proof.overlay_multiproof(input, targets).map_err(ProviderError::from)
     }
 
     fn witness(&self, mut input: TrieInput, target: HashedPostState) -> ProviderResult<Vec<Bytes>> {
-        input.prepend(self.revert_state()?);
+        input.prepend(self.revert_state()?.into());
         TrieWitness::overlay_witness(self.tx(), input, target)
             .map_err(ProviderError::from)
             .map(|hm| hm.into_values().collect())

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

@@ -12,7 +12,7 @@ use reth_trie::{
     updates::TrieUpdates,
     witness::TrieWitness,
     AccountProof, HashedPostState, HashedStorage, KeccakKeyHasher, MultiProof, MultiProofTargets,
-    StateRoot, StorageMultiProof, StorageRoot, TrieInput,
+    StateRoot, StorageMultiProof, StorageRoot, TrieInput, TrieInputSorted,
 };
 use reth_trie_db::{
     DatabaseProof, DatabaseStateRoot, DatabaseStorageProof, DatabaseStorageRoot,
@@ -60,29 +60,32 @@ impl<Provider: BlockHashReader> BlockHashReader for LatestStateProviderRef<'_, P
 
 impl<Provider: DBProvider + Sync> StateRootProvider for LatestStateProviderRef<'_, Provider> {
     fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
-        StateRoot::overlay_root(self.tx(), hashed_state)
+        StateRoot::overlay_root(self.tx(), &hashed_state.into_sorted())
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
     fn state_root_from_nodes(&self, input: TrieInput) -> ProviderResult<B256> {
-        StateRoot::overlay_root_from_nodes(self.tx(), input)
+        StateRoot::overlay_root_from_nodes(self.tx(), TrieInputSorted::from_unsorted(input))
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
     fn state_root_with_updates(
         &self,
         hashed_state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        StateRoot::overlay_root_with_updates(self.tx(), hashed_state)
+        StateRoot::overlay_root_with_updates(self.tx(), &hashed_state.into_sorted())
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
     fn state_root_from_nodes_with_updates(
         &self,
         input: TrieInput,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        StateRoot::overlay_root_from_nodes_with_updates(self.tx(), input)
-            .map_err(|err| ProviderError::Database(err.into()))
+        StateRoot::overlay_root_from_nodes_with_updates(
+            self.tx(),
+            TrieInputSorted::from_unsorted(input),
+        )
+        .map_err(|err| ProviderError::Database(err.into()))
     }
 }
 

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

@@ -14,7 +14,7 @@ use reth_trie::{
     hashed_cursor::{HashedCursorFactory, HashedPostStateCursorFactory},
     trie_cursor::{InMemoryTrieCursorFactory, TrieCursorFactory},
     updates::TrieUpdatesSorted,
-    HashedPostState, HashedPostStateSorted, KeccakKeyHasher,
+    HashedPostStateSorted, KeccakKeyHasher,
 };
 use reth_trie_db::{
     DatabaseHashedCursorFactory, DatabaseHashedPostState, DatabaseTrieCursorFactory,
@@ -234,13 +234,10 @@ where
                 let _guard = debug_span!(target: "providers::state::overlay", "Retrieving hashed state reverts").entered();
 
                 let start = Instant::now();
-                // TODO(mediocregopher) make from_reverts return sorted
-                // https://github.com/paradigmxyz/reth/issues/19382
-                let res = HashedPostState::from_reverts::<KeccakKeyHasher>(
+                let res = HashedPostStateSorted::from_reverts::<KeccakKeyHasher>(
                     provider.tx_ref(),
                     from_block + 1..,
-                )?
-                .into_sorted();
+                )?;
                 retrieve_hashed_state_reverts_duration = start.elapsed();
                 res
             };

crates/storage/provider/src/writer/mod.rs

@@ -917,7 +917,7 @@ mod tests {
             assert_eq!(
                 StateRoot::overlay_root(
                     tx,
-                    provider_factory.hashed_post_state(&state.bundle_state)
+                    &provider_factory.hashed_post_state(&state.bundle_state).into_sorted()
                 )
                 .unwrap(),
                 state_root(expected.clone().into_iter().map(|(address, (account, storage))| (