feat(engine): parallel sparse storage roots (#13269)

crates/engine/tree/src/tree/root.rs

@@ -1,6 +1,7 @@
 //! State root task related functionality.
 
 use alloy_primitives::map::{HashMap, HashSet};
+use rayon::iter::{ParallelBridge, ParallelIterator};
 use reth_evm::system_calls::OnStateHook;
 use reth_provider::{
     providers::ConsistentDbView, BlockReader, DBProvider, DatabaseProviderFactory,
@@ -567,13 +568,19 @@ fn update_sparse_trie(
     trie.reveal_multiproof(targets, multiproof)?;
 
     // Update storage slots with new values and calculate storage roots.
-    for (address, storage) in state.storages {
+    let (tx, rx) = mpsc::channel();
+    state
+        .storages
+        .into_iter()
+        .map(|(address, storage)| (address, storage, trie.take_storage_trie(&address)))
+        .par_bridge()
+        .map(|(address, storage, storage_trie)| {
             trace!(target: "engine::root::sparse", ?address, "Updating storage");
-        let storage_trie = trie.storage_trie_mut(&address).ok_or(SparseTrieError::Blind)?;
+            let mut storage_trie = storage_trie.ok_or(SparseTrieError::Blind)?;
 
             if storage.wiped {
                 trace!(target: "engine::root::sparse", ?address, "Wiping storage");
-            storage_trie.wipe();
+                storage_trie.wipe()?;
             }
 
             for (slot, value) in storage.storage {
@@ -581,7 +588,6 @@ fn update_sparse_trie(
                 if value.is_zero() {
                     trace!(target: "engine::root::sparse", ?address, ?slot, "Removing storage slot");
 
-                // TODO: handle blinded node error
                     storage_trie.remove_leaf(&slot_nibbles)?;
                 } else {
                     trace!(target: "engine::root::sparse", ?address, ?slot, "Updating storage slot");
@@ -591,6 +597,16 @@ fn update_sparse_trie(
             }
 
             storage_trie.root();
+
+            SparseStateTrieResult::Ok((address, storage_trie))
+        })
+        .for_each_init(|| tx.clone(), |tx, result| {
+            tx.send(result).unwrap()
+        });
+    drop(tx);
+    for result in rx {
+        let (address, storage_trie) = result?;
+        trie.insert_storage_trie(address, storage_trie);
     }
 
     // Update accounts with new values

crates/evm/execution-errors/src/trie.rs

@@ -107,14 +107,14 @@ pub enum SparseTrieError {
         /// Path to the node.
         path: Nibbles,
         /// Node that was at the path when revealing.
-        node: Box<dyn core::fmt::Debug>,
+        node: Box<dyn core::fmt::Debug + Send>,
     },
     /// RLP error.
     #[error(transparent)]
     Rlp(#[from] alloy_rlp::Error),
     /// Other.
     #[error(transparent)]
-    Other(#[from] Box<dyn core::error::Error>),
+    Other(#[from] Box<dyn core::error::Error + Send>),
 }
 
 /// Trie witness errors.

crates/trie/sparse/src/state.rs

@@ -97,9 +97,26 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
     /// Returns mutable reference to storage sparse trie if it was revealed.
     pub fn storage_trie_mut(
         &mut self,
-        account: &B256,
+        address: &B256,
     ) -> Option<&mut RevealedSparseTrie<F::StorageNodeProvider>> {
-        self.storages.get_mut(account).and_then(|e| e.as_revealed_mut())
+        self.storages.get_mut(address).and_then(|e| e.as_revealed_mut())
+    }
+
+    /// Takes the storage trie for the provided address.
+    pub fn take_storage_trie(
+        &mut self,
+        address: &B256,
+    ) -> Option<SparseTrie<F::StorageNodeProvider>> {
+        self.storages.remove(address)
+    }
+
+    /// Inserts storage trie for the provided address.
+    pub fn insert_storage_trie(
+        &mut self,
+        address: B256,
+        storage_trie: SparseTrie<F::StorageNodeProvider>,
+    ) {
+        self.storages.insert(address, storage_trie);
     }
 
     /// Reveal unknown trie paths from provided leaf path and its proof for the account.