Add incremental trie cache optimization for flashblocks state root calculation

crates/op-rbuilder/src/builder/payload.rs

@@ -46,7 +46,7 @@ use reth_revm::{
     State, database::StateProviderDatabase, db::states::bundle_state::BundleRetention,
 };
 use reth_transaction_pool::TransactionPool;
-use reth_trie::{HashedPostState, updates::TrieUpdates};
+use reth_trie::{HashedPostState, TrieInput, updates::TrieUpdates};
 use revm::Database;
 use std::{collections::BTreeMap, ops::Deref, sync::Arc, time::Instant};
 use tokio::sync::{mpsc, oneshot, watch};
@@ -108,6 +108,9 @@ pub(super) struct FlashblocksState {
     /// Index into ExecutionInfo tracking the last consumed flashblock
     /// Used for slicing transactions/receipts per flashblock
     last_flashblock_tx_index: usize,
+    /// Cached trie updates from previous flashblock for incremental state root calculation.
+    /// None only for the first flashblock; populated after each subsequent state root calculation.
+    prev_trie_updates: Option<Arc<TrieUpdates>>,
 }
 
 struct FallbackBuildOutput<Cache, Transition> {
@@ -157,6 +160,7 @@ impl FlashblocksState {
             da_footprint_per_batch: self.da_footprint_per_batch,
             disable_state_root: self.disable_state_root,
             last_flashblock_tx_index: self.last_flashblock_tx_index,
+            prev_trie_updates: self.prev_trie_updates.clone(),
         }
     }
 
@@ -1199,32 +1203,81 @@ where
     // calculate the state root
     let state_root_start_time = Instant::now();
     let mut state_root = B256::ZERO;
-    let mut trie_output = TrieUpdates::default();
     let mut hashed_state = HashedPostState::default();
+    let mut trie_updates_to_cache: Option<Arc<TrieUpdates>> = None;
 
     if calculate_state_root {
         let state_provider = state.database.as_ref();
+
+        // prev_trie_updates is None for the first flashblock.
+        let prev_trie = fb_state
+            .as_deref()
+            .and_then(|s| s.prev_trie_updates.clone());
+        let flashblock_index = fb_state
+            .as_deref()
+            .map(|s| s.flashblock_index())
+            .unwrap_or(0);
+
         hashed_state = state_provider.hashed_post_state(&state.bundle_state);
-        (state_root, trie_output) = {
+
+        let trie_output;
+        (state_root, trie_output) = if let Some(prev_trie) = prev_trie {
+            // Incremental path: Use cached trie from previous flashblock
+            debug!(
+                target: "payload_builder",
+                flashblock_index,
+                "Using incremental state root calculation with cached trie"
+            );
+
+            let trie_input = TrieInput::new(
+                (*prev_trie).clone(),
+                hashed_state.clone(),
+                hashed_state.construct_prefix_sets(), // Don't freeze - need TriePrefixSetsMut
+            );
+
+            state_provider
+                .state_root_from_nodes_with_updates(trie_input)
+                .map_err(PayloadBuilderError::other)?
+        } else {
+            debug!(
+                target: "payload_builder",
+                flashblock_index,
+                "Using full state root calculation"
+            );
+
             state
                 .database
                 .as_ref()
                 .state_root_with_updates(hashed_state.clone())
                 .inspect_err(|err| {
-                    warn!(target: "payload_builder",
-                    parent_header=%ctx.parent().hash(),
+                    warn!(
+                        target: "payload_builder",
+                        parent_header=%ctx.parent().hash(),
                         %err,
                         "failed to calculate state root for payload"
                     );
                 })?
         };
+
+        // Cache trie updates to apply in fb_state later (avoids mut on fb_state parameter).
+        // Wrap in Arc once so the same allocation is reused for both `executed` and fb_state.
+        trie_updates_to_cache = Some(Arc::new(trie_output));
+
         let state_root_calculation_time = state_root_start_time.elapsed();
         ctx.metrics
             .state_root_calculation_duration
             .record(state_root_calculation_time);
         ctx.metrics
             .state_root_calculation_gauge
             .set(state_root_calculation_time);
+
+        debug!(
+            target: "payload_builder",
+            flashblock_index,
+            state_root = %state_root,
+            duration_ms = state_root_calculation_time.as_millis(),
+            "State root calculation completed"
+        );
     }
 
     let mut requests_hash = None;
@@ -1316,7 +1369,11 @@ where
     let executed = BuiltPayloadExecutedBlock {
         recovered_block: Arc::new(recovered_block),
         execution_output: Arc::new(execution_output),
-        trie_updates: either::Either::Left(Arc::new(trie_output)),
+        trie_updates: either::Either::Left(
+            trie_updates_to_cache
+                .clone()
+                .unwrap_or_else(|| Arc::new(TrieUpdates::default())),
+        ),
         hashed_state: either::Either::Left(Arc::new(hashed_state)),
     };
     debug!(
@@ -1337,6 +1394,9 @@ where
 
     // pick the new transactions from the info field and update the last flashblock index
     let (new_transactions, new_receipts) = if let Some(fb_state) = fb_state {
+        if let Some(updates) = trie_updates_to_cache.take() {
+            fb_state.prev_trie_updates = Some(updates);
+        }
         let new_txs = fb_state.slice_new_transactions(&info.executed_transactions);
         let new_receipts = fb_state.slice_new_receipts(&info.receipts);
         fb_state.set_last_flashblock_tx_index(info.executed_transactions.len());