clone

Cargo.lock

@@ -11178,6 +11178,7 @@ dependencies = [
  "bincode 1.3.3",
  "bytes",
  "codspeed-criterion-compat",
+ "dashmap 6.1.0",
  "derive_more",
  "hash-db",
  "itertools 0.14.0",

crates/engine/tree/Cargo.toml

@@ -44,6 +44,8 @@ alloy-primitives.workspace = true
 alloy-rlp.workspace = true
 alloy-rpc-types-engine.workspace = true
 
+parking_lot.workspace = true
+
 revm.workspace = true
 revm-primitives.workspace = true
 
@@ -65,7 +67,6 @@ schnellru.workspace = true
 rayon.workspace = true
 tracing.workspace = true
 derive_more.workspace = true
-parking_lot.workspace = true
 crossbeam-channel.workspace = true
 
 # optional deps for test-utils

crates/engine/tree/src/tree/payload_processor/multiproof.rs

@@ -225,7 +225,7 @@ struct MultiproofInput {
     proof_targets: MultiProofTargets,
     proof_sequence_number: u64,
     state_root_message_sender: CrossbeamSender<MultiProofMessage>,
-    multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
+    multi_added_removed_keys: Option<MultiAddedRemovedKeys>,
 }
 
 impl MultiproofInput {
@@ -608,9 +608,6 @@ impl MultiProofTask {
         // we still want to optimistically fetch extension children for the leaf addition case.
         self.multi_added_removed_keys.touch_accounts(targets.keys().copied());
 
-        // Clone+Arc MultiAddedRemovedKeys for sharing with the dispatched multiproof tasks
-        let multi_added_removed_keys = Arc::new(self.multi_added_removed_keys.clone());
-
         self.metrics.prefetch_proof_targets_accounts_histogram.record(targets.len() as f64);
         self.metrics
             .prefetch_proof_targets_storages_histogram
@@ -636,7 +633,7 @@ impl MultiProofTask {
                     proof_targets,
                     proof_sequence_number: self.proof_sequencer.next_sequence(),
                     state_root_message_sender: self.tx.clone(),
-                    multi_added_removed_keys: Some(multi_added_removed_keys.clone()),
+                    multi_added_removed_keys: Some(self.multi_added_removed_keys.clone()),
                 });
             },
         );
@@ -704,9 +701,6 @@ impl MultiProofTask {
             state_updates += 1;
         }
 
-        // Clone+Arc MultiAddedRemovedKeys for sharing with the dispatched multiproof tasks
-        let multi_added_removed_keys = Arc::new(self.multi_added_removed_keys.clone());
-
         let chunking_len = not_fetched_state_update.chunking_length();
         let mut spawned_proof_targets = MultiProofTargets::default();
         let available_account_workers =
@@ -725,7 +719,7 @@ impl MultiProofTask {
                 let proof_targets = get_proof_targets(
                     &hashed_state_update,
                     &self.fetched_proof_targets,
-                    &multi_added_removed_keys,
+                    &self.multi_added_removed_keys,
                 );
                 spawned_proof_targets.extend_ref(&proof_targets);
 
@@ -735,7 +729,7 @@ impl MultiProofTask {
                     proof_targets,
                     proof_sequence_number: self.proof_sequencer.next_sequence(),
                     state_root_message_sender: self.tx.clone(),
-                    multi_added_removed_keys: Some(multi_added_removed_keys.clone()),
+                    multi_added_removed_keys: Some(self.multi_added_removed_keys.clone()),
                 });
             },
         );
