From cbf31556a98ba805308e5df83e08e02076190a58 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 27 Feb 2025 13:50:04 +0100 Subject: [PATCH] chore(trie): convert `MultiProofTargets` into unit struct (#14748) --- crates/engine/tree/src/tree/mod.rs | 3 +-- crates/engine/tree/src/tree/root.rs | 13 +++------ crates/trie/common/src/hashed_state.rs | 3 +-- crates/trie/common/src/proofs.rs | 37 ++++++++++++++++++++++++-- crates/trie/db/tests/witness.rs | 17 +++++++++--- crates/trie/trie/src/proof/blinded.rs | 9 +++---- crates/trie/trie/src/proof/mod.rs | 2 +- 7 files changed, 58 insertions(+), 26 deletions(-) diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 8c026d063c..d48bc51c7e 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/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 diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index 4d80a24060..770ffe3cf7 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/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::*; diff --git a/crates/trie/common/src/hashed_state.rs b/crates/trie/common/src/hashed_state.rs index 0065d50187..fb2297207c 100644 --- a/crates/trie/common/src/hashed_state.rs +++ b/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()); } diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index 9eff5c90e3..c712dba2da 100644 --- a/crates/trie/common/src/proofs.rs +++ b/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; +#[derive(Deref, DerefMut, IntoIterator, Clone, PartialEq, Eq, Default, Debug)] +pub struct MultiProofTargets(B256Map); + +impl FromIterator<(B256, B256Set)> for MultiProofTargets { + fn from_iter>(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 diff --git a/crates/trie/db/tests/witness.rs b/crates/trie/db/tests/witness.rs index a73eaae863..5dfa1c3e4a 100644 --- a/crates/trie/db/tests/witness.rs +++ b/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]), )])) diff --git a/crates/trie/trie/src/proof/blinded.rs b/crates/trie/trie/src/proof/blinded.rs index 9e802cf525..e5ac6bf2a0 100644 --- a/crates/trie/trie/src/proof/blinded.rs +++ b/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, 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()) diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs index 6ffcf67cd2..b9cd3a99ee 100644 --- a/crates/trie/trie/src/proof/mod.rs +++ b/crates/trie/trie/src/proof/mod.rs @@ -93,7 +93,7 @@ where slots: &[B256], ) -> Result { Ok(self - .multiproof(HashMap::from_iter([( + .multiproof(MultiProofTargets::from_iter([( keccak256(address), slots.iter().map(keccak256).collect(), )]))?