github/reth
1
1
Fork 0
mirror of https://github.com/paradigmxyz/reth.git synced 2026-01-09 15:28:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

perf(trie): consolidate branch node mask fields into single tuple

Browse Source 
Consolidate separate `branch_node_hash_masks` and `branch_node_tree_masks`
HashMap fields into a single `branch_node_masks: BranchNodeMasksMap` field
across all multiproof structs (MultiProof, StorageMultiProof, DecodedMultiProof,
DecodedStorageMultiProof).

Uses a simple type alias `BranchNodeMasks = (TrieMask, TrieMask)` instead of a
dedicated struct for minimal overhead.

This reduces struct overhead by 50% (80 -> 40 bytes) and improves:
- Insert: 1.7-1.8x faster
- Iterate: ~2x faster
- Lookup: ~12% faster
- Clone: ~2.7x faster
This commit is contained in:
yongkangc
2025-12-29 01:10:28 +00:00
parent 05b3a8668c
commit d0f1190def
8 changed files with 131 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/trie/common/src/lib.rs
View File

@@ -42,7 +42,7 @@ mod subnode;
pub use subnode::StoredSubNode; pub use subnode::StoredSubNode;
mod trie; mod trie;
pub use trie::{ProofTrieNode, TrieMasks}; pub use trie::{BranchNodeMasks, BranchNodeMasksMap, ProofTrieNode, TrieMasks};
/// The implementation of a container for storing intermediate changes to a trie. /// The implementation of a container for storing intermediate changes to a trie.
/// The container indicates when the trie has been modified. /// The container indicates when the trie has been modified.

76
crates/trie/common/src/proofs.rs
View File

