feat: trieUpdates and hashedState added to notification

crates/chain-state/src/in_memory.rs

@@ -822,6 +822,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()),
                     );
                     chain
                 }));
@@ -832,13 +834,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()),
                     );
                     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()),
                     );
                     chain
                 }));
@@ -1424,13 +1430,29 @@ mod tests {
         // Test commit notification
         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();
+        expected_trie_updates
+            .insert(0, Arc::new(TrieUpdates::from((*block0.trie_updates).clone())));
+        expected_trie_updates
+            .insert(1, Arc::new(TrieUpdates::from((*block1.trie_updates).clone())));
+
+        // Build expected hashed state map
+        let mut expected_hashed_state: BTreeMap<BlockNumber, Arc<HashedPostState>> =
+            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())));
+
         assert_eq!(
             chain_commit.to_chain_notification(),
             CanonStateNotification::Commit {
                 new: Arc::new(Chain::new(
                     vec![block0.recovered_block().clone(), block1.recovered_block().clone()],
                     sample_execution_outcome.clone(),
-                    None
+                    expected_trie_updates,
+                    expected_hashed_state
                 ))
             }
         );
@@ -1441,18 +1463,39 @@ mod tests {
             old: vec![block1.clone(), block2.clone()],
         };
 
+        // 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())));
+
+        // 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())));
+        // 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())));
+        // 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())));
         assert_eq!(
             chain_reorg.to_chain_notification(),
             CanonStateNotification::Reorg {
                 old: Arc::new(Chain::new(
                     vec![block1.recovered_block().clone(), block2.recovered_block().clone()],
                     sample_execution_outcome.clone(),
-                    None
+                    old_trie_updates,
+                    old_hashed_state
                 )),
                 new: Arc::new(Chain::new(
                     vec![block1a.recovered_block().clone(), block2a.recovered_block().clone()],
                     sample_execution_outcome,
-                    None
+                    new_trie_updates,
+                    new_hashed_state
                 ))
             }
         );

crates/chain-state/src/notifications.rs

@@ -242,6 +242,7 @@ mod tests {
     use reth_ethereum_primitives::{Receipt, TransactionSigned, TxType};
     use reth_execution_types::ExecutionOutcome;
     use reth_primitives_traits::SealedBlock;
+    use std::collections::BTreeMap;
 
     #[test]
     fn test_commit_notification() {
@@ -260,7 +261,8 @@ mod tests {
         let chain: Arc<Chain> = Arc::new(Chain::new(
             vec![block1.clone(), block2.clone()],
             ExecutionOutcome::default(),
-            None,
+            BTreeMap::new(),
+            BTreeMap::new(),
         ));
 
         // Create a commit notification
@@ -295,12 +297,17 @@ mod tests {
         block3.set_block_number(3);
         block3.set_hash(block3_hash);
 
-        let old_chain: Arc<Chain> =
-            Arc::new(Chain::new(vec![block1.clone()], ExecutionOutcome::default(), None));
+        let old_chain: Arc<Chain> = Arc::new(Chain::new(
+            vec![block1.clone()],
+            ExecutionOutcome::default(),
+            BTreeMap::new(),
+            BTreeMap::new(),
+        ));
         let new_chain = Arc::new(Chain::new(
             vec![block2.clone(), block3.clone()],
             ExecutionOutcome::default(),
-            None,
+            BTreeMap::new(),
+            BTreeMap::new(),
         ));
 
         // Create a reorg notification
@@ -362,8 +369,12 @@ mod tests {
         let execution_outcome = ExecutionOutcome { receipts, ..Default::default() };
 
         // Create a new chain segment with `block1` and `block2` and the execution outcome.
-        let new_chain: Arc<Chain> =
-            Arc::new(Chain::new(vec![block1.clone(), block2.clone()], execution_outcome, None));
+        let new_chain: Arc<Chain> = Arc::new(Chain::new(
+            vec![block1.clone(), block2.clone()],
+            execution_outcome,
+            BTreeMap::new(),
+            BTreeMap::new(),
+        ));
 
         // Create a commit notification containing the new chain segment.
         let notification = CanonStateNotification::Commit { new: new_chain };
@@ -420,8 +431,12 @@ mod tests {
             ExecutionOutcome { receipts: old_receipts, ..Default::default() };
 
         // Create an old chain segment to be reverted, containing `old_block1`.
-        let old_chain: Arc<Chain> =
-            Arc::new(Chain::new(vec![old_block1.clone()], old_execution_outcome, None));
+        let old_chain: Arc<Chain> = Arc::new(Chain::new(
+            vec![old_block1.clone()],
+            old_execution_outcome,
+            BTreeMap::new(),
+            BTreeMap::new(),
+        ));
 
         // Define block2 for the new chain segment, which will be committed.
         let mut body = BlockBody::<TransactionSigned>::default();
@@ -449,7 +464,12 @@ mod tests {
             ExecutionOutcome { receipts: new_receipts, ..Default::default() };
 
         // Create a new chain segment to be committed, containing `new_block1`.
-        let new_chain = Arc::new(Chain::new(vec![new_block1.clone()], new_execution_outcome, None));
+        let new_chain = Arc::new(Chain::new(
+            vec![new_block1.clone()],
+            new_execution_outcome,
+            BTreeMap::new(),
+            BTreeMap::new(),
+        ));
 
         // Create a reorg notification with both reverted (old) and committed (new) chain segments.
         let notification = CanonStateNotification::Reorg { old: old_chain, new: new_chain };