chore(trie): add e2e for reorg

crates/chain-state/src/in_memory.rs

@@ -919,13 +919,25 @@ impl<N: NodePrimitives<SignedTx: SignedTransaction>> NewCanonicalChain<N> {
     pub fn to_chain_notification(&self) -> CanonStateNotification<N> {
         match self {
             Self::Commit { new } => {
-                let new = Arc::new(new.iter().fold(Chain::default(), |mut chain, exec| {
-                    chain.append_block(
-                        exec.recovered_block().clone(),
-                        exec.execution_outcome().clone(),
-                    );
-                    chain
-                }));
+                // For single-block chains, preserve trie updates for testing
+                // This maintains backwards compatibility while allowing tests to access trie
+                // updates
+                let new = if new.len() == 1 {
+                    let block = new.first().unwrap();
+                    Arc::new(Chain::from_block(
+                        block.recovered_block().clone(),
+                        block.execution_outcome().clone(),
+                        block.trie_updates().cloned(),
+                    ))
+                } else {
+                    Arc::new(new.iter().fold(Chain::default(), |mut chain, exec| {
+                        chain.append_block(
+                            exec.recovered_block().clone(),
+                            exec.execution_outcome().clone(),
+                        );
+                        chain
+                    }))
+                };
                 CanonStateNotification::Commit { new }
             }
             Self::Reorg { new, old } => {

crates/e2e-test-utils/Cargo.toml

@@ -22,7 +22,7 @@ reth-rpc-eth-api.workspace = true
 reth-rpc-api = { workspace = true, features = ["client"] }
 reth-payload-builder = { workspace = true, features = ["test-utils"] }
 reth-payload-builder-primitives.workspace = true
-reth-payload-primitives.workspace = true
+reth-payload-primitives = { workspace = true, features = ["op"] }
 reth-provider.workspace = true
 reth-node-api.workspace = true
 reth-node-core.workspace = true
@@ -41,6 +41,7 @@ reth-consensus.workspace = true
 reth-primitives.workspace = true
 reth-db-common.workspace = true
 reth-primitives-traits.workspace = true
+reth-trie-common.workspace = true
 
 revm.workspace = true
 tempfile.workspace = true
@@ -71,4 +72,4 @@ derive_more.workspace = true
 
 [[test]]
 name = "e2e_testsuite"
-path = "tests/e2e-testsuite/main.rs"
+path = "tests/e2e_testsuite/main.rs"

crates/e2e-test-utils/src/node.rs

@@ -1,4 +1,7 @@
-use crate::{network::NetworkTestContext, payload::PayloadTestContext, rpc::RpcTestContext};
+use crate::{
+    network::NetworkTestContext, payload::PayloadTestContext, rpc::RpcTestContext,
+    testsuite::TrieUpdateEvent,
+};
 use alloy_consensus::{transaction::TxHashRef, BlockHeader};
 use alloy_eips::BlockId;
 use alloy_primitives::{BlockHash, BlockNumber, Bytes, Sealable, B256};
@@ -17,13 +20,14 @@ use reth_node_builder::{rpc::RethRpcAddOns, FullNode, NodeTypes};
 
 use reth_payload_primitives::{BuiltPayload, PayloadBuilderAttributes};
 use reth_provider::{
-    BlockReader, BlockReaderIdExt, CanonStateNotificationStream, CanonStateSubscriptions,
-    HeaderProvider, StageCheckpointReader,
+    BlockReader, BlockReaderIdExt, CanonStateNotification, CanonStateNotificationStream,
+    CanonStateSubscriptions, HeaderProvider, StageCheckpointReader,
 };
 use reth_rpc_builder::auth::AuthServerHandle;
 use reth_rpc_eth_api::helpers::{EthApiSpec, EthTransactions, TraceExt};
 use reth_stages_types::StageId;
 use std::pin::Pin;
+use tokio::sync::mpsc;
 use tokio_stream::StreamExt;
 use url::Url;
 
@@ -304,6 +308,101 @@ where
         self.inner.auth_server_handle().clone()
     }
 
+    /// Start a background task that forwards trie updates from canonical state notifications
+    /// to the test environment via the provided channel
+    pub fn start_trie_update_forwarder(
+        &mut self,
+        node_idx: usize,
+        tx: mpsc::UnboundedSender<TrieUpdateEvent>,
+    ) where
+        AddOns: RethRpcAddOns<Node> + 'static,
+    {
+        // Take the canonical stream (we can't clone it, so we move it into the task)
+        // This means each node can only have one forwarder, which is fine for our use case
+        let canonical_stream = std::mem::replace(
+            &mut self.canonical_stream,
+            self.inner.provider.canonical_state_stream(),
+        );
+
+        tokio::spawn(async move {
+            println!("Starting trie update forwarder for node {}", node_idx);
+            tracing::debug!("Starting trie update forwarder for node {}", node_idx);
+
+            let mut stream = canonical_stream;
+            while let Some(notification) = stream.next().await {
+                println!(
+                    "Node {} received canonical state notification: {:?}",
+                    node_idx,
+                    std::mem::discriminant(&notification)
+                );
+                match notification {
+                    CanonStateNotification::Commit { new } => {
+                        // Extract trie updates from committed chain
+                        if let Some(trie_updates) = new.trie_updates() {
+                            // Get the tip block hash from the chain
+                            if let Some(tip_header) = new.headers().last() {
+                                let tip_hash: alloy_primitives::B256 = tip_header.hash();
+
+                                let event = TrieUpdateEvent {
+                                    node_idx,
+                                    block_hash: tip_hash,
+                                    trie_updates: trie_updates.clone(),
+                                };
+
+                                // Send the event, ignore errors if receiver is dropped
+                                let _ = tx.send(event);
+
+                                tracing::debug!(
+                                    "Forwarded trie updates for committed block {} from node {}",
+                                    tip_hash,
+                                    node_idx
+                                );
+                            }
+                        } else {
+                            println!("Node {}: No trie updates in committed chain", node_idx);
+                            tracing::debug!(
+                                "No trie updates in committed chain for node {}",
+                                node_idx
+                            );
+                        }
+                    }
+                    CanonStateNotification::Reorg { old: _old, new } => {
+                        // Extract trie updates from the new canonical chain after reorg
+                        if let Some(trie_updates) = new.trie_updates() {
+                            // Get all blocks in the new chain and forward their trie updates
+                            for header in new.headers() {
+                                let block_hash: alloy_primitives::B256 = header.hash();
+
+                                let event = TrieUpdateEvent {
+                                    node_idx,
+                                    block_hash,
+                                    trie_updates: trie_updates.clone(),
+                                };
+
+                                // Send the event, ignore errors if receiver is dropped
+                                let _ = tx.send(event);
+
+                                tracing::debug!(
+                                    "Forwarded trie updates for reorged block {} from node {}",
+                                    block_hash,
+                                    node_idx
+                                );
+                            }
+                        } else {
+                            println!("Node {}: No trie updates in reorged chain", node_idx);
+                            tracing::debug!(
+                                "No trie updates in reorged chain for node {}",
+                                node_idx
+                            );
+                        }
+                    }
+                }
+            }
+
+            tracing::debug!("Trie update forwarder ended for node {}", node_idx);
+        });
+    }
+
     /// Creates a [`crate::testsuite::NodeClient`] from this test context.
     ///
     /// This helper method extracts the necessary handles and creates a client