@@ -1,18 +1,18 @@
//! Merkle trie proofs. //! Merkle trie proofs.
use crate::{Nibbles, TrieAccount}; use crate::{BranchNodeMasksMap, Nibbles, TrieAccount};
use alloc::{borrow::Cow, vec::Vec}; use alloc::{borrow::Cow, vec::Vec};
use alloy_consensus::constants::KECCAK_EMPTY; use alloy_consensus::constants::KECCAK_EMPTY;
use alloy_primitives::{ use alloy_primitives::{
keccak256, keccak256,
map::{hash_map, B256Map, B256Set, HashMap}, map::{hash_map, B256Map, B256Set},
Address, Bytes, B256, U256, Address, Bytes, B256, U256,
}; };
use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE}; use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE};
use alloy_trie::{ use alloy_trie::{
nodes::TrieNode, nodes::TrieNode,
proof::{verify_proof, DecodedProofNodes, ProofNodes, ProofVerificationError}, proof::{verify_proof, DecodedProofNodes, ProofNodes, ProofVerificationError},
TrieMask, EMPTY_ROOT_HASH, EMPTY_ROOT_HASH,
}; };
use derive_more::{Deref, DerefMut, IntoIterator}; use derive_more::{Deref, DerefMut, IntoIterator};
use itertools::Itertools; use itertools::Itertools;
@@ -173,10 +173,9 @@ impl Iterator for ChunkedMultiProofTargets {
pub struct MultiProof { pub struct MultiProof {
/// State trie multiproof for requested accounts. /// State trie multiproof for requested accounts.
pub account_subtree: ProofNodes, pub account_subtree: ProofNodes,
/// The hash masks of the branch nodes in the account proof. /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the account
pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>, /// proof.
/// The tree masks of the branch nodes in the account proof. pub branch_node_masks: BranchNodeMasksMap,
pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
/// Storage trie multiproofs. /// Storage trie multiproofs.
pub storages: B256Map<StorageMultiProof>, pub storages: B256Map<StorageMultiProof>,
} }
@@ -185,8 +184,7 @@ impl MultiProof {
/// Returns true if the multiproof is empty. /// Returns true if the multiproof is empty.
pub fn is_empty(&self) -> bool { pub fn is_empty(&self) -> bool {
self.account_subtree.is_empty() && self.account_subtree.is_empty() &&
self.branch_node_hash_masks.is_empty() && self.branch_node_masks.is_empty() &&
self.branch_node_tree_masks.is_empty() &&
self.storages.is_empty() self.storages.is_empty()
} }
@@ -267,9 +265,7 @@ impl MultiProof {
/// proofs. /// proofs.
pub fn extend(&mut self, other: Self) { pub fn extend(&mut self, other: Self) {
self.account_subtree.extend_from(other.account_subtree); self.account_subtree.extend_from(other.account_subtree);
self.branch_node_masks.extend(other.branch_node_masks);
self.branch_node_hash_masks.extend(other.branch_node_hash_masks);
self.branch_node_tree_masks.extend(other.branch_node_tree_masks);
for (hashed_address, storage) in other.storages { for (hashed_address, storage) in other.storages {
match self.storages.entry(hashed_address) { match self.storages.entry(hashed_address) {
@@ -277,8 +273,7 @@ impl MultiProof {
debug_assert_eq!(entry.get().root, storage.root); debug_assert_eq!(entry.get().root, storage.root);
let entry = entry.get_mut(); let entry = entry.get_mut();
entry.subtree.extend_from(storage.subtree); entry.subtree.extend_from(storage.subtree);
entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks); entry.branch_node_masks.extend(storage.branch_node_masks);
entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks);
} }
hash_map::Entry::Vacant(entry) => { hash_map::Entry::Vacant(entry) => {
entry.insert(storage); entry.insert(storage);
@@ -302,10 +297,9 @@ impl MultiProof {
pub struct DecodedMultiProof { pub struct DecodedMultiProof {
/// State trie multiproof for requested accounts. /// State trie multiproof for requested accounts.
pub account_subtree: DecodedProofNodes, pub account_subtree: DecodedProofNodes,
/// The hash masks of the branch nodes in the account proof. /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the account
pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>, /// proof.
/// The tree masks of the branch nodes in the account proof. pub branch_node_masks: BranchNodeMasksMap,
pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
/// Storage trie multiproofs. /// Storage trie multiproofs.
pub storages: B256Map<DecodedStorageMultiProof>, pub storages: B256Map<DecodedStorageMultiProof>,
} }
@@ -314,8 +308,7 @@ impl DecodedMultiProof {
/// Returns true if the multiproof is empty. /// Returns true if the multiproof is empty.
pub fn is_empty(&self) -> bool { pub fn is_empty(&self) -> bool {
self.account_subtree.is_empty() && self.account_subtree.is_empty() &&
self.branch_node_hash_masks.is_empty() && self.branch_node_masks.is_empty() &&
self.branch_node_tree_masks.is_empty() &&
self.storages.is_empty() self.storages.is_empty()
} }
@@ -395,9 +388,7 @@ impl DecodedMultiProof {
/// proofs. /// proofs.
pub fn extend(&mut self, other: Self) { pub fn extend(&mut self, other: Self) {
self.account_subtree.extend_from(other.account_subtree); self.account_subtree.extend_from(other.account_subtree);
self.branch_node_masks.extend(other.branch_node_masks);
self.branch_node_hash_masks.extend(other.branch_node_hash_masks);
self.branch_node_tree_masks.extend(other.branch_node_tree_masks);
for (hashed_address, storage) in other.storages { for (hashed_address, storage) in other.storages {
match self.storages.entry(hashed_address) { match self.storages.entry(hashed_address) {
@@ -405,8 +396,7 @@ impl DecodedMultiProof {
debug_assert_eq!(entry.get().root, storage.root); debug_assert_eq!(entry.get().root, storage.root);
let entry = entry.get_mut(); let entry = entry.get_mut();
entry.subtree.extend_from(storage.subtree); entry.subtree.extend_from(storage.subtree);
entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks); entry.branch_node_masks.extend(storage.branch_node_masks);
entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks);
} }
hash_map::Entry::Vacant(entry) => { hash_map::Entry::Vacant(entry) => {
entry.insert(storage); entry.insert(storage);
@@ -437,12 +427,7 @@ impl TryFrom<MultiProof> for DecodedMultiProof {
.into_iter() .into_iter()
.map(|(address, storage)| Ok((address, storage.try_into()?))) .map(|(address, storage)| Ok((address, storage.try_into()?)))
.collect::<Result<B256Map<_>, alloy_rlp::Error>>()?; .collect::<Result<B256Map<_>, alloy_rlp::Error>>()?;
Ok(Self { Ok(Self { account_subtree, branch_node_masks: multi_proof.branch_node_masks, storages })
account_subtree,
branch_node_hash_masks: multi_proof.branch_node_hash_masks,
branch_node_tree_masks: multi_proof.branch_node_tree_masks,
storages,
})
} }
} }
@@ -453,10 +438,9 @@ pub struct StorageMultiProof {
pub root: B256, pub root: B256,
/// Storage multiproof for requested slots. /// Storage multiproof for requested slots.
pub subtree: ProofNodes, pub subtree: ProofNodes,
/// The hash masks of the branch nodes in the storage proof. /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the storage
pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>, /// proof.
/// The tree masks of the branch nodes in the storage proof. pub branch_node_masks: BranchNodeMasksMap,
pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
} }
impl StorageMultiProof { impl StorageMultiProof {
@@ -468,8 +452,7 @@ impl StorageMultiProof {
Nibbles::default(), Nibbles::default(),
Bytes::from([EMPTY_STRING_CODE]), Bytes::from([EMPTY_STRING_CODE]),
)]), )]),
branch_node_hash_masks: HashMap::default(), branch_node_masks: BranchNodeMasksMap::default(),
branch_node_tree_masks: HashMap::default(),
} }
} }
@@ -508,10 +491,9 @@ pub struct DecodedStorageMultiProof {
pub root: B256, pub root: B256,
/// Storage multiproof for requested slots. /// Storage multiproof for requested slots.
pub subtree: DecodedProofNodes, pub subtree: DecodedProofNodes,
/// The hash masks of the branch nodes in the storage proof. /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the storage
pub branch_node_hash_masks: HashMap<Nibbles, TrieMask>, /// proof.
/// The tree masks of the branch nodes in the storage proof. pub branch_node_masks: BranchNodeMasksMap,
pub branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
} }
impl DecodedStorageMultiProof { impl DecodedStorageMultiProof {
@@ -520,8 +502,7 @@ impl DecodedStorageMultiProof {
Self { Self {
root: EMPTY_ROOT_HASH, root: EMPTY_ROOT_HASH,
subtree: DecodedProofNodes::from_iter([(Nibbles::default(), TrieNode::EmptyRoot)]), subtree: DecodedProofNodes::from_iter([(Nibbles::default(), TrieNode::EmptyRoot)]),
branch_node_hash_masks: HashMap::default(), branch_node_masks: BranchNodeMasksMap::default(),
branch_node_tree_masks: HashMap::default(),
} }
} }
@@ -560,8 +541,7 @@ impl TryFrom<StorageMultiProof> for DecodedStorageMultiProof {
Ok(Self { Ok(Self {
root: multi_proof.root, root: multi_proof.root,
subtree, subtree,
branch_node_hash_masks: multi_proof.branch_node_hash_masks, branch_node_masks: multi_proof.branch_node_masks,
branch_node_tree_masks: multi_proof.branch_node_tree_masks,
}) })
} }
} }
@@ -910,8 +890,7 @@ mod tests {
StorageMultiProof { StorageMultiProof {
root, root,
subtree: subtree1, subtree: subtree1,
branch_node_hash_masks: HashMap::default(), branch_node_masks: BranchNodeMasksMap::default(),
branch_node_tree_masks: HashMap::default(),
}, },
); );
@@ -925,8 +904,7 @@ mod tests {
StorageMultiProof { StorageMultiProof {
root, root,
subtree: subtree2, subtree: subtree2,
branch_node_hash_masks: HashMap::default(), branch_node_masks: BranchNodeMasksMap::default(),
branch_node_tree_masks: HashMap::default(),
}, },
); );

