chore(chain-state): Receive sorted updates to exex

crates/chain-state/src/in_memory.rs

@@ -927,8 +927,8 @@ impl<N: NodePrimitives<SignedTx: SignedTransaction>> NewCanonicalChain<N> {
                     chain.append_block(
                         exec.recovered_block().clone(),
                         exec.execution_outcome().clone(),
-                        Arc::new((*exec.trie_updates()).clone().into()),
-                        Arc::new((*exec.hashed_state()).clone().into()),
+                        exec.trie_updates(),
+                        exec.hashed_state(),
                     );
                     chain
                 }));
@@ -939,17 +939,17 @@ impl<N: NodePrimitives<SignedTx: SignedTransaction>> NewCanonicalChain<N> {
                     chain.append_block(
                         exec.recovered_block().clone(),
                         exec.execution_outcome().clone(),
-                        Arc::new((*exec.trie_updates()).clone().into()),
-                        Arc::new((*exec.hashed_state()).clone().into()),
+                        exec.trie_updates(),
+                        exec.hashed_state(),
                     );
                     chain
                 }));
                 let old = Arc::new(old.iter().fold(Chain::default(), |mut chain, exec| {
                     chain.append_block(
                         exec.recovered_block().clone(),
                         exec.execution_outcome().clone(),
-                        Arc::new((*exec.trie_updates()).clone().into()),
-                        Arc::new((*exec.hashed_state()).clone().into()),
+                        exec.trie_updates(),
+                        exec.hashed_state(),
                     );
                     chain
                 }));
