perf(storage): batch trie updates across blocks in save_blocks

crates/chain-state/src/lazy_overlay.rs

@@ -10,11 +10,6 @@ use reth_trie::{updates::TrieUpdatesSorted, HashedPostStateSorted, TrieInputSort
 use std::sync::{Arc, OnceLock};
 use tracing::{debug, trace};
 
-/// Threshold for switching from `extend_ref` loop to `merge_batch`.
-///
-/// Benchmarked crossover: `extend_ref` wins up to ~64 blocks, `merge_batch` wins beyond.
-const MERGE_BATCH_THRESHOLD: usize = 64;
-
 /// Inputs captured for lazy overlay computation.
 #[derive(Clone)]
 struct LazyOverlayInputs {
@@ -128,57 +123,28 @@ impl LazyOverlay {
 
     /// Merge all blocks' trie data into a single [`TrieInputSorted`].
     ///
-    /// Blocks are ordered newest to oldest. We iterate oldest to newest so that
-    /// newer values override older ones.
+    /// Blocks are ordered newest to oldest. Uses hybrid merge algorithm that
+    /// switches between `extend_ref` (small batches) and k-way merge (large batches).
     fn merge_blocks(blocks: &[DeferredTrieData]) -> TrieInputSorted {
         if blocks.is_empty() {
             return TrieInputSorted::default();
         }
 
-        // Single block: use its data directly
-        if blocks.len() == 1 {
-            let data = blocks[0].wait_cloned();
-            return TrieInputSorted {
-                state: Arc::clone(&data.hashed_state),
-                nodes: Arc::clone(&data.trie_updates),
-                prefix_sets: Default::default(),
-            };
-        }
-
-        if blocks.len() < MERGE_BATCH_THRESHOLD {
-            // Small k: extend_ref loop is faster
-            // Iterate oldest->newest so newer values override older ones
-            let mut blocks_iter = blocks.iter().rev();
-            let first = blocks_iter.next().expect("blocks is non-empty");
-            let data = first.wait_cloned();
-
-            let mut state = Arc::clone(&data.hashed_state);
-            let mut nodes = Arc::clone(&data.trie_updates);
-            let state_mut = Arc::make_mut(&mut state);
-            let nodes_mut = Arc::make_mut(&mut nodes);
-
-            for block in blocks_iter {
-                let data = block.wait_cloned();
-                state_mut.extend_ref(data.hashed_state.as_ref());
-                nodes_mut.extend_ref(data.trie_updates.as_ref());
-            }
-
-            TrieInputSorted { state, nodes, prefix_sets: Default::default() }
-        } else {
-            // Large k: merge_batch is faster (O(n log k) via k-way merge)
-            let trie_data: Vec<_> = blocks.iter().map(|b| b.wait_cloned()).collect();
-
-            let merged_state = HashedPostStateSorted::merge_batch(
-                trie_data.iter().map(|d| d.hashed_state.as_ref()),
-            );
-            let merged_nodes =
-                TrieUpdatesSorted::merge_batch(trie_data.iter().map(|d| d.trie_updates.as_ref()));
-
-            TrieInputSorted {
-                state: Arc::new(merged_state),
-                nodes: Arc::new(merged_nodes),
-                prefix_sets: Default::default(),
-            }
+        // Collect all trie data first (blocks are newest-to-oldest)
+        let trie_data: Vec<_> = blocks.iter().map(|b| b.wait_cloned()).collect();
+
+        // Use hybrid merge which handles the threshold internally
+        let merged_state = HashedPostStateSorted::merge_batch_hybrid(
+            trie_data.iter().map(|d| d.hashed_state.as_ref()),
+        );
+        let merged_nodes = TrieUpdatesSorted::merge_batch_hybrid(
+            trie_data.iter().map(|d| d.trie_updates.as_ref()),
+        );
+
+        TrieInputSorted {
+            state: Arc::new(merged_state),
+            nodes: Arc::new(merged_nodes),
+            prefix_sets: Default::default(),
         }
     }
 }

crates/storage/provider/src/providers/database/provider.rs

@@ -555,11 +555,24 @@ impl<TX: DbTx + DbTxMut + 'static, N: NodeTypesForProvider> DatabaseProvider<TX,
                     let start = Instant::now();
                     self.write_hashed_state(&trie_data.hashed_state)?;
                     timings.write_hashed_state += start.elapsed();
+                }
+            }
 
-                    let start = Instant::now();
-                    self.write_trie_updates_sorted(&trie_data.trie_updates)?;
-                    timings.write_trie_updates += start.elapsed();
+            // Write all trie updates in a single batch.
+            // This reduces cursor open/close overhead from N calls to 1.
+            // Uses hybrid algorithm: extend_ref for small batches, k-way merge for large.
+            if save_mode.with_state() {
+                let start = Instant::now();
+                // Collect Arc refs first to extend their lifetime
+                let trie_updates: Vec<_> = blocks.iter().map(|b| b.trie_updates()).collect();
+                // merge_batch_hybrid expects newest-to-oldest, so reverse the iterator
+                let merged = TrieUpdatesSorted::merge_batch_hybrid(
+                    trie_updates.iter().rev().map(|arc| arc.as_ref()),
+                );
+                if !merged.is_empty() {
+                    self.write_trie_updates_sorted(&merged)?;
                 }
+                timings.write_trie_updates += start.elapsed();
             }
 
             // Full mode: update history indices

crates/trie/common/src/hashed_state.rs

@@ -690,6 +690,47 @@ impl HashedPostStateSorted {
         Self { accounts, storages }
     }
 
+    /// Hybrid batch-merge sorted hashed post states. Iterator yields **newest to oldest**.
+    ///
+    /// Uses a hybrid algorithm that switches between `extend_ref` (for small batches)
+    /// and k-way `merge_batch` (for large batches) based on benchmarked crossover point.
+    ///
+    /// - Small k (< threshold): O(n * k) via `extend_ref` loop, but with low constant factors
+    /// - Large k (≥ threshold): O(n log k) via k-way merge
+    ///
+    /// The threshold is tuned based on benchmarks where `extend_ref` wins up to ~64 items.
+    pub fn merge_batch_hybrid<'a>(states: impl IntoIterator<Item = &'a Self>) -> Self {
+        const MERGE_BATCH_THRESHOLD: usize = 64;
+
+        let states: Vec<_> = states.into_iter().collect();
+
+        if states.is_empty() {
+            return Self::default();
+        }
+
+        if states.len() == 1 {
+            return states[0].clone();
+        }
+
+        if states.len() < MERGE_BATCH_THRESHOLD {
+            // Small k: extend_ref loop is faster.
+            // States are newest-to-oldest, so iterate in reverse (oldest-to-newest)
+            // to let newer values override older ones.
+            let mut iter = states.iter().rev();
+            let first = iter.next().expect("states is non-empty");
+            let mut result = (*first).clone();
+
+            for state in iter {
+                result.extend_ref(state);
+            }
+
+            result
+        } else {
+            // Large k: merge_batch is faster (O(n log k) via k-way merge)
+            Self::merge_batch(states)
+        }
+    }
+
     /// Clears all accounts and storage data.
     pub fn clear(&mut self) {
         self.accounts.clear();

crates/trie/common/src/updates.rs

@@ -686,6 +686,47 @@ impl TrieUpdatesSorted {
 
         Self { account_nodes, storage_tries }
     }
+
+    /// Hybrid batch-merge sorted trie updates. Iterator yields **newest to oldest**.
+    ///
+    /// Uses a hybrid algorithm that switches between `extend_ref` (for small batches)
+    /// and k-way `merge_batch` (for large batches) based on benchmarked crossover point.
+    ///
+    /// - Small k (< threshold): O(n * k) via `extend_ref` loop, but with low constant factors
+    /// - Large k (≥ threshold): O(n log k) via k-way merge
+    ///
+    /// The threshold is tuned based on benchmarks where `extend_ref` wins up to ~64 items.
+    pub fn merge_batch_hybrid<'a>(updates: impl IntoIterator<Item = &'a Self>) -> Self {
+        const MERGE_BATCH_THRESHOLD: usize = 64;
+
+        let updates: Vec<_> = updates.into_iter().collect();
+
+        if updates.is_empty() {
+            return Self::default();
+        }
+
+        if updates.len() == 1 {
+            return updates[0].clone();
+        }
+
+        if updates.len() < MERGE_BATCH_THRESHOLD {
+            // Small k: extend_ref loop is faster.
+            // Updates are newest-to-oldest, so iterate in reverse (oldest-to-newest)
+            // to let newer values override older ones.
+            let mut iter = updates.iter().rev();
+            let first = iter.next().expect("updates is non-empty");
+            let mut result = (*first).clone();
+
+            for update in iter {
+                result.extend_ref(update);
+            }
+
+            result
+        } else {
+            // Large k: merge_batch is faster (O(n log k) via k-way merge)
+            Self::merge_batch(updates)
+        }
+    }
 }
 
 impl AsRef<Self> for TrieUpdatesSorted {