17
crates/trie/common/src/trie.rs
View File

@@ -1,8 +1,25 @@
//! Types related to sparse trie nodes and masks. //! Types related to sparse trie nodes and masks.
use crate::Nibbles; use crate::Nibbles;
use alloy_primitives::map::HashMap;
use alloy_trie::{nodes::TrieNode, TrieMask}; use alloy_trie::{nodes::TrieNode, TrieMask};
/// Branch node masks containing `hash_mask` and `tree_mask`.
///
/// Consolidates `hash_mask` and `tree_mask` into a single struct, reducing `HashMap` overhead
/// when storing masks by path. Instead of two separate `HashMap<Nibbles, TrieMask>`,
/// we use a single `HashMap<Nibbles, BranchNodeMasks>`.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct BranchNodeMasks {
/// Hash mask indicating which children are stored as hashes.
pub hash_mask: TrieMask,
/// Tree mask indicating which children are complete subtrees.
pub tree_mask: TrieMask,
}
/// Map from nibble path to branch node masks.
pub type BranchNodeMasksMap = HashMap<Nibbles, BranchNodeMasks>;
/// Struct for passing around branch node mask information. /// Struct for passing around branch node mask information.
/// ///
/// Branch nodes can have up to 16 children (one for each nibble). /// Branch nodes can have up to 16 children (one for each nibble).

