From d0f1190def4b66f742db4da248210a13c3a5dc78 Mon Sep 17 00:00:00 2001 From: yongkangc Date: Mon, 29 Dec 2025 01:10:28 +0000 Subject: [PATCH] perf(trie): consolidate branch node mask fields into single tuple 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 --- crates/trie/common/src/lib.rs | 2 +- crates/trie/common/src/proofs.rs | 76 +++++++------------- crates/trie/common/src/trie.rs | 17 +++++ crates/trie/parallel/src/proof_task.rs | 18 ++--- crates/trie/sparse/src/state.rs | 96 +++++++++++-------------- crates/trie/trie/src/proof/mod.rs | 43 ++++++----- crates/trie/trie/src/proof/trie_node.rs | 10 +-- crates/trie/trie/src/proof_v2/mod.rs | 16 ++--- 8 files changed, 131 insertions(+), 147 deletions(-) diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index f212dd2910..4a80be284c 100644 --- a/crates/trie/common/src/lib.rs +++ b/crates/trie/common/src/lib.rs @@ -42,7 +42,7 @@ mod subnode; pub use subnode::StoredSubNode; 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 container indicates when the trie has been modified. diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index a8e0bb59b9..4aa4a6200b 100644 --- a/crates/trie/common/src/proofs.rs +++ b/crates/trie/common/src/proofs.rs @@ -1,18 +1,18 @@ //! Merkle trie proofs. -use crate::{Nibbles, TrieAccount}; +use crate::{BranchNodeMasksMap, Nibbles, TrieAccount}; use alloc::{borrow::Cow, vec::Vec}; use alloy_consensus::constants::KECCAK_EMPTY; use alloy_primitives::{ keccak256, - map::{hash_map, B256Map, B256Set, HashMap}, + map::{hash_map, B256Map, B256Set}, Address, Bytes, B256, U256, }; use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE}; use alloy_trie::{ nodes::TrieNode, proof::{verify_proof, DecodedProofNodes, ProofNodes, ProofVerificationError}, - TrieMask, EMPTY_ROOT_HASH, + EMPTY_ROOT_HASH, }; use derive_more::{Deref, DerefMut, IntoIterator}; use itertools::Itertools; @@ -173,10 +173,9 @@ impl Iterator for ChunkedMultiProofTargets { pub struct MultiProof { /// State trie multiproof for requested accounts. pub account_subtree: ProofNodes, - /// The hash masks of the branch nodes in the account proof. - pub branch_node_hash_masks: HashMap, - /// The tree masks of the branch nodes in the account proof. - pub branch_node_tree_masks: HashMap, + /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the account + /// proof. + pub branch_node_masks: BranchNodeMasksMap, /// Storage trie multiproofs. pub storages: B256Map, } @@ -185,8 +184,7 @@ impl MultiProof { /// Returns true if the multiproof is empty. pub fn is_empty(&self) -> bool { self.account_subtree.is_empty() && - self.branch_node_hash_masks.is_empty() && - self.branch_node_tree_masks.is_empty() && + self.branch_node_masks.is_empty() && self.storages.is_empty() } @@ -267,9 +265,7 @@ impl MultiProof { /// proofs. pub fn extend(&mut self, other: Self) { self.account_subtree.extend_from(other.account_subtree); - - self.branch_node_hash_masks.extend(other.branch_node_hash_masks); - self.branch_node_tree_masks.extend(other.branch_node_tree_masks); + self.branch_node_masks.extend(other.branch_node_masks); for (hashed_address, storage) in other.storages { match self.storages.entry(hashed_address) { @@ -277,8 +273,7 @@ impl MultiProof { debug_assert_eq!(entry.get().root, storage.root); let entry = entry.get_mut(); entry.subtree.extend_from(storage.subtree); - entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks); - entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks); + entry.branch_node_masks.extend(storage.branch_node_masks); } hash_map::Entry::Vacant(entry) => { entry.insert(storage); @@ -302,10 +297,9 @@ impl MultiProof { pub struct DecodedMultiProof { /// State trie multiproof for requested accounts. pub account_subtree: DecodedProofNodes, - /// The hash masks of the branch nodes in the account proof. - pub branch_node_hash_masks: HashMap, - /// The tree masks of the branch nodes in the account proof. - pub branch_node_tree_masks: HashMap, + /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the account + /// proof. + pub branch_node_masks: BranchNodeMasksMap, /// Storage trie multiproofs. pub storages: B256Map, } @@ -314,8 +308,7 @@ impl DecodedMultiProof { /// Returns true if the multiproof is empty. pub fn is_empty(&self) -> bool { self.account_subtree.is_empty() && - self.branch_node_hash_masks.is_empty() && - self.branch_node_tree_masks.is_empty() && + self.branch_node_masks.is_empty() && self.storages.is_empty() } @@ -395,9 +388,7 @@ impl DecodedMultiProof { /// proofs. pub fn extend(&mut self, other: Self) { self.account_subtree.extend_from(other.account_subtree); - - self.branch_node_hash_masks.extend(other.branch_node_hash_masks); - self.branch_node_tree_masks.extend(other.branch_node_tree_masks); + self.branch_node_masks.extend(other.branch_node_masks); for (hashed_address, storage) in other.storages { match self.storages.entry(hashed_address) { @@ -405,8 +396,7 @@ impl DecodedMultiProof { debug_assert_eq!(entry.get().root, storage.root); let entry = entry.get_mut(); entry.subtree.extend_from(storage.subtree); - entry.branch_node_hash_masks.extend(storage.branch_node_hash_masks); - entry.branch_node_tree_masks.extend(storage.branch_node_tree_masks); + entry.branch_node_masks.extend(storage.branch_node_masks); } hash_map::Entry::Vacant(entry) => { entry.insert(storage); @@ -437,12 +427,7 @@ impl TryFrom for DecodedMultiProof { .into_iter() .map(|(address, storage)| Ok((address, storage.try_into()?))) .collect::, alloy_rlp::Error>>()?; - Ok(Self { - account_subtree, - branch_node_hash_masks: multi_proof.branch_node_hash_masks, - branch_node_tree_masks: multi_proof.branch_node_tree_masks, - storages, - }) + Ok(Self { account_subtree, branch_node_masks: multi_proof.branch_node_masks, storages }) } } @@ -453,10 +438,9 @@ pub struct StorageMultiProof { pub root: B256, /// Storage multiproof for requested slots. pub subtree: ProofNodes, - /// The hash masks of the branch nodes in the storage proof. - pub branch_node_hash_masks: HashMap, - /// The tree masks of the branch nodes in the storage proof. - pub branch_node_tree_masks: HashMap, + /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the storage + /// proof. + pub branch_node_masks: BranchNodeMasksMap, } impl StorageMultiProof { @@ -468,8 +452,7 @@ impl StorageMultiProof { Nibbles::default(), Bytes::from([EMPTY_STRING_CODE]), )]), - branch_node_hash_masks: HashMap::default(), - branch_node_tree_masks: HashMap::default(), + branch_node_masks: BranchNodeMasksMap::default(), } } @@ -508,10 +491,9 @@ pub struct DecodedStorageMultiProof { pub root: B256, /// Storage multiproof for requested slots. pub subtree: DecodedProofNodes, - /// The hash masks of the branch nodes in the storage proof. - pub branch_node_hash_masks: HashMap, - /// The tree masks of the branch nodes in the storage proof. - pub branch_node_tree_masks: HashMap, + /// Consolidated branch node masks (`hash_mask`, `tree_mask`) for each path in the storage + /// proof. + pub branch_node_masks: BranchNodeMasksMap, } impl DecodedStorageMultiProof { @@ -520,8 +502,7 @@ impl DecodedStorageMultiProof { Self { root: EMPTY_ROOT_HASH, subtree: DecodedProofNodes::from_iter([(Nibbles::default(), TrieNode::EmptyRoot)]), - branch_node_hash_masks: HashMap::default(), - branch_node_tree_masks: HashMap::default(), + branch_node_masks: BranchNodeMasksMap::default(), } } @@ -560,8 +541,7 @@ impl TryFrom for DecodedStorageMultiProof { Ok(Self { root: multi_proof.root, subtree, - branch_node_hash_masks: multi_proof.branch_node_hash_masks, - branch_node_tree_masks: multi_proof.branch_node_tree_masks, + branch_node_masks: multi_proof.branch_node_masks, }) } } @@ -910,8 +890,7 @@ mod tests { StorageMultiProof { root, subtree: subtree1, - branch_node_hash_masks: HashMap::default(), - branch_node_tree_masks: HashMap::default(), + branch_node_masks: BranchNodeMasksMap::default(), }, ); @@ -925,8 +904,7 @@ mod tests { StorageMultiProof { root, subtree: subtree2, - branch_node_hash_masks: HashMap::default(), - branch_node_tree_masks: HashMap::default(), + branch_node_masks: BranchNodeMasksMap::default(), }, ); diff --git a/crates/trie/common/src/trie.rs b/crates/trie/common/src/trie.rs index 8794839301..1d2706c858 100644 --- a/crates/trie/common/src/trie.rs +++ b/crates/trie/common/src/trie.rs @@ -1,8 +1,25 @@ //! Types related to sparse trie nodes and masks. use crate::Nibbles; +use alloy_primitives::map::HashMap; 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`, +/// we use a single `HashMap`. +#[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; + /// Struct for passing around branch node mask information. /// /// Branch nodes can have up to 16 children (one for each nibble). diff --git a/crates/trie/parallel/src/proof_task.rs b/crates/trie/parallel/src/proof_task.rs index 58dc99fc37..73e91e98fc 100644 --- a/crates/trie/parallel/src/proof_task.rs +++ b/crates/trie/parallel/src/proof_task.rs @@ -58,6 +58,7 @@ use reth_trie_common::{ added_removed_keys::MultiAddedRemovedKeys, prefix_set::{PrefixSet, PrefixSetMut}, proof::{DecodedProofNodes, ProofRetainer}, + BranchNodeMasks, BranchNodeMasksMap, }; use reth_trie_sparse::provider::{RevealedNode, TrieNodeProvider, TrieNodeProviderFactory}; use std::{ @@ -1421,14 +1422,16 @@ where let account_subtree_raw_nodes = hash_builder.take_proof_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(); - ( - updated_branch_nodes.iter().map(|(path, node)| (*path, node.hash_mask)).collect(), - updated_branch_nodes.into_iter().map(|(path, node)| (path, node.tree_mask)).collect(), - ) + updated_branch_nodes + .into_iter() + .map(|(path, node)| { + (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask }) + }) + .collect() } else { - (Default::default(), Default::default()) + BranchNodeMasksMap::default() }; // Extend tracker with accumulated metrics from account cursors @@ -1437,8 +1440,7 @@ where Ok(DecodedMultiProof { account_subtree: decoded_account_subtree, - branch_node_hash_masks, - branch_node_tree_masks, + branch_node_masks, storages: collected_decoded_storages, }) } diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index f02c748430..fc6a7710b2 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -5,7 +5,7 @@ use crate::{ }; use alloc::{collections::VecDeque, vec::Vec}; use alloy_primitives::{ - map::{B256Map, HashMap, HashSet}, + map::{B256Map, HashSet}, Bytes, B256, }; use alloy_rlp::{Decodable, Encodable}; @@ -15,8 +15,8 @@ use reth_primitives_traits::Account; use reth_trie_common::{ proof::ProofNodes, updates::{StorageTrieUpdates, TrieUpdates}, - DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles, ProofTrieNode, RlpNode, - StorageMultiProof, TrieAccount, TrieMask, TrieMasks, TrieNode, EMPTY_ROOT_HASH, + BranchNodeMasksMap, DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles, + ProofTrieNode, RlpNode, StorageMultiProof, TrieAccount, TrieMasks, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; use tracing::{instrument, trace}; @@ -247,19 +247,10 @@ where &mut self, multiproof: DecodedMultiProof, ) -> SparseStateTrieResult<()> { - let DecodedMultiProof { - account_subtree, - storages, - branch_node_hash_masks, - branch_node_tree_masks, - } = multiproof; + let DecodedMultiProof { account_subtree, storages, branch_node_masks } = multiproof; // first reveal the account proof nodes - self.reveal_decoded_account_multiproof( - account_subtree, - branch_node_hash_masks, - branch_node_tree_masks, - )?; + self.reveal_decoded_account_multiproof(account_subtree, branch_node_masks)?; #[cfg(not(feature = "std"))] // If nostd then serially reveal storage proof nodes for each storage trie @@ -329,31 +320,24 @@ where pub fn reveal_account_multiproof( &mut self, account_subtree: ProofNodes, - branch_node_hash_masks: HashMap, - branch_node_tree_masks: HashMap, + branch_node_masks: BranchNodeMasksMap, ) -> SparseStateTrieResult<()> { // decode the multiproof first let decoded_multiproof = account_subtree.try_into()?; - self.reveal_decoded_account_multiproof( - decoded_multiproof, - branch_node_hash_masks, - branch_node_tree_masks, - ) + self.reveal_decoded_account_multiproof(decoded_multiproof, branch_node_masks) } /// Reveals a decoded account multiproof. pub fn reveal_decoded_account_multiproof( &mut self, account_subtree: DecodedProofNodes, - branch_node_hash_masks: HashMap, - branch_node_tree_masks: HashMap, + branch_node_masks: BranchNodeMasksMap, ) -> SparseStateTrieResult<()> { let FilterMappedProofNodes { root_node, nodes, new_nodes, metric_values: _metric_values } = filter_map_revealed_nodes( account_subtree, &mut self.revealed_account_paths, - &branch_node_hash_masks, - &branch_node_tree_masks, + &branch_node_masks, )?; #[cfg(feature = "metrics")] { @@ -426,8 +410,7 @@ where filter_map_revealed_nodes( storage_subtree.subtree, revealed_nodes, - &storage_subtree.branch_node_hash_masks, - &storage_subtree.branch_node_tree_masks, + &storage_subtree.branch_node_masks, )?; if let Some(root_node) = root_node { @@ -959,8 +942,7 @@ struct FilterMappedProofNodes { fn filter_map_revealed_nodes( proof_nodes: DecodedProofNodes, revealed_nodes: &mut HashSet, - branch_node_hash_masks: &HashMap, - branch_node_tree_masks: &HashMap, + branch_node_masks: &BranchNodeMasksMap, ) -> SparseStateTrieResult { let mut result = FilterMappedProofNodes { 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 // according to the state mask. result.new_nodes += branch.state_mask.count_ones() as usize; - TrieMasks { - hash_mask: branch_node_hash_masks.get(&path).copied(), - tree_mask: branch_node_tree_masks.get(&path).copied(), + if let Some(branch_masks) = branch_node_masks.get(&path) { + TrieMasks { + hash_mask: Some(branch_masks.hash_mask), + tree_mask: Some(branch_masks.tree_mask), + } + } else { + TrieMasks::none() } } TrieNode::Extension(_) => { @@ -1042,7 +1028,7 @@ mod tests { use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH}; use reth_trie_common::{ proof::{ProofNodes, ProofRetainer}, - BranchNode, LeafNode, StorageMultiProof, TrieMask, + BranchNode, BranchNodeMasks, LeafNode, StorageMultiProof, TrieMask, }; #[test] @@ -1149,8 +1135,7 @@ mod tests { (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()), (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()), ]), - branch_node_hash_masks: Default::default(), - branch_node_tree_masks: Default::default(), + branch_node_masks: 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_2, &alloy_rlp::encode_fixed_size(&value_2)); + use alloy_trie::TrieMask; + let storage_root = storage_hash_builder.root(); let storage_proof_nodes = storage_hash_builder.take_proof_nodes(); - let storage_branch_node_hash_masks = HashMap::from_iter([ - (Nibbles::default(), TrieMask::new(0b010)), - (Nibbles::from_nibbles([0x1]), TrieMask::new(0b11)), + let storage_branch_node_masks = BranchNodeMasksMap::from_iter([ + ( + 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"); @@ -1257,19 +1250,20 @@ mod tests { .reveal_decoded_multiproof( MultiProof { account_subtree: proof_nodes, - branch_node_hash_masks: HashMap::from_iter([( + branch_node_masks: BranchNodeMasksMap::from_iter([( 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([ ( address_1, StorageMultiProof { root, subtree: storage_proof_nodes.clone(), - branch_node_hash_masks: storage_branch_node_hash_masks.clone(), - branch_node_tree_masks: HashMap::default(), + branch_node_masks: storage_branch_node_masks.clone(), }, ), ( @@ -1277,8 +1271,7 @@ mod tests { StorageMultiProof { root, subtree: storage_proof_nodes, - branch_node_hash_masks: storage_branch_node_hash_masks, - branch_node_tree_masks: HashMap::default(), + branch_node_masks: storage_branch_node_masks, }, ), ]), @@ -1366,16 +1359,11 @@ mod tests { (Nibbles::from_nibbles([0x1]), leaf.clone()), ]); - let branch_node_hash_masks = HashMap::default(); - let branch_node_tree_masks = HashMap::default(); + let branch_node_masks = BranchNodeMasksMap::default(); - let decoded = filter_map_revealed_nodes( - proof_nodes, - &mut revealed_nodes, - &branch_node_hash_masks, - &branch_node_tree_masks, - ) - .unwrap(); + let decoded = + filter_map_revealed_nodes(proof_nodes, &mut revealed_nodes, &branch_node_masks) + .unwrap(); assert_eq!( decoded, diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs index c0b9012355..cba010f2e1 100644 --- a/crates/trie/trie/src/proof/mod.rs +++ b/crates/trie/trie/src/proof/mod.rs @@ -11,14 +11,15 @@ use crate::{ }; use alloy_primitives::{ keccak256, - map::{B256Map, B256Set, HashMap, HashSet}, + map::{B256Map, B256Set, HashSet}, Address, B256, }; use alloy_rlp::{BufMut, Encodable}; use alloy_trie::proof::AddedRemovedKeys; use reth_execution_errors::trie::StateProofError; use reth_trie_common::{ - proof::ProofRetainer, AccountProof, MultiProof, MultiProofTargets, StorageMultiProof, + proof::ProofRetainer, AccountProof, BranchNodeMasks, BranchNodeMasksMap, MultiProof, + MultiProofTargets, StorageMultiProof, }; mod trie_node; @@ -180,20 +181,19 @@ where } let _ = hash_builder.root(); 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(); - ( - updated_branch_nodes.iter().map(|(path, node)| (*path, node.hash_mask)).collect(), - updated_branch_nodes - .into_iter() - .map(|(path, node)| (path, node.tree_mask)) - .collect(), - ) + updated_branch_nodes + .into_iter() + .map(|(path, node)| { + (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask }) + }) + .collect() } 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 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(); - ( - updated_branch_nodes.iter().map(|(path, node)| (*path, node.hash_mask)).collect(), - updated_branch_nodes - .into_iter() - .map(|(path, node)| (path, node.tree_mask)) - .collect(), - ) + updated_branch_nodes + .into_iter() + .map(|(path, node)| { + (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask }) + }) + .collect() } 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 }) } } diff --git a/crates/trie/trie/src/proof/trie_node.rs b/crates/trie/trie/src/proof/trie_node.rs index f6c62b866e..2c7dad6fd2 100644 --- a/crates/trie/trie/src/proof/trie_node.rs +++ b/crates/trie/trie/src/proof/trie_node.rs @@ -79,8 +79,9 @@ where .multiproof(targets) .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; let node = proof.account_subtree.into_inner().remove(path); - let tree_mask = proof.branch_node_tree_masks.remove(path); - let hash_mask = proof.branch_node_hash_masks.remove(path); + let masks = proof.branch_node_masks.remove(path); + let hash_mask = masks.map(|m| m.hash_mask); + let tree_mask = masks.map(|m| m.tree_mask); trace!( target: "trie::proof::blinded", @@ -131,8 +132,9 @@ where .storage_multiproof(targets) .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?; let node = proof.subtree.into_inner().remove(path); - let tree_mask = proof.branch_node_tree_masks.remove(path); - let hash_mask = proof.branch_node_hash_masks.remove(path); + let masks = proof.branch_node_masks.remove(path); + let hash_mask = masks.map(|m| m.hash_mask); + let tree_mask = masks.map(|m| m.tree_mask); trace!( target: "trie::proof::blinded", diff --git a/crates/trie/trie/src/proof_v2/mod.rs b/crates/trie/trie/src/proof_v2/mod.rs index a9cf0b2d51..eb4ee5db7e 100644 --- a/crates/trie/trie/src/proof_v2/mod.rs +++ b/crates/trie/trie/src/proof_v2/mod.rs @@ -1648,17 +1648,15 @@ mod tests { // though we never store the root node so the masks for it aren't really valid. let masks = if path.is_empty() { TrieMasks::none() - } else { + } else if let Some(branch_masks) = + proof_legacy_result.branch_node_masks.get(path) + { TrieMasks { - hash_mask: proof_legacy_result - .branch_node_hash_masks - .get(path) - .copied(), - tree_mask: proof_legacy_result - .branch_node_tree_masks - .get(path) - .copied(), + hash_mask: Some(branch_masks.hash_mask), + tree_mask: Some(branch_masks.tree_mask), } + } else { + TrieMasks::none() }; ProofTrieNode { path: *path, node, masks }