From 008cb25d004239886979e2773414ee51cbd90d1f Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com> Date: Fri, 13 Dec 2024 09:02:39 +0000 Subject: [PATCH] chore(trie): make `MultiProofTargets` a type alias (#13364) --- crates/engine/tree/src/tree/root.rs | 35 +++++++++++++++++++++------ crates/trie/common/src/proofs.rs | 20 +-------------- crates/trie/db/tests/witness.rs | 19 +++++---------- crates/trie/sparse/src/state.rs | 3 +-- crates/trie/trie/src/proof/blinded.rs | 2 +- crates/trie/trie/src/proof/mod.rs | 8 +++--- 6 files changed, 40 insertions(+), 47 deletions(-) diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index 5fc042f6c8..f9d16e0fe4 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/crates/engine/tree/src/tree/root.rs @@ -326,7 +326,7 @@ where let hashed_state_update = evm_state_to_hashed_post_state(update); let proof_targets = get_proof_targets(&hashed_state_update, fetched_proof_targets); - fetched_proof_targets.extend_ref(&proof_targets); + extend_multi_proof_targets_ref(fetched_proof_targets, &proof_targets); // Dispatch proof gathering for this state update scope.spawn(move |_| { @@ -379,12 +379,16 @@ where None } else { // Merge all ready proofs and state updates - ready_proofs.into_iter().reduce(|mut acc, (state_update, targets, proof)| { - acc.0.extend(state_update); - acc.1.extend(targets); - acc.2.extend(proof); - acc - }) + ready_proofs.into_iter().reduce( + |(mut acc_state_update, mut acc_targets, mut acc_proof), + (state_update, targets, proof)| { + acc_state_update.extend(state_update); + extend_multi_proof_targets(&mut acc_targets, targets); + acc_proof.extend(proof); + + (acc_state_update, acc_targets, acc_proof) + }, + ) } } @@ -483,7 +487,10 @@ where ) { if self.sparse_trie.is_none() { current_state_update.extend(combined_state_update); - current_proof_targets.extend(combined_proof_targets); + extend_multi_proof_targets( + &mut current_proof_targets, + combined_proof_targets, + ); current_multiproof.extend(combined_proof); } else { self.spawn_root_calculation( @@ -683,6 +690,18 @@ fn update_sparse_trie< Ok((trie, elapsed)) } +fn extend_multi_proof_targets(targets: &mut MultiProofTargets, other: MultiProofTargets) { + for (address, slots) in other { + targets.entry(address).or_default().extend(slots); + } +} + +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/proofs.rs b/crates/trie/common/src/proofs.rs index 8455d1e8ac..eb3626d90e 100644 --- a/crates/trie/common/src/proofs.rs +++ b/crates/trie/common/src/proofs.rs @@ -13,29 +13,11 @@ use alloy_trie::{ proof::{verify_proof, ProofNodes, ProofVerificationError}, TrieMask, EMPTY_ROOT_HASH, }; -use derive_more::derive::{Deref, DerefMut, From, Into, IntoIterator}; use itertools::Itertools; use reth_primitives_traits::Account; /// Proof targets map. -#[derive(Debug, Default, Clone, Deref, DerefMut, From, Into, IntoIterator)] -pub struct MultiProofTargets(B256HashMap); - -impl MultiProofTargets { - /// Extends the proof targets map with another one. - pub fn extend(&mut self, other: Self) { - for (address, slots) in other.0 { - self.0.entry(address).or_default().extend(slots); - } - } - - /// Extends the proof targets map with another one by reference. - pub fn extend_ref(&mut self, other: &Self) { - for (address, slots) in &other.0 { - self.0.entry(*address).or_default().extend(slots); - } - } -} +pub type MultiProofTargets = B256HashMap; /// 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 c9732bef49..385f6269f3 100644 --- a/crates/trie/db/tests/witness.rs +++ b/crates/trie/db/tests/witness.rs @@ -39,9 +39,7 @@ 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]))]).into(), - ) + .multiproof(HashMap::from_iter([(hashed_address, HashSet::from_iter([hashed_slot]))])) .unwrap(); let witness = TrieWitness::from_tx(provider.tx_ref()) @@ -79,9 +77,7 @@ 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]))]).into(), - ) + .multiproof(HashMap::from_iter([(hashed_address, HashSet::from_iter([hashed_slot]))])) .unwrap(); let witness = @@ -126,13 +122,10 @@ 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([( - hashed_address, - HashSet::from_iter([hashed_slot1, hashed_slot2]), - )]) - .into(), - ) + .multiproof(HashMap::from_iter([( + hashed_address, + HashSet::from_iter([hashed_slot1, hashed_slot2]), + )])) .unwrap(); let witness = TrieWitness::from_tx(provider.tx_ref()) diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 5112e64a22..c6c741e002 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -560,8 +560,7 @@ mod tests { HashMap::from_iter([ (address_1, HashSet::from_iter([slot_1, slot_2])), (address_2, HashSet::from_iter([slot_1, slot_2])), - ]) - .into(), + ]), MultiProof { account_subtree: proof_nodes, branch_node_hash_masks: HashMap::from_iter([( diff --git a/crates/trie/trie/src/proof/blinded.rs b/crates/trie/trie/src/proof/blinded.rs index eb713c05aa..33a1a43b57 100644 --- a/crates/trie/trie/src/proof/blinded.rs +++ b/crates/trie/trie/src/proof/blinded.rs @@ -91,7 +91,7 @@ where let proof = Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone()) .with_prefix_sets_mut(self.prefix_sets.as_ref().clone()) - .multiproof(targets.into()) + .multiproof(targets) .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; Ok(proof.account_subtree.into_inner().remove(path)) diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs index fe31f7baa9..7467785cc8 100644 --- a/crates/trie/trie/src/proof/mod.rs +++ b/crates/trie/trie/src/proof/mod.rs @@ -94,10 +94,10 @@ where slots: &[B256], ) -> Result { Ok(self - .multiproof( - HashMap::from_iter([(keccak256(address), slots.iter().map(keccak256).collect())]) - .into(), - )? + .multiproof(HashMap::from_iter([( + keccak256(address), + slots.iter().map(keccak256).collect(), + )]))? .account_proof(address, slots)?) }