18
crates/trie/parallel/src/proof_task.rs
View File

@@ -58,6 +58,7 @@ use reth_trie_common::{
added_removed_keys::MultiAddedRemovedKeys, added_removed_keys::MultiAddedRemovedKeys,
prefix_set::{PrefixSet, PrefixSetMut}, prefix_set::{PrefixSet, PrefixSetMut},
proof::{DecodedProofNodes, ProofRetainer}, proof::{DecodedProofNodes, ProofRetainer},
BranchNodeMasks, BranchNodeMasksMap,
}; };
use reth_trie_sparse::provider::{RevealedNode, TrieNodeProvider, TrieNodeProviderFactory}; use reth_trie_sparse::provider::{RevealedNode, TrieNodeProvider, TrieNodeProviderFactory};
use std::{ use std::{
@@ -1421,14 +1422,16 @@ where
let account_subtree_raw_nodes = hash_builder.take_proof_nodes(); let account_subtree_raw_nodes = hash_builder.take_proof_nodes();
let decoded_account_subtree = DecodedProofNodes::try_from(account_subtree_raw_nodes)?; let decoded_account_subtree = DecodedProofNodes::try_from(account_subtree_raw_nodes)?;
let (branch_node_hash_masks, branch_node_tree_masks) = if ctx.collect_branch_node_masks { let branch_node_masks = if ctx.collect_branch_node_masks {
let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default(); let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
( updated_branch_nodes
updated_branch_nodes.iter().map(|(path, node)| (*path, node.hash_mask)).collect(), .into_iter()
updated_branch_nodes.into_iter().map(|(path, node)| (path, node.tree_mask)).collect(), .map(|(path, node)| {
) (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask })
})
.collect()
} else { } else {
(Default::default(), Default::default()) BranchNodeMasksMap::default()
}; };
// Extend tracker with accumulated metrics from account cursors // Extend tracker with accumulated metrics from account cursors
@@ -1437,8 +1440,7 @@ where
Ok(DecodedMultiProof { Ok(DecodedMultiProof {
account_subtree: decoded_account_subtree, account_subtree: decoded_account_subtree,
branch_node_hash_masks, branch_node_masks,
branch_node_tree_masks,
storages: collected_decoded_storages, storages: collected_decoded_storages,
}) })
} }

96
crates/trie/sparse/src/state.rs
View File

