diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index e2ed6aa147..a8c455265d 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/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,30 +568,45 @@ 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 { - trace!(target: "engine::root::sparse", ?address, "Updating storage"); - let storage_trie = trie.storage_trie_mut(&address).ok_or(SparseTrieError::Blind)?; + 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 mut storage_trie = storage_trie.ok_or(SparseTrieError::Blind)?; - if storage.wiped { - trace!(target: "engine::root::sparse", ?address, "Wiping storage"); - storage_trie.wipe(); - } - - for (slot, value) in storage.storage { - let slot_nibbles = Nibbles::unpack(slot); - 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"); - storage_trie - .update_leaf(slot_nibbles, alloy_rlp::encode_fixed_size(&value).to_vec())?; + if storage.wiped { + trace!(target: "engine::root::sparse", ?address, "Wiping storage"); + storage_trie.wipe()?; } - } - storage_trie.root(); + for (slot, value) in storage.storage { + let slot_nibbles = Nibbles::unpack(slot); + if value.is_zero() { + trace!(target: "engine::root::sparse", ?address, ?slot, "Removing storage slot"); + + storage_trie.remove_leaf(&slot_nibbles)?; + } else { + trace!(target: "engine::root::sparse", ?address, ?slot, "Updating storage slot"); + storage_trie + .update_leaf(slot_nibbles, alloy_rlp::encode_fixed_size(&value).to_vec())?; + } + } + + 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 diff --git a/crates/evm/execution-errors/src/trie.rs b/crates/evm/execution-errors/src/trie.rs index 8d04f97e8e..ba1bfcc023 100644 --- a/crates/evm/execution-errors/src/trie.rs +++ b/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, + node: Box, }, /// RLP error. #[error(transparent)] Rlp(#[from] alloy_rlp::Error), /// Other. #[error(transparent)] - Other(#[from] Box), + Other(#[from] Box), } /// Trie witness errors. diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 6638632f0a..6f5db4eda7 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -97,9 +97,26 @@ impl SparseStateTrie { /// 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> { - 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> { + self.storages.remove(address) + } + + /// Inserts storage trie for the provided address. + pub fn insert_storage_trie( + &mut self, + address: B256, + storage_trie: SparseTrie, + ) { + self.storages.insert(address, storage_trie); } /// Reveal unknown trie paths from provided leaf path and its proof for the account.