@@ -1536,19 +1536,24 @@ mod tests {
         let chain_commit = NewCanonicalChain::Commit { new: vec![block0.clone(), block1.clone()] };
 
         // Build expected trie updates map
-        let mut expected_trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdates>> = BTreeMap::new();
+        let mut expected_trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>> =
+            BTreeMap::new();
         expected_trie_updates
-            .insert(0, Arc::new(TrieUpdates::from((*block0.trie_updates()).clone())));
+            .insert(0, Arc::new(TrieUpdates::from((*block0.trie_updates()).clone()).into_sorted()));
         expected_trie_updates
-            .insert(1, Arc::new(TrieUpdates::from((*block1.trie_updates()).clone())));
+            .insert(1, Arc::new(TrieUpdates::from((*block1.trie_updates()).clone()).into_sorted()));
 
         // Build expected hashed state map
-        let mut expected_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostState>> =
+        let mut expected_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostStateSorted>> =
             BTreeMap::new();
-        expected_hashed_state
-            .insert(0, Arc::new(HashedPostState::from((*block0.hashed_state()).clone())));
-        expected_hashed_state
-            .insert(1, Arc::new(HashedPostState::from((*block1.hashed_state()).clone())));
+        expected_hashed_state.insert(
+            0,
+            Arc::new(HashedPostState::from((*block0.hashed_state()).clone()).into_sorted()),
+        );
+        expected_hashed_state.insert(
+            1,
+            Arc::new(HashedPostState::from((*block1.hashed_state()).clone()).into_sorted()),
+        );
 
         assert_eq!(
             chain_commit.to_chain_notification(),
@@ -1569,26 +1574,44 @@ mod tests {
         };
 
         // Build expected trie updates for old chain
-        let mut old_trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdates>> = BTreeMap::new();
-        old_trie_updates.insert(1, Arc::new(TrieUpdates::from((*block1.trie_updates()).clone())));
-        old_trie_updates.insert(2, Arc::new(TrieUpdates::from((*block2.trie_updates()).clone())));
+        let mut old_trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>> = BTreeMap::new();
+        old_trie_updates
+            .insert(1, Arc::new(TrieUpdates::from((*block1.trie_updates()).clone()).into_sorted()));
+        old_trie_updates
+            .insert(2, Arc::new(TrieUpdates::from((*block2.trie_updates()).clone()).into_sorted()));
 
         // Build expected trie updates for new chain
-        let mut new_trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdates>> = BTreeMap::new();
-        new_trie_updates.insert(1, Arc::new(TrieUpdates::from((*block1a.trie_updates()).clone())));
-        new_trie_updates.insert(2, Arc::new(TrieUpdates::from((*block2a.trie_updates()).clone())));
+        let mut new_trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>> = BTreeMap::new();
+        new_trie_updates.insert(
+            1,
+            Arc::new(TrieUpdates::from((*block1a.trie_updates()).clone()).into_sorted()),
+        );
+        new_trie_updates.insert(
+            2,
+            Arc::new(TrieUpdates::from((*block2a.trie_updates()).clone()).into_sorted()),
+        );
         // Build expected hashed state for old chain
-        let mut old_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostState>> = BTreeMap::new();
-        old_hashed_state
-            .insert(1, Arc::new(HashedPostState::from((*block1.hashed_state()).clone())));
-        old_hashed_state
-            .insert(2, Arc::new(HashedPostState::from((*block2.hashed_state()).clone())));
+        let mut old_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostStateSorted>> =
+            BTreeMap::new();
+        old_hashed_state.insert(
+            1,
+            Arc::new(HashedPostState::from((*block1.hashed_state()).clone()).into_sorted()),
+        );
+        old_hashed_state.insert(
+            2,
+            Arc::new(HashedPostState::from((*block2.hashed_state()).clone()).into_sorted()),
+        );
         // Build expected hashed state for new chain
-        let mut new_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostState>> = BTreeMap::new();
-        new_hashed_state
-            .insert(1, Arc::new(HashedPostState::from((*block1a.hashed_state()).clone())));
-        new_hashed_state
-            .insert(2, Arc::new(HashedPostState::from((*block2a.hashed_state()).clone())));
+        let mut new_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostStateSorted>> =
+            BTreeMap::new();
+        new_hashed_state.insert(
+            1,
+            Arc::new(HashedPostState::from((*block1a.hashed_state()).clone()).into_sorted()),
+        );
+        new_hashed_state.insert(
+            2,
+            Arc::new(HashedPostState::from((*block2a.hashed_state()).clone()).into_sorted()),
+        );
         assert_eq!(
             chain_reorg.to_chain_notification(),
             CanonStateNotification::Reorg {

crates/evm/execution-types/src/chain.rs

@@ -10,7 +10,7 @@ use reth_primitives_traits::{
     transaction::signed::SignedTransaction, Block, BlockBody, NodePrimitives, RecoveredBlock,
     SealedHeader,
 };
-use reth_trie_common::{updates::TrieUpdates, HashedPostState};
+use reth_trie_common::{updates::TrieUpdatesSorted, HashedPostStateSorted};
 
 /// A chain of blocks and their final state.
 ///
@@ -35,9 +35,9 @@ pub struct Chain<N: NodePrimitives = reth_ethereum_primitives::EthPrimitives> {
     /// Additionally, it includes the individual state changes that led to the current state.
     execution_outcome: ExecutionOutcome<N::Receipt>,
     /// State trie updates for each block in the chain, keyed by block number.
-    trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdates>>,
+    trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>>,
     /// Hashed post state for each block in the chain, keyed by block number.
-    hashed_state: BTreeMap<BlockNumber, Arc<HashedPostState>>,
+    hashed_state: BTreeMap<BlockNumber, Arc<HashedPostStateSorted>>,
 }
 
 impl<N: NodePrimitives> Default for Chain<N> {
@@ -60,8 +60,8 @@ impl<N: NodePrimitives> Chain<N> {
     pub fn new(
         blocks: impl IntoIterator<Item = RecoveredBlock<N::Block>>,
         execution_outcome: ExecutionOutcome<N::Receipt>,
-        trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdates>>,
-        hashed_state: BTreeMap<BlockNumber, Arc<HashedPostState>>,
+        trie_updates: BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>>,
+        hashed_state: BTreeMap<BlockNumber, Arc<HashedPostStateSorted>>,
     ) -> Self {
         let blocks =
             blocks.into_iter().map(|b| (b.header().number(), b)).collect::<BTreeMap<_, _>>();
@@ -74,8 +74,8 @@ impl<N: NodePrimitives> Chain<N> {
     pub fn from_block(
         block: RecoveredBlock<N::Block>,
         execution_outcome: ExecutionOutcome<N::Receipt>,
-        trie_updates: Arc<TrieUpdates>,
-        hashed_state: Arc<HashedPostState>,
+        trie_updates: Arc<TrieUpdatesSorted>,
+        hashed_state: Arc<HashedPostStateSorted>,
     ) -> Self {
         let block_number = block.header().number();
         let trie_updates_map = BTreeMap::from([(block_number, trie_updates)]);
@@ -99,12 +99,12 @@ impl<N: NodePrimitives> Chain<N> {
     }
 
     /// Get all trie updates for this chain.
-    pub const fn trie_updates(&self) -> &BTreeMap<BlockNumber, Arc<TrieUpdates>> {
+    pub const fn trie_updates(&self) -> &BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>> {
         &self.trie_updates
     }
 
     /// Get trie updates for a specific block number.
-    pub fn trie_updates_at(&self, block_number: BlockNumber) -> Option<&Arc<TrieUpdates>> {
+    pub fn trie_updates_at(&self, block_number: BlockNumber) -> Option<&Arc<TrieUpdatesSorted>> {
         self.trie_updates.get(&block_number)
     }
 
@@ -114,12 +114,15 @@ impl<N: NodePrimitives> Chain<N> {
     }
 
     /// Get all hashed states for this chain.
-    pub const fn hashed_state(&self) -> &BTreeMap<BlockNumber, Arc<HashedPostState>> {
+    pub const fn hashed_state(&self) -> &BTreeMap<BlockNumber, Arc<HashedPostStateSorted>> {
         &self.hashed_state
     }
 
     /// Get hashed state for a specific block number.
-    pub fn hashed_state_at(&self, block_number: BlockNumber) -> Option<&Arc<HashedPostState>> {
+    pub fn hashed_state_at(
+        &self,
+        block_number: BlockNumber,
+    ) -> Option<&Arc<HashedPostStateSorted>> {
         self.hashed_state.get(&block_number)
     }
 
@@ -181,8 +184,8 @@ impl<N: NodePrimitives> Chain<N> {
     ) -> (
         ChainBlocks<'static, N::Block>,
         ExecutionOutcome<N::Receipt>,
-        BTreeMap<BlockNumber, Arc<TrieUpdates>>,
-        BTreeMap<BlockNumber, Arc<HashedPostState>>,
+        BTreeMap<BlockNumber, Arc<TrieUpdatesSorted>>,
+        BTreeMap<BlockNumber, Arc<HashedPostStateSorted>>,
     ) {
         (
             ChainBlocks { blocks: Cow::Owned(self.blocks) },
@@ -301,8 +304,8 @@ impl<N: NodePrimitives> Chain<N> {
         &mut self,
         block: RecoveredBlock<N::Block>,
         execution_outcome: ExecutionOutcome<N::Receipt>,
-        trie_updates: Arc<TrieUpdates>,
-        hashed_state: Arc<HashedPostState>,
+        trie_updates: Arc<TrieUpdatesSorted>,
+        hashed_state: Arc<HashedPostStateSorted>,
     ) {
         let block_number = block.header().number();
         self.blocks.insert(block_number, block);
@@ -491,9 +494,9 @@ pub(super) mod serde_bincode_compat {
         #[serde(default, rename = "trie_updates_legacy")]
         _trie_updates_legacy: Option<TrieUpdates<'a>>,
         #[serde(default)]
-        trie_updates: BTreeMap<BlockNumber, Arc<super::TrieUpdates>>,
+        trie_updates: BTreeMap<BlockNumber, Arc<super::TrieUpdatesSorted>>,
         #[serde(default)]
-        hashed_state: BTreeMap<BlockNumber, Arc<super::HashedPostState>>,
+        hashed_state: BTreeMap<BlockNumber, Arc<super::HashedPostStateSorted>>,
     }
 
     #[derive(Debug)]

crates/optimism/exex/src/lib.rs

@@ -20,7 +20,7 @@ use reth_optimism_trie::{
     live::LiveTrieCollector, OpProofStoragePrunerTask, OpProofsStorage, OpProofsStore,
 };
 use reth_provider::{BlockReader, TransactionVariant};
-use reth_trie::{updates::TrieUpdates, HashedPostState};
+use reth_trie::{updates::TrieUpdatesSorted, HashedPostStateSorted};
 use std::{sync::Arc, time::Duration};
 use tracing::{debug, info};
 
@@ -253,8 +253,8 @@ where
                 collector
                     .store_block_updates(
                         block.block_with_parent(),
-                        trie_updates.clone(),
-                        hashed_state.clone(),
+                        (**trie_updates).clone(),
+                        (**hashed_state).clone(),
                     )
                     .await?;
 
@@ -306,8 +306,11 @@ where
         }
 
         // find the common ancestor
-        let mut block_updates: Vec<(BlockWithParent, Arc<TrieUpdates>, Arc<HashedPostState>)> =
-            Vec::with_capacity(new.len());
+        let mut block_updates: Vec<(
+            BlockWithParent,
+            Arc<TrieUpdatesSorted>,
+            Arc<HashedPostStateSorted>,
+        )> = Vec::with_capacity(new.len());
         for block_number in new.blocks().keys() {
             // verify if the fork point matches
             if old.fork_block() != new.fork_block() {

crates/optimism/trie/Cargo.toml

@@ -45,7 +45,6 @@ eyre.workspace = true
 strum.workspace = true
 tracing.workspace = true
 derive_more.workspace = true
-itertools.workspace = true
 
 [dev-dependencies]
 reth-codecs = { workspace = true, features = ["test-utils"] }

crates/optimism/trie/src/api.rs

@@ -8,25 +8,25 @@ use reth_primitives_traits::Account;
 use reth_trie::{
     hashed_cursor::{HashedCursor, HashedStorageCursor},
     trie_cursor::{TrieCursor, TrieStorageCursor},
-    updates::TrieUpdates,
-    BranchNodeCompact, HashedPostState, Nibbles,
+    updates::TrieUpdatesSorted,
+    BranchNodeCompact, HashedPostStateSorted, Nibbles,
 };
 use std::{fmt::Debug, time::Duration};
 
 /// Diff of trie updates and post state for a block.
 #[derive(Debug, Clone, Default)]
 pub struct BlockStateDiff {
     /// Trie updates for branch nodes
-    pub trie_updates: TrieUpdates,
+    pub sorted_trie_updates: TrieUpdatesSorted,
     /// Post state for leaf nodes (accounts and storage)
-    pub post_state: HashedPostState,
+    pub sorted_post_state: HashedPostStateSorted,
 }
 
 impl BlockStateDiff {
     /// Extend the [` BlockStateDiff`] from other latest [`BlockStateDiff`]
-    pub fn extend(&mut self, other: Self) {
-        self.trie_updates.extend(other.trie_updates);
-        self.post_state.extend(other.post_state);
+    pub fn extend_ref(&mut self, other: &Self) {
+        self.sorted_trie_updates.extend_ref(&other.sorted_trie_updates);
+        self.sorted_post_state.extend_ref(&other.sorted_post_state);
     }
 }