@@ -5,7 +5,7 @@ use crate::{
}; };
use alloc::{collections::VecDeque, vec::Vec}; use alloc::{collections::VecDeque, vec::Vec};
use alloy_primitives::{ use alloy_primitives::{
map::{B256Map, HashMap, HashSet}, map::{B256Map, HashSet},
Bytes, B256, Bytes, B256,
}; };
use alloy_rlp::{Decodable, Encodable}; use alloy_rlp::{Decodable, Encodable};
@@ -15,8 +15,8 @@ use reth_primitives_traits::Account;
use reth_trie_common::{ use reth_trie_common::{
proof::ProofNodes, proof::ProofNodes,
updates::{StorageTrieUpdates, TrieUpdates}, updates::{StorageTrieUpdates, TrieUpdates},
DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles, ProofTrieNode, RlpNode, BranchNodeMasksMap, DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles,
StorageMultiProof, TrieAccount, TrieMask, TrieMasks, TrieNode, EMPTY_ROOT_HASH, ProofTrieNode, RlpNode, StorageMultiProof, TrieAccount, TrieMasks, TrieNode, EMPTY_ROOT_HASH,
TRIE_ACCOUNT_RLP_MAX_SIZE, TRIE_ACCOUNT_RLP_MAX_SIZE,
}; };
use tracing::{instrument, trace}; use tracing::{instrument, trace};
@@ -247,19 +247,10 @@ where
&mut self, &mut self,
multiproof: DecodedMultiProof, multiproof: DecodedMultiProof,
) -> SparseStateTrieResult<()> { ) -> SparseStateTrieResult<()> {
let DecodedMultiProof { let DecodedMultiProof { account_subtree, storages, branch_node_masks } = multiproof;
account_subtree,
storages,
branch_node_hash_masks,
branch_node_tree_masks,
} = multiproof;
// first reveal the account proof nodes // first reveal the account proof nodes
self.reveal_decoded_account_multiproof( self.reveal_decoded_account_multiproof(account_subtree, branch_node_masks)?;
account_subtree,
branch_node_hash_masks,
branch_node_tree_masks,
)?;
#[cfg(not(feature = "std"))] #[cfg(not(feature = "std"))]
// If nostd then serially reveal storage proof nodes for each storage trie // If nostd then serially reveal storage proof nodes for each storage trie
@@ -329,31 +320,24 @@ where
pub fn reveal_account_multiproof( pub fn reveal_account_multiproof(
&mut self, &mut self,
account_subtree: ProofNodes, account_subtree: ProofNodes,
branch_node_hash_masks: HashMap<Nibbles, TrieMask>, branch_node_masks: BranchNodeMasksMap,
branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
) -> SparseStateTrieResult<()> { ) -> SparseStateTrieResult<()> {
// decode the multiproof first // decode the multiproof first
let decoded_multiproof = account_subtree.try_into()?; let decoded_multiproof = account_subtree.try_into()?;
self.reveal_decoded_account_multiproof( self.reveal_decoded_account_multiproof(decoded_multiproof, branch_node_masks)
decoded_multiproof,
branch_node_hash_masks,
branch_node_tree_masks,
)
} }
/// Reveals a decoded account multiproof. /// Reveals a decoded account multiproof.
pub fn reveal_decoded_account_multiproof( pub fn reveal_decoded_account_multiproof(
&mut self, &mut self,
account_subtree: DecodedProofNodes, account_subtree: DecodedProofNodes,
branch_node_hash_masks: HashMap<Nibbles, TrieMask>, branch_node_masks: BranchNodeMasksMap,
branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
) -> SparseStateTrieResult<()> { ) -> SparseStateTrieResult<()> {
let FilterMappedProofNodes { root_node, nodes, new_nodes, metric_values: _metric_values } = let FilterMappedProofNodes { root_node, nodes, new_nodes, metric_values: _metric_values } =
filter_map_revealed_nodes( filter_map_revealed_nodes(
account_subtree, account_subtree,
&mut self.revealed_account_paths, &mut self.revealed_account_paths,
&branch_node_hash_masks, &branch_node_masks,
&branch_node_tree_masks,
)?; )?;
#[cfg(feature = "metrics")] #[cfg(feature = "metrics")]
{ {
@@ -426,8 +410,7 @@ where
filter_map_revealed_nodes( filter_map_revealed_nodes(
storage_subtree.subtree, storage_subtree.subtree,
revealed_nodes, revealed_nodes,
&storage_subtree.branch_node_hash_masks, &storage_subtree.branch_node_masks,
&storage_subtree.branch_node_tree_masks,
)?; )?;
if let Some(root_node) = root_node { if let Some(root_node) = root_node {
@@ -959,8 +942,7 @@ struct FilterMappedProofNodes {
fn filter_map_revealed_nodes( fn filter_map_revealed_nodes(
proof_nodes: DecodedProofNodes, proof_nodes: DecodedProofNodes,
revealed_nodes: &mut HashSet<Nibbles>, revealed_nodes: &mut HashSet<Nibbles>,
branch_node_hash_masks: &HashMap<Nibbles, TrieMask>, branch_node_masks: &BranchNodeMasksMap,
branch_node_tree_masks: &HashMap<Nibbles, TrieMask>,
) -> SparseStateTrieResult<FilterMappedProofNodes> { ) -> SparseStateTrieResult<FilterMappedProofNodes> {
let mut result = FilterMappedProofNodes { let mut result = FilterMappedProofNodes {
root_node: None, root_node: None,
@@ -991,9 +973,13 @@ fn filter_map_revealed_nodes(
// If it's a branch node, increase the number of new nodes by the number of children // If it's a branch node, increase the number of new nodes by the number of children
// according to the state mask. // according to the state mask.
result.new_nodes += branch.state_mask.count_ones() as usize; result.new_nodes += branch.state_mask.count_ones() as usize;
TrieMasks { if let Some(branch_masks) = branch_node_masks.get(&path) {
hash_mask: branch_node_hash_masks.get(&path).copied(), TrieMasks {
tree_mask: branch_node_tree_masks.get(&path).copied(), hash_mask: Some(branch_masks.hash_mask),
tree_mask: Some(branch_masks.tree_mask),
}
} else {
TrieMasks::none()
} }
} }
TrieNode::Extension(_) => { TrieNode::Extension(_) => {
@@ -1042,7 +1028,7 @@ mod tests {
use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH}; use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH};
use reth_trie_common::{ use reth_trie_common::{
proof::{ProofNodes, ProofRetainer}, proof::{ProofNodes, ProofRetainer},
BranchNode, LeafNode, StorageMultiProof, TrieMask, BranchNode, BranchNodeMasks, LeafNode, StorageMultiProof, TrieMask,
}; };
#[test] #[test]
@@ -1149,8 +1135,7 @@ mod tests {
(Nibbles::from_nibbles([0x0]), leaf_1.clone().into()), (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
(Nibbles::from_nibbles([0x1]), leaf_1.clone().into()), (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
]), ]),
branch_node_hash_masks: Default::default(), branch_node_masks: Default::default(),
branch_node_tree_masks: Default::default(),
}, },
)]), )]),
..Default::default() ..Default::default()
@@ -1227,11 +1212,19 @@ mod tests {
storage_hash_builder.add_leaf(slot_path_1, &alloy_rlp::encode_fixed_size(&value_1)); storage_hash_builder.add_leaf(slot_path_1, &alloy_rlp::encode_fixed_size(&value_1));
storage_hash_builder.add_leaf(slot_path_2, &alloy_rlp::encode_fixed_size(&value_2)); storage_hash_builder.add_leaf(slot_path_2, &alloy_rlp::encode_fixed_size(&value_2));
use alloy_trie::TrieMask;
let storage_root = storage_hash_builder.root(); let storage_root = storage_hash_builder.root();
let storage_proof_nodes = storage_hash_builder.take_proof_nodes(); let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
let storage_branch_node_hash_masks = HashMap::from_iter([ let storage_branch_node_masks = BranchNodeMasksMap::from_iter([
(Nibbles::default(), TrieMask::new(0b010)), (
(Nibbles::from_nibbles([0x1]), TrieMask::new(0b11)), Nibbles::default(),
BranchNodeMasks { hash_mask: TrieMask::new(0b010), tree_mask: TrieMask::default() },
),
(
Nibbles::from_nibbles([0x1]),
BranchNodeMasks { hash_mask: TrieMask::new(0b11), tree_mask: TrieMask::default() },
),
]); ]);
let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000"); let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
@@ -1257,19 +1250,20 @@ mod tests {
.reveal_decoded_multiproof( .reveal_decoded_multiproof(
MultiProof { MultiProof {
account_subtree: proof_nodes, account_subtree: proof_nodes,
branch_node_hash_masks: HashMap::from_iter([( branch_node_masks: BranchNodeMasksMap::from_iter([(
Nibbles::from_nibbles([0x1]), Nibbles::from_nibbles([0x1]),
TrieMask::new(0b00), BranchNodeMasks {
hash_mask: TrieMask::new(0b00),
tree_mask: TrieMask::default(),
},
)]), )]),
branch_node_tree_masks: HashMap::default(),
storages: HashMap::from_iter([ storages: HashMap::from_iter([
( (
address_1, address_1,
StorageMultiProof { StorageMultiProof {
root, root,
subtree: storage_proof_nodes.clone(), subtree: storage_proof_nodes.clone(),
branch_node_hash_masks: storage_branch_node_hash_masks.clone(), branch_node_masks: storage_branch_node_masks.clone(),
branch_node_tree_masks: HashMap::default(),
}, },
), ),
( (
@@ -1277,8 +1271,7 @@ mod tests {
StorageMultiProof { StorageMultiProof {
root, root,
subtree: storage_proof_nodes, subtree: storage_proof_nodes,
branch_node_hash_masks: storage_branch_node_hash_masks, branch_node_masks: storage_branch_node_masks,
branch_node_tree_masks: HashMap::default(),
}, },
), ),
]), ]),
@@ -1366,16 +1359,11 @@ mod tests {
(Nibbles::from_nibbles([0x1]), leaf.clone()), (Nibbles::from_nibbles([0x1]), leaf.clone()),
]); ]);
let branch_node_hash_masks = HashMap::default(); let branch_node_masks = BranchNodeMasksMap::default();
let branch_node_tree_masks = HashMap::default();
let decoded = filter_map_revealed_nodes( let decoded =
proof_nodes, filter_map_revealed_nodes(proof_nodes, &mut revealed_nodes, &branch_node_masks)
&mut revealed_nodes, .unwrap();
&branch_node_hash_masks,
&branch_node_tree_masks,
)
.unwrap();
assert_eq!( assert_eq!(
decoded, decoded,

43
crates/trie/trie/src/proof/mod.rs
View File

@@ -11,14 +11,15 @@ use crate::{
}; };
use alloy_primitives::{ use alloy_primitives::{
keccak256, keccak256,
map::{B256Map, B256Set, HashMap, HashSet}, map::{B256Map, B256Set, HashSet},
Address, B256, Address, B256,
}; };
use alloy_rlp::{BufMut, Encodable}; use alloy_rlp::{BufMut, Encodable};
use alloy_trie::proof::AddedRemovedKeys; use alloy_trie::proof::AddedRemovedKeys;
use reth_execution_errors::trie::StateProofError; use reth_execution_errors::trie::StateProofError;
use reth_trie_common::{ use reth_trie_common::{
proof::ProofRetainer, AccountProof, MultiProof, MultiProofTargets, StorageMultiProof, proof::ProofRetainer, AccountProof, BranchNodeMasks, BranchNodeMasksMap, MultiProof,
MultiProofTargets, StorageMultiProof,
}; };
mod trie_node; mod trie_node;
@@ -180,20 +181,19 @@ where
} }
let _ = hash_builder.root(); let _ = hash_builder.root();
let account_subtree = hash_builder.take_proof_nodes(); let account_subtree = hash_builder.take_proof_nodes();
let (branch_node_hash_masks, branch_node_tree_masks) = if self.collect_branch_node_masks { let branch_node_masks = if self.collect_branch_node_masks {
let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default(); let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
( updated_branch_nodes
updated_branch_nodes.iter().map(|(path, node)| (*path, node.hash_mask)).collect(), .into_iter()
updated_branch_nodes .map(|(path, node)| {
.into_iter() (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask })
.map(|(path, node)| (path, node.tree_mask)) })
.collect(), .collect()
)
} else { } else {
(HashMap::default(), HashMap::default()) BranchNodeMasksMap::default()
}; };
Ok(MultiProof { account_subtree, branch_node_hash_masks, branch_node_tree_masks, storages }) Ok(MultiProof { account_subtree, branch_node_masks, storages })
} }
} }
@@ -398,19 +398,18 @@ where
let root = hash_builder.root(); let root = hash_builder.root();
let subtree = hash_builder.take_proof_nodes(); let subtree = hash_builder.take_proof_nodes();
let (branch_node_hash_masks, branch_node_tree_masks) = if self.collect_branch_node_masks { let branch_node_masks = if self.collect_branch_node_masks {
let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default(); let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default();
( updated_branch_nodes
updated_branch_nodes.iter().map(|(path, node)| (*path, node.hash_mask)).collect(), .into_iter()
updated_branch_nodes .map(|(path, node)| {
.into_iter() (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask })
.map(|(path, node)| (path, node.tree_mask)) })
.collect(), .collect()
)
} else { } else {
(HashMap::default(), HashMap::default()) BranchNodeMasksMap::default()
}; };
Ok(StorageMultiProof { root, subtree, branch_node_hash_masks, branch_node_tree_masks }) Ok(StorageMultiProof { root, subtree, branch_node_masks })
} }
} }

