perf(trie): parallelize merge_ancestors_into_overlay (#21473)

Co-authored-by: Amp <amp@ampcode.com>
Co-authored-by: DaniPopes <57450786+DaniPopes@users.noreply.github.com>

Author: 3 people

crates/chain-state/src/deferred_trie.rs

@@ -274,8 +274,54 @@ impl DeferredTrieData {
     /// In normal operation, the parent always has a cached overlay and this
     /// function is never called.
     ///
-    /// Iterates ancestors oldest -> newest, then extends with current block's data,
-    /// so later state takes precedence.
+    /// When the `rayon` feature is enabled:
+    /// 1. Collects ancestor data (states and updates)
+    /// 2. Merges states and trie updates in parallel using k-way merge
+    #[cfg(feature = "rayon")]
+    fn merge_ancestors_into_overlay(
+        ancestors: &[Self],
+        sorted_hashed_state: &HashedPostStateSorted,
+        sorted_trie_updates: &TrieUpdatesSorted,
+    ) -> TrieInputSorted {
+        // Early exit: no ancestors means just wrap current block's data
+        if ancestors.is_empty() {
+            return TrieInputSorted::new(
+                Arc::new(sorted_trie_updates.clone()),
+                Arc::new(sorted_hashed_state.clone()),
+                Default::default(),
+            );
+        }
+
+        // Collect ancestor data in reverse (newest to oldest) for merge_slice
+        let (states, updates): (Vec<_>, Vec<_>) = ancestors
+            .iter()
+            .rev()
+            .map(|a| {
+                // Note: we can assume that this trie data has already been computed
+                let data = a.wait_cloned();
+                (data.hashed_state, data.trie_updates)
+            })
+            .unzip();
+
+        // Merge state and nodes in parallel using k-way merge
+        let (state, nodes) = rayon::join(
+            || {
+                let mut merged = HashedPostStateSorted::merge_slice(&states);
+                merged.extend_ref_and_sort(sorted_hashed_state);
+                merged
+            },
+            || {
+                let mut merged = TrieUpdatesSorted::merge_slice(&updates);
+                merged.extend_ref_and_sort(sorted_trie_updates);
+                merged
+            },
+        );
+
+        TrieInputSorted::new(Arc::new(nodes), Arc::new(state), Default::default())
+    }
+
+    /// Sequential fallback when rayon is not available.
+    #[cfg(not(feature = "rayon"))]
     fn merge_ancestors_into_overlay(
         ancestors: &[Self],
         sorted_hashed_state: &HashedPostStateSorted,
@@ -293,18 +339,8 @@ impl DeferredTrieData {
             nodes_mut.extend_ref_and_sort(ancestor_data.trie_updates.as_ref());
         }
 
-        // Extend with current block's sorted data last (takes precedence)
-        #[cfg(feature = "rayon")]
-        rayon::join(
-            || state_mut.extend_ref_and_sort(sorted_hashed_state),
-            || nodes_mut.extend_ref_and_sort(sorted_trie_updates),
-        );
-
-        #[cfg(not(feature = "rayon"))]
-        {
-            state_mut.extend_ref_and_sort(sorted_hashed_state);
-            nodes_mut.extend_ref_and_sort(sorted_trie_updates);
-        }
+        state_mut.extend_ref_and_sort(sorted_hashed_state);
+        nodes_mut.extend_ref_and_sort(sorted_trie_updates);
 
         overlay
     }

crates/trie/common/src/hashed_state.rs

@@ -614,16 +614,29 @@ impl HashedPostStateSorted {
     /// For small batches, uses `extend_ref_and_sort` loop.
     /// For large batches, uses k-way merge for O(n log k) complexity.
     pub fn merge_batch<T: AsRef<Self> + From<Self>>(iter: impl IntoIterator<Item = T>) -> T {
+        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
+        match items.len() {
+            0 => Self::default().into(),
+            1 => items.into_iter().next().expect("len == 1"),
+            _ => Self::merge_slice(&items).into(),
+        }
+    }
+
+    /// Batch-merge sorted hashed post states from a slice. Slice is **newest to oldest**.
+    ///
+    /// This variant takes a slice reference directly, avoiding iterator collection overhead.
+    /// For small batches, uses `extend_ref_and_sort` loop.
+    /// For large batches, uses k-way merge for O(n log k) complexity.
+    pub fn merge_slice<T: AsRef<Self>>(items: &[T]) -> Self {
         const THRESHOLD: usize = 30;
 
-        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
         let k = items.len();
 
         if k == 0 {
-            return Self::default().into();
+            return Self::default();
         }
         if k == 1 {
-            return items.into_iter().next().expect("k == 1");
+            return items[0].as_ref().clone();
         }
 
         if k < THRESHOLD {
@@ -633,7 +646,7 @@ impl HashedPostStateSorted {
             for next in iter {
                 acc.extend_ref_and_sort(next.as_ref());
             }
-            return acc.into();
+            return acc;
         }
 
         // Large k: k-way merge.
@@ -647,7 +660,7 @@ impl HashedPostStateSorted {
 
         let mut acc: B256Map<StorageAcc<'_>> = B256Map::default();
 
-        for item in &items {
+        for item in items {
             for (addr, storage) in &item.as_ref().storages {
                 let entry = acc.entry(*addr).or_insert_with(|| StorageAcc {
                     wiped: false,
@@ -675,7 +688,7 @@ impl HashedPostStateSorted {
             })
             .collect();
 
-        Self { accounts, storages }.into()
+        Self { accounts, storages }
     }
 
     /// Clears all accounts and storage data.

crates/trie/common/src/updates.rs

@@ -632,16 +632,29 @@ impl TrieUpdatesSorted {
     /// For small batches, uses `extend_ref_and_sort` loop.
     /// For large batches, uses k-way merge for O(n log k) complexity.
     pub fn merge_batch<T: AsRef<Self> + From<Self>>(iter: impl IntoIterator<Item = T>) -> T {
+        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
+        match items.len() {
+            0 => Self::default().into(),
+            1 => items.into_iter().next().expect("len == 1"),
+            _ => Self::merge_slice(&items).into(),
+        }
+    }
+
+    /// Batch-merge sorted trie updates from a slice. Slice is **newest to oldest**.
+    ///
+    /// This variant takes a slice reference directly, avoiding iterator collection overhead.
+    /// For small batches, uses `extend_ref_and_sort` loop.
+    /// For large batches, uses k-way merge for O(n log k) complexity.
+    pub fn merge_slice<T: AsRef<Self>>(items: &[T]) -> Self {
         const THRESHOLD: usize = 30;
 
-        let items: alloc::vec::Vec<_> = iter.into_iter().collect();
         let k = items.len();
 
         if k == 0 {
-            return Self::default().into();
+            return Self::default();
         }
         if k == 1 {
-            return items.into_iter().next().expect("k == 1");
+            return items[0].as_ref().clone();
         }
 
         if k < THRESHOLD {
@@ -651,7 +664,7 @@ impl TrieUpdatesSorted {
             for next in iter {
                 acc.extend_ref_and_sort(next.as_ref());
             }
-            return acc.into();
+            return acc;
         }
 
         // Large k: k-way merge.
@@ -666,7 +679,7 @@ impl TrieUpdatesSorted {
 
         let mut acc: B256Map<StorageAcc<'_>> = B256Map::default();
 
-        for item in &items {
+        for item in items {
             for (addr, storage) in &item.as_ref().storage_tries {
                 let entry = acc.entry(*addr).or_insert_with(|| StorageAcc {
                     is_deleted: false,
@@ -695,7 +708,7 @@ impl TrieUpdatesSorted {
             })
             .collect();
 
-        Self { account_nodes, storage_tries }.into()
+        Self { account_nodes, storage_tries }
     }
 }