@@ -1217,7 +1211,7 @@ fn get_proof_targets(
             .keys()
             .filter(|slot| {
                 !fetched.is_some_and(|f| f.contains(*slot)) ||
-                    storage_added_removed_keys.is_some_and(|k| k.is_removed(slot))
+                    storage_added_removed_keys.as_ref().is_some_and(|k| k.is_removed(slot))
             })
             .peekable();
 

crates/trie/common/Cargo.toml

@@ -28,6 +28,7 @@ bytes = { workspace = true, optional = true }
 derive_more.workspace = true
 itertools = { workspace = true, features = ["use_alloc"] }
 nybbles = { workspace = true, features = ["rlp"] }
+dashmap.workspace = true
 
 # reth
 revm-database.workspace = true

crates/trie/common/src/added_removed_keys.rs

@@ -1,14 +1,17 @@
 //! Tracking of keys having been added and removed from the tries.
 
+use std::sync::Arc;
+
 use crate::HashedPostState;
-use alloy_primitives::{map::B256Map, B256};
+use alloy_primitives::B256;
 use alloy_trie::proof::AddedRemovedKeys;
+use dashmap::{mapref::one::Ref, DashMap};
 
 /// Tracks added and removed keys across account and storage tries.
 #[derive(Debug, Clone)]
 pub struct MultiAddedRemovedKeys {
     account: AddedRemovedKeys,
-    storages: B256Map<AddedRemovedKeys>,
+    storages: Arc<DashMap<B256, AddedRemovedKeys, alloy_primitives::map::DefaultHashBuilder>>,
 }
 
 /// Returns [`AddedRemovedKeys`] with default parameters. This is necessary while we are not yet
@@ -50,7 +53,7 @@ impl MultiAddedRemovedKeys {
                 continue
             }
 
-            let storage_removed_keys =
+            let mut storage_removed_keys =
                 self.storages.entry(*hashed_address).or_insert_with(default_added_removed_keys);
 
             for (key, val) in &storage.storage {
@@ -68,7 +71,7 @@ impl MultiAddedRemovedKeys {
     }
 
     /// Returns a [`AddedRemovedKeys`] for the storage trie of a particular account, if any.
-    pub fn get_storage(&self, hashed_address: &B256) -> Option<&AddedRemovedKeys> {
+    pub fn get_storage(&self, hashed_address: &B256) -> Option<Ref<'_, B256, AddedRemovedKeys>> {
         self.storages.get(hashed_address)
     }
 

crates/trie/common/src/hashed_state.rs

@@ -215,7 +215,9 @@ impl HashedPostState {
                     let mut storage_not_in_targets = HashedStorage::default();
                     storage.storage.retain(|&slot, value| {
                         if storage_in_targets.contains(&slot) &&
-                            !storage_added_removed_keys.is_some_and(|k| k.is_removed(&slot))
+                            !storage_added_removed_keys
+                                .as_ref()
+                                .is_some_and(|k| k.is_removed(&slot))
                         {
                             return true
                         }

crates/trie/parallel/src/proof.rs

@@ -30,7 +30,7 @@ pub struct ParallelProof {
     /// Flag indicating whether to include branch node masks in the proof.
     collect_branch_node_masks: bool,
     /// Provided by the user to give the necessary context to retain extra proofs.
-    multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
+    multi_added_removed_keys: Option<MultiAddedRemovedKeys>,
     /// Handle to the proof worker pools.
     proof_worker_handle: ProofWorkerHandle,
     /// Whether to use V2 storage proofs.
@@ -72,7 +72,7 @@ impl ParallelProof {
     /// extra proofs needed to add and remove leaf nodes from the tries.
     pub fn with_multi_added_removed_keys(
         mut self,
-        multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
+        multi_added_removed_keys: Option<MultiAddedRemovedKeys>,
     ) -> Self {
         self.multi_added_removed_keys = multi_added_removed_keys;
         self

crates/trie/parallel/src/proof_task.rs

@@ -447,8 +447,9 @@ where
 
         // Get or create added/removed keys context
         let multi_added_removed_keys =
-            multi_added_removed_keys.unwrap_or_else(|| Arc::new(MultiAddedRemovedKeys::new()));
-        let added_removed_keys = multi_added_removed_keys.get_storage(&hashed_address);
+            multi_added_removed_keys.unwrap_or_else(MultiAddedRemovedKeys::new);
+        let added_removed_keys =
+            multi_added_removed_keys.get_storage(&hashed_address).map(|k| k.as_ref().clone());
 
         let span = debug_span!(
             target: "trie::proof_task",
@@ -1556,7 +1557,7 @@ fn dispatch_storage_proofs(
     targets: &MultiProofTargets,
     storage_prefix_sets: &mut B256Map<PrefixSet>,
     with_branch_node_masks: bool,
-    multi_added_removed_keys: Option<&Arc<MultiAddedRemovedKeys>>,
+    multi_added_removed_keys: Option<&MultiAddedRemovedKeys>,
     use_v2_proofs: bool,
 ) -> Result<B256Map<CrossbeamReceiver<StorageProofResultMessage>>, ParallelStateRootError> {
     let mut storage_proof_receivers =
@@ -1613,7 +1614,7 @@ pub enum StorageProofInput {
         /// Whether or not to collect branch node masks
         with_branch_node_masks: bool,
         /// Provided by the user to give the necessary context to retain extra proofs.
-        multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
+        multi_added_removed_keys: Option<MultiAddedRemovedKeys>,
     },
     /// V2 storage proof variant
     V2 {
@@ -1632,7 +1633,7 @@ impl StorageProofInput {
         prefix_set: PrefixSet,
         target_slots: B256Set,
         with_branch_node_masks: bool,
-        multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
+        multi_added_removed_keys: Option<MultiAddedRemovedKeys>,
     ) -> Self {
         Self::Legacy {
             hashed_address,
@@ -1668,7 +1669,7 @@ pub struct AccountMultiproofInput {
     /// Whether or not to collect branch node masks.
     pub collect_branch_node_masks: bool,
     /// Provided by the user to give the necessary context to retain extra proofs.
-    pub multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
+    pub multi_added_removed_keys: Option<MultiAddedRemovedKeys>,
     /// Context for sending the proof result.
     pub proof_result_sender: ProofResultContext,
     /// Whether to use V2 storage proofs.
@@ -1684,7 +1685,7 @@ struct AccountMultiproofParams<'a> {
     /// Whether or not to collect branch node masks.
     collect_branch_node_masks: bool,
     /// Provided by the user to give the necessary context to retain extra proofs.
-    multi_added_removed_keys: Option<&'a Arc<MultiAddedRemovedKeys>>,
+    multi_added_removed_keys: Option<&'a MultiAddedRemovedKeys>,
     /// Receivers for storage proofs being computed in parallel.
     storage_proof_receivers: B256Map<CrossbeamReceiver<StorageProofResultMessage>>,
     /// Cached storage roots. This will be used to read storage roots for missed leaves, as well as