10
crates/trie/trie/src/proof/trie_node.rs
View File

@@ -79,8 +79,9 @@ where
.multiproof(targets) .multiproof(targets)
.map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?;
let node = proof.account_subtree.into_inner().remove(path); let node = proof.account_subtree.into_inner().remove(path);
let tree_mask = proof.branch_node_tree_masks.remove(path); let masks = proof.branch_node_masks.remove(path);
let hash_mask = proof.branch_node_hash_masks.remove(path); let hash_mask = masks.map(|m| m.hash_mask);
let tree_mask = masks.map(|m| m.tree_mask);
trace!( trace!(
target: "trie::proof::blinded", target: "trie::proof::blinded",
@@ -131,8 +132,9 @@ where
.storage_multiproof(targets) .storage_multiproof(targets)
.map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?;
let node = proof.subtree.into_inner().remove(path); let node = proof.subtree.into_inner().remove(path);
let tree_mask = proof.branch_node_tree_masks.remove(path); let masks = proof.branch_node_masks.remove(path);
let hash_mask = proof.branch_node_hash_masks.remove(path); let hash_mask = masks.map(|m| m.hash_mask);
let tree_mask = masks.map(|m| m.tree_mask);
trace!( trace!(
target: "trie::proof::blinded", target: "trie::proof::blinded",

16
crates/trie/trie/src/proof_v2/mod.rs
View File

@@ -1648,17 +1648,15 @@ mod tests {
// though we never store the root node so the masks for it aren't really valid. // though we never store the root node so the masks for it aren't really valid.
let masks = if path.is_empty() { let masks = if path.is_empty() {
TrieMasks::none() TrieMasks::none()
} else { } else if let Some(branch_masks) =
proof_legacy_result.branch_node_masks.get(path)
{
TrieMasks { TrieMasks {
hash_mask: proof_legacy_result hash_mask: Some(branch_masks.hash_mask),
.branch_node_hash_masks tree_mask: Some(branch_masks.tree_mask),
.get(path)
.copied(),
tree_mask: proof_legacy_result
.branch_node_tree_masks
.get(path)
.copied(),
} }
} else {
TrieMasks::none()
}; };
ProofTrieNode { path: *path, node, masks } ProofTrieNode { path: *path, node, masks }