chore(trie): convert MultiProofTargets into unit struct (#14748)

crates/engine/tree/src/tree/mod.rs

@@ -2791,8 +2791,7 @@ where
 
             let Some(state_root_sender) = state_root_sender else { return };
 
-            let mut targets =
-                MultiProofTargets::with_capacity_and_hasher(res.state.len(), Default::default());
+            let mut targets = MultiProofTargets::with_capacity(res.state.len());
             let mut storage_targets = 0;
             for (addr, account) in res.state {
                 // if the account was not touched, or if the account was selfdestructed, do not

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

@@ -31,6 +31,7 @@ use reth_trie_sparse::{
 use revm_primitives::{keccak256, B256};
 use std::{
     collections::{BTreeMap, VecDeque},
+    ops::DerefMut,
     sync::{
         mpsc::{self, channel, Receiver, Sender},
         Arc,
@@ -707,7 +708,7 @@ where
     /// Handles request for proof prefetch.
     fn on_prefetch_proof(&mut self, targets: MultiProofTargets) {
         let proof_targets = self.get_prefetch_proof_targets(targets);
-        extend_multi_proof_targets_ref(&mut self.fetched_proof_targets, &proof_targets);
+        self.fetched_proof_targets.extend_ref(&proof_targets);
 
         self.multiproof_manager.spawn_or_queue(MultiProofInput {
             config: self.config.clone(),
@@ -745,7 +746,7 @@ where
         });
 
         // For all non-subset remaining targets, we have to calculate the difference
-        for (hashed_address, target_storage) in &mut targets {
+        for (hashed_address, target_storage) in targets.deref_mut() {
             let Some(fetched_storage) = self.fetched_proof_targets.get(hashed_address) else {
                 // this means the account has not been fetched yet, so we must fetch everything
                 // associated with this account
@@ -780,7 +781,7 @@ where
     ) {
         let hashed_state_update = evm_state_to_hashed_post_state(update);
         let proof_targets = get_proof_targets(&hashed_state_update, &self.fetched_proof_targets);
-        extend_multi_proof_targets_ref(&mut self.fetched_proof_targets, &proof_targets);
+        self.fetched_proof_targets.extend_ref(&proof_targets);
 
         self.multiproof_manager.spawn_or_queue(MultiProofInput {
             config: self.config.clone(),
@@ -1282,12 +1283,6 @@ where
     Ok(elapsed)
 }
 
-fn extend_multi_proof_targets_ref(targets: &mut MultiProofTargets, other: &MultiProofTargets) {
-    for (address, slots) in other {
-        targets.entry(*address).or_default().extend(slots);
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

crates/trie/common/src/hashed_state.rs

@@ -159,8 +159,7 @@ impl HashedPostState {
     /// Create multiproof targets for this state.
     pub fn multi_proof_targets(&self) -> MultiProofTargets {
         // Pre-allocate minimum capacity for the targets.
-        let mut targets =
-            MultiProofTargets::with_capacity_and_hasher(self.accounts.len(), Default::default());
+        let mut targets = MultiProofTargets::with_capacity(self.accounts.len());
         for hashed_address in self.accounts.keys() {
             targets.insert(*hashed_address, Default::default());
         }

crates/trie/common/src/proofs.rs

@@ -1,7 +1,7 @@
 //! Merkle trie proofs.
 
 use crate::{Nibbles, TrieAccount};
-use alloc::vec::Vec;
+use alloc::{borrow::Cow, vec::Vec};
 use alloy_consensus::constants::KECCAK_EMPTY;
 use alloy_primitives::{
     keccak256,
@@ -14,11 +14,44 @@ use alloy_trie::{
     proof::{verify_proof, DecodedProofNodes, ProofNodes, ProofVerificationError},
     TrieMask, EMPTY_ROOT_HASH,
 };
+use derive_more::{Deref, DerefMut, IntoIterator};
 use itertools::Itertools;
 use reth_primitives_traits::Account;
 
 /// Proof targets map.
-pub type MultiProofTargets = B256Map<B256Set>;
+#[derive(Deref, DerefMut, IntoIterator, Clone, PartialEq, Eq, Default, Debug)]
+pub struct MultiProofTargets(B256Map<B256Set>);
+
+impl FromIterator<(B256, B256Set)> for MultiProofTargets {
+    fn from_iter<T: IntoIterator<Item = (B256, B256Set)>>(iter: T) -> Self {
+        Self(B256Map::from_iter(iter))
+    }
+}
+
+impl MultiProofTargets {
+    /// Creates an empty `MultiProofTargets` with at least the specified capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self(B256Map::with_capacity_and_hasher(capacity, Default::default()))
+    }
+
+    /// Extend multi proof targets with contents of other.
+    pub fn extend(&mut self, other: Self) {
+        self.extend_inner(Cow::Owned(other));
+    }
+
+    /// Extend multi proof targets with contents of other.
+    ///
+    /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`.
+    pub fn extend_ref(&mut self, other: &Self) {
+        self.extend_inner(Cow::Borrowed(other));
+    }
+
+    fn extend_inner(&mut self, other: Cow<'_, Self>) {
+        for (hashed_address, hashed_slots) in other.iter() {
+            self.entry(*hashed_address).or_default().extend(hashed_slots);
+        }
+    }
+}
 
 /// The state multiproof of target accounts and multiproofs of their storage tries.
 /// Multiproof is effectively a state subtrie that only contains the nodes

crates/trie/db/tests/witness.rs

@@ -11,7 +11,10 @@ use reth_db::{cursor::DbCursorRW, tables};
 use reth_db_api::transaction::DbTxMut;
 use reth_primitives_traits::{Account, StorageEntry};
 use reth_provider::{test_utils::create_test_provider_factory, HashingWriter};
-use reth_trie::{proof::Proof, witness::TrieWitness, HashedPostState, HashedStorage, StateRoot};
+use reth_trie::{
+    proof::Proof, witness::TrieWitness, HashedPostState, HashedStorage, MultiProofTargets,
+    StateRoot,
+};
 use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness};
 
 #[test]
@@ -39,7 +42,10 @@ fn includes_empty_node_preimage() {
 
     let state_root = StateRoot::from_tx(provider.tx_ref()).root().unwrap();
     let multiproof = Proof::from_tx(provider.tx_ref())
-        .multiproof(HashMap::from_iter([(hashed_address, HashSet::from_iter([hashed_slot]))]))
+        .multiproof(MultiProofTargets::from_iter([(
+            hashed_address,
+            HashSet::from_iter([hashed_slot]),
+        )]))
         .unwrap();
 
     let witness = TrieWitness::from_tx(provider.tx_ref())
@@ -77,7 +83,10 @@ fn includes_nodes_for_destroyed_storage_nodes() {
 
     let state_root = StateRoot::from_tx(provider.tx_ref()).root().unwrap();
     let multiproof = Proof::from_tx(provider.tx_ref())
-        .multiproof(HashMap::from_iter([(hashed_address, HashSet::from_iter([hashed_slot]))]))
+        .multiproof(MultiProofTargets::from_iter([(
+            hashed_address,
+            HashSet::from_iter([hashed_slot]),
+        )]))
         .unwrap();
 
     let witness =
@@ -122,7 +131,7 @@ fn correctly_decodes_branch_node_values() {
 
     let state_root = StateRoot::from_tx(provider.tx_ref()).root().unwrap();
     let multiproof = Proof::from_tx(provider.tx_ref())
-        .multiproof(HashMap::from_iter([(
+        .multiproof(MultiProofTargets::from_iter([(
             hashed_address,
             HashSet::from_iter([hashed_slot1, hashed_slot2]),
         )]))

crates/trie/trie/src/proof/blinded.rs

@@ -1,11 +1,8 @@
 use super::{Proof, StorageProof};
 use crate::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory};
-use alloy_primitives::{
-    map::{HashMap, HashSet},
-    B256,
-};
+use alloy_primitives::{map::HashSet, B256};
 use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind};
-use reth_trie_common::{prefix_set::TriePrefixSetsMut, Nibbles};
+use reth_trie_common::{prefix_set::TriePrefixSetsMut, MultiProofTargets, Nibbles};
 use reth_trie_sparse::blinded::{
     pad_path_to_key, BlindedProvider, BlindedProviderFactory, RevealedNode,
 };
@@ -90,7 +87,7 @@ where
     fn blinded_node(&mut self, path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError> {
         let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now);
 
-        let targets = HashMap::from_iter([(pad_path_to_key(path), HashSet::default())]);
+        let targets = MultiProofTargets::from_iter([(pad_path_to_key(path), HashSet::default())]);
         let mut proof =
             Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
                 .with_prefix_sets_mut(self.prefix_sets.as_ref().clone())

crates/trie/trie/src/proof/mod.rs

@@ -93,7 +93,7 @@ where
         slots: &[B256],
     ) -> Result<AccountProof, StateProofError> {
         Ok(self
-            .multiproof(HashMap::from_iter([(
+            .multiproof(MultiProofTargets::from_iter([(
                 keccak256(address),
                 slots.iter().map(keccak256).collect(),
             )]))?