From b1424b13fe94a2fbcd3beb16e4517cebde9617f1 Mon Sep 17 00:00:00 2001 From: Alexey Shekhirin Date: Wed, 4 Dec 2024 14:57:04 +0000 Subject: [PATCH] use hash builder updates --- crates/trie/parallel/src/proof.rs | 16 +- crates/trie/sparse/src/state.rs | 221 +++++++++----------- crates/trie/sparse/src/trie.rs | 55 ++--- crates/trie/trie/src/node_iter.rs | 13 +- crates/trie/trie/src/proof/mod.rs | 33 ++- crates/trie/trie/src/trie_cursor/subnode.rs | 7 - crates/trie/trie/src/walker.rs | 6 - 7 files changed, 138 insertions(+), 213 deletions(-) diff --git a/crates/trie/parallel/src/proof.rs b/crates/trie/parallel/src/proof.rs index a8125528b706..a519fe2947b3 100644 --- a/crates/trie/parallel/src/proof.rs +++ b/crates/trie/parallel/src/proof.rs @@ -158,11 +158,10 @@ where // Create a hash builder to rebuild the root node since it is not available in the database. let retainer: ProofRetainer = targets.keys().map(Nibbles::unpack).collect(); - let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); + let mut hash_builder = + HashBuilder::default().with_proof_retainer(retainer).with_updates(true); let mut storages = HashMap::default(); - let mut branch_node_hash_masks = HashMap::default(); - let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE); let mut account_node_iter = TrieNodeIter::new( walker, @@ -173,7 +172,6 @@ where { match account_node { TrieElement::Branch(node) => { - branch_node_hash_masks.insert(node.key.clone(), node.hash_mask); hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); } TrieElement::Leaf(hashed_address, account) => { @@ -225,11 +223,11 @@ where #[cfg(feature = "metrics")] self.metrics.record_state_trie(tracker.finish()); - Ok(MultiProof { - account_subtree: hash_builder.take_proof_nodes(), - storages, - branch_node_hash_masks, - }) + let account_subtree = hash_builder.take_proof_nodes(); + let branch_node_hash_masks = + hash_builder.split().1.into_iter().map(|(path, node)| (path, node.hash_mask)).collect(); + + Ok(MultiProof { account_subtree, branch_node_hash_masks, storages }) } } diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 4b0b713ce2a3..9b4b38002511 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -13,8 +13,7 @@ use reth_primitives_traits::Account; use reth_tracing::tracing::trace; use reth_trie_common::{ updates::{StorageTrieUpdates, TrieUpdates}, - MultiProof, Nibbles, TrieAccount, TrieMask, TrieNode, EMPTY_ROOT_HASH, - TRIE_ACCOUNT_RLP_MAX_SIZE, + MultiProof, Nibbles, TrieAccount, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; use std::{fmt, iter::Peekable}; @@ -103,6 +102,75 @@ impl SparseStateTrie { self.storages.get_mut(account).and_then(|e| e.as_revealed_mut()) } + /// Reveal unknown trie paths from provided leaf path and its proof for the account. + /// NOTE: This method does not extensively validate the proof. + pub fn reveal_account( + &mut self, + account: B256, + proof: impl IntoIterator, + ) -> SparseStateTrieResult<()> { + if self.is_account_revealed(&account) { + return Ok(()); + } + + let mut proof = proof.into_iter().peekable(); + + let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) }; + + // Reveal root node if it wasn't already. + let trie = self.state.reveal_root_with_provider( + self.provider_factory.account_node_provider(), + root_node, + self.retain_updates, + )?; + + // Reveal the remaining proof nodes. + for (path, bytes) in proof { + let node = TrieNode::decode(&mut &bytes[..])?; + trie.reveal_node(path, node)?; + } + + // Mark leaf path as revealed. + self.revealed.entry(account).or_default(); + + Ok(()) + } + + /// Reveal unknown trie paths from provided leaf path and its proof for the storage slot. + /// NOTE: This method does not extensively validate the proof. + pub fn reveal_storage_slot( + &mut self, + account: B256, + slot: B256, + proof: impl IntoIterator, + ) -> SparseStateTrieResult<()> { + if self.is_storage_slot_revealed(&account, &slot) { + return Ok(()); + } + + let mut proof = proof.into_iter().peekable(); + + let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) }; + + // Reveal root node if it wasn't already. + let trie = self.storages.entry(account).or_default().reveal_root_with_provider( + self.provider_factory.storage_node_provider(account), + root_node, + self.retain_updates, + )?; + + // Reveal the remaining proof nodes. + for (path, bytes) in proof { + let node = TrieNode::decode(&mut &bytes[..])?; + trie.reveal_node(path, node)?; + } + + // Mark leaf path as revealed. + self.revealed.entry(account).or_default().insert(slot); + + Ok(()) + } + /// Reveal unknown trie paths from multiproof and the list of included accounts and slots. /// NOTE: This method does not extensively validate the proof. pub fn reveal_multiproof( @@ -113,14 +181,11 @@ impl SparseStateTrie { let account_subtree = multiproof.account_subtree.into_nodes_sorted(); let mut account_nodes = account_subtree.into_iter().peekable(); - if let Some((root_node, hash_mask)) = - self.validate_root_node(&mut account_nodes, &multiproof.branch_node_hash_masks)? - { + if let Some(root_node) = self.validate_root_node(&mut account_nodes)? { // Reveal root node if it wasn't already. let trie = self.state.reveal_root_with_provider( self.provider_factory.account_node_provider(), root_node, - hash_mask, self.retain_updates, )?; @@ -128,48 +193,27 @@ impl SparseStateTrie { for (path, bytes) in account_nodes { let node = TrieNode::decode(&mut &bytes[..])?; trace!(target: "trie::sparse", ?path, ?node, "Revealing account node"); - - let hash_mask = if let TrieNode::Branch(_) = node { - let Some(hash_mask) = multiproof.branch_node_hash_masks.get(&path) else { - return Err(SparseStateTrieError::MissingBranchNodeHashMask { path }); - }; - Some(*hash_mask) - } else { - None - }; - trie.reveal_node(path, node, hash_mask)?; + trie.reveal_node(path, node)?; } } for (account, storage_subtree) in multiproof.storages { - let subtree = storage_subtree.subtree.into_nodes_sorted(); - let mut nodes = subtree.into_iter().peekable(); + let storage_subtree = storage_subtree.subtree.into_nodes_sorted(); + let mut storage_nodes = storage_subtree.into_iter().peekable(); - if let Some((root_node, hash_mask)) = - self.validate_root_node(&mut nodes, &storage_subtree.branch_node_hash_masks)? - { + if let Some(root_node) = self.validate_root_node(&mut storage_nodes)? { // Reveal root node if it wasn't already. let trie = self.storages.entry(account).or_default().reveal_root_with_provider( self.provider_factory.storage_node_provider(account), root_node, - hash_mask, self.retain_updates, )?; // Reveal the remaining proof nodes. - for (path, bytes) in nodes { + for (path, bytes) in storage_nodes { let node = TrieNode::decode(&mut &bytes[..])?; trace!(target: "trie::sparse", ?account, ?path, ?node, "Revealing storage node"); - let hash_mask = if let TrieNode::Branch(_) = node { - let Some(hash_mask) = storage_subtree.branch_node_hash_masks.get(&path) - else { - return Err(SparseStateTrieError::MissingBranchNodeHashMask { path }); - }; - Some(*hash_mask) - } else { - None - }; - trie.reveal_node(path, node, hash_mask)?; + trie.reveal_node(path, node)?; } } } @@ -185,8 +229,7 @@ impl SparseStateTrie { fn validate_root_node>( &self, proof: &mut Peekable, - branch_node_hash_masks: &HashMap, - ) -> SparseStateTrieResult)>> { + ) -> SparseStateTrieResult> { let mut proof = proof.into_iter().peekable(); // Validate root node. @@ -201,16 +244,7 @@ impl SparseStateTrie { return Err(SparseStateTrieError::InvalidRootNode { path, node }) } - let hash_mask = if matches!(root_node, TrieNode::Branch(_)) { - let Some(hash_mask) = branch_node_hash_masks.get(&path) else { - return Err(SparseStateTrieError::MissingBranchNodeHashMask { path }) - }; - Some(*hash_mask) - } else { - None - }; - - Ok(Some((root_node, hash_mask))) + Ok(Some(root_node)) } /// Update the account leaf node. @@ -360,40 +394,27 @@ mod tests { use rand::{rngs::StdRng, Rng, SeedableRng}; use reth_primitives_traits::Account; use reth_trie::{updates::StorageTrieUpdates, HashBuilder, TrieAccount, EMPTY_ROOT_HASH}; - use reth_trie_common::{proof::ProofRetainer, BranchNode, StorageMultiProof}; + use reth_trie_common::proof::ProofRetainer; #[test] - fn validate_root_node_first_node_not_root() { + fn validate_proof_first_node_not_root() { let sparse = SparseStateTrie::default(); let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))]; assert_matches!( - sparse.validate_root_node(&mut proof.into_iter().peekable(), &HashMap::default()), + sparse.validate_root_node(&mut proof.into_iter().peekable()), Err(SparseStateTrieError::InvalidRootNode { .. }) ); } #[test] - fn validate_root_node_invalid_proof_with_empty_root() { + fn validate_proof_invalid_proof_with_empty_root() { let sparse = SparseStateTrie::default(); let proof = [ (Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])), (Nibbles::from_nibbles([0x1]), Bytes::new()), ]; assert_matches!( - sparse.validate_root_node(&mut proof.into_iter().peekable(), &HashMap::default()), - Err(SparseStateTrieError::InvalidRootNode { .. }) - ); - } - - #[test] - fn validate_root_node_no_hash_mask() { - let sparse = SparseStateTrie::default(); - - let mut rlp = Vec::new(); - BranchNode::default().encode(&mut rlp); - let proof = [(Nibbles::default(), Bytes::from(rlp))]; - assert_matches!( - sparse.validate_root_node(&mut proof.into_iter().peekable(), &HashMap::default()), + sparse.validate_root_node(&mut proof.into_iter().peekable()), Err(SparseStateTrieError::InvalidRootNode { .. }) ); } @@ -402,16 +423,13 @@ mod tests { fn reveal_account_empty() { let retainer = ProofRetainer::from_iter([Nibbles::default()]); let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); - let root = hash_builder.root(); + hash_builder.root(); let proofs = hash_builder.take_proof_nodes(); assert_eq!(proofs.len(), 1); let mut sparse = SparseStateTrie::default(); assert_eq!(sparse.state, SparseTrie::Blind); - - let targets = HashMap::from_iter([(root, HashSet::default())]); - let multiproof = MultiProof { account_subtree: proofs, ..Default::default() }; - sparse.reveal_multiproof(targets, multiproof).unwrap(); + sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap(); assert_eq!(sparse.state, SparseTrie::revealed_empty()); } @@ -419,27 +437,15 @@ mod tests { fn reveal_storage_slot_empty() { let retainer = ProofRetainer::from_iter([Nibbles::default()]); let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); - let root = hash_builder.root(); + hash_builder.root(); let proofs = hash_builder.take_proof_nodes(); assert_eq!(proofs.len(), 1); let mut sparse = SparseStateTrie::default(); assert!(sparse.storages.is_empty()); - - let targets = - HashMap::from_iter([(B256::default(), HashSet::from_iter([B256::default()]))]); - let multiproof = MultiProof { - storages: HashMap::from_iter([( - B256::default(), - StorageMultiProof { - root, - subtree: proofs, - branch_node_hash_masks: HashMap::default(), - }, - )]), - ..Default::default() - }; - sparse.reveal_multiproof(targets, multiproof).unwrap(); + sparse + .reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner()) + .unwrap(); assert_eq!( sparse.storages, HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())]) @@ -475,11 +481,9 @@ mod tests { storage_hash_builder.add_leaf(slot_path_2.clone(), &alloy_rlp::encode_fixed_size(&value_2)); 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 proof_nodes = storage_hash_builder.take_proof_nodes(); + let storage_proof_1 = proof_nodes.matching_nodes_sorted(&slot_path_1); + let storage_proof_2 = proof_nodes.matching_nodes_sorted(&slot_path_2); let address_1 = b256!("1000000000000000000000000000000000000000000000000000000000000000"); let address_path_1 = Nibbles::unpack(address_1); @@ -500,39 +504,16 @@ mod tests { let root = hash_builder.root(); let proof_nodes = hash_builder.take_proof_nodes(); + let proof_1 = proof_nodes.matching_nodes_sorted(&address_path_1); + let proof_2 = proof_nodes.matching_nodes_sorted(&address_path_2); let mut sparse = SparseStateTrie::default().with_updates(true); - sparse.reveal_multiproof( - HashMap::from_iter([ - (address_1, HashSet::from_iter([slot_1, slot_2])), - (address_2, HashSet::from_iter([slot_1, slot_2])), - ]), - MultiProof { - account_subtree: proof_nodes, - branch_node_hash_masks: HashMap::from_iter([( - Nibbles::from_nibbles([0x1]), - TrieMask::new(0b00), - )]), - storages: HashMap::from_iter([ - ( - address_1, - StorageMultiProof { - root, - subtree: storage_proof_nodes.clone(), - branch_node_hash_masks: storage_branch_node_hash_masks.clone(), - }, - ), - ( - address_2, - StorageMultiProof { - root, - subtree: storage_proof_nodes, - branch_node_hash_masks: storage_branch_node_hash_masks, - }, - ), - ]), - }, - ); + sparse.reveal_account(address_1, proof_1).unwrap(); + sparse.reveal_account(address_2, proof_2).unwrap(); + sparse.reveal_storage_slot(address_1, slot_1, storage_proof_1.clone()).unwrap(); + sparse.reveal_storage_slot(address_1, slot_2, storage_proof_2.clone()).unwrap(); + sparse.reveal_storage_slot(address_2, slot_1, storage_proof_1).unwrap(); + sparse.reveal_storage_slot(address_2, slot_2, storage_proof_2).unwrap(); assert_eq!(sparse.root(), Some(root)); diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 297d12e790b4..df5dd25486c1 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -59,10 +59,9 @@ impl SparseTrie { pub fn reveal_root( &mut self, root: TrieNode, - hash_mask: Option, retain_updates: bool, ) -> SparseTrieResult<&mut RevealedSparseTrie> { - self.reveal_root_with_provider(Default::default(), root, hash_mask, retain_updates) + self.reveal_root_with_provider(Default::default(), root, retain_updates) } } @@ -90,14 +89,12 @@ impl

SparseTrie

{ &mut self, provider: P, root: TrieNode, - hash_mask: Option, retain_updates: bool, ) -> SparseTrieResult<&mut RevealedSparseTrie

> { if self.is_blind() { *self = Self::Revealed(Box::new(RevealedSparseTrie::from_provider_and_root( provider, root, - hash_mask, retain_updates, )?)) } @@ -156,8 +153,6 @@ pub struct RevealedSparseTrie

{ provider: P, /// All trie nodes. nodes: HashMap, - /// All branch node hash masks. - branch_node_hash_masks: HashMap, /// All leaf values. values: HashMap>, /// Prefix set. @@ -172,7 +167,6 @@ impl

fmt::Debug for RevealedSparseTrie

{ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("RevealedSparseTrie") .field("nodes", &self.nodes) - .field("branch_hash_masks", &self.branch_node_hash_masks) .field("values", &self.values) .field("prefix_set", &self.prefix_set) .field("updates", &self.updates) @@ -186,7 +180,6 @@ impl Default for RevealedSparseTrie { Self { provider: Default::default(), nodes: HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]), - branch_node_hash_masks: HashMap::default(), values: HashMap::default(), prefix_set: PrefixSetMut::default(), updates: None, @@ -197,22 +190,17 @@ impl Default for RevealedSparseTrie { impl RevealedSparseTrie { /// Create new revealed sparse trie from the given root node. - pub fn from_root( - node: TrieNode, - hash_mask: Option, - retain_updates: bool, - ) -> SparseTrieResult { + pub fn from_root(node: TrieNode, retain_updates: bool) -> SparseTrieResult { let mut this = Self { provider: Default::default(), nodes: HashMap::default(), - branch_node_hash_masks: HashMap::default(), values: HashMap::default(), prefix_set: PrefixSetMut::default(), rlp_buf: Vec::new(), updates: None, } .with_updates(retain_updates); - this.reveal_node(Nibbles::default(), node, hash_mask)?; + this.reveal_node(Nibbles::default(), node)?; Ok(this) } } @@ -222,20 +210,18 @@ impl

RevealedSparseTrie

{ pub fn from_provider_and_root( provider: P, node: TrieNode, - hash_mask: Option, retain_updates: bool, ) -> SparseTrieResult { let mut this = Self { provider, nodes: HashMap::default(), - branch_node_hash_masks: HashMap::default(), values: HashMap::default(), prefix_set: PrefixSetMut::default(), rlp_buf: Vec::new(), updates: None, } .with_updates(retain_updates); - this.reveal_node(Nibbles::default(), node, hash_mask)?; + this.reveal_node(Nibbles::default(), node)?; Ok(this) } @@ -244,7 +230,6 @@ impl

RevealedSparseTrie

{ RevealedSparseTrie { provider, nodes: self.nodes, - branch_node_hash_masks: self.branch_node_hash_masks, values: self.values, prefix_set: self.prefix_set, updates: self.updates, @@ -276,16 +261,7 @@ impl

RevealedSparseTrie

{ } /// Reveal the trie node only if it was not known already. - pub fn reveal_node( - &mut self, - path: Nibbles, - node: TrieNode, - hash_mask: Option, - ) -> SparseTrieResult<()> { - if let Some(hash_mask) = hash_mask { - self.branch_node_hash_masks.insert(path.clone(), hash_mask); - } - + pub fn reveal_node(&mut self, path: Nibbles, node: TrieNode) -> SparseTrieResult<()> { match node { TrieNode::EmptyRoot => { debug_assert!(path.is_empty()); @@ -369,7 +345,7 @@ impl

RevealedSparseTrie

{ return Ok(()) } - self.reveal_node(path, TrieNode::decode(&mut &child[..])?, None) + self.reveal_node(path, TrieNode::decode(&mut &child[..])?) } /// Update the leaf node with provided value. @@ -957,8 +933,7 @@ where if let Some(node) = self.provider.blinded_node(child_path.clone())? { let decoded = TrieNode::decode(&mut &node[..])?; trace!(target: "trie::sparse", ?child_path, ?decoded, "Revealing remaining blinded branch child"); - // TODO: fetch hash mask - self.reveal_node(child_path.clone(), decoded, None)?; + self.reveal_node(child_path.clone(), decoded)?; } } @@ -1240,7 +1215,7 @@ mod tests { state: impl IntoIterator + Clone, destroyed_accounts: HashSet, proof_targets: impl IntoIterator, - ) -> (B256, TrieUpdates, ProofNodes, HashMap) { + ) -> (B256, TrieUpdates, ProofNodes) { let mut account_rlp = Vec::new(); let mut hash_builder = HashBuilder::default() @@ -1264,11 +1239,9 @@ mod tests { ), ); - let mut branch_node_hash_masks = HashMap::default(); while let Some(node) = node_iter.try_next().unwrap() { match node { TrieElement::Branch(branch) => { - branch_node_hash_masks.insert(branch.key, branch.hash_mask); hash_builder.add_branch(branch.key, branch.value, branch.children_are_in_trie); } TrieElement::Leaf(key, account) => { @@ -1287,7 +1260,7 @@ mod tests { let removed_keys = node_iter.walker.take_removed_keys(); trie_updates.finalize(hash_builder, removed_keys, destroyed_accounts); - (root, trie_updates, proof_nodes, branch_node_hash_masks) + (root, trie_updates, proof_nodes) } /// Assert that the sparse trie nodes and the proof nodes from the hash builder are equal. @@ -1771,8 +1744,8 @@ mod tests { // Branch (Mask = 11) // ├── 0 -> Hash (Path = 0) // └── 1 -> Leaf (Path = 1) - sparse.reveal_node(Nibbles::default(), branch, Some(TrieMask::new(0b01))).unwrap(); - sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + sparse.reveal_node(Nibbles::default(), branch).unwrap(); + sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf)).unwrap(); // Removing a blinded leaf should result in an error assert_matches!( @@ -1802,8 +1775,8 @@ mod tests { // Branch (Mask = 11) // ├── 0 -> Hash (Path = 0) // └── 1 -> Leaf (Path = 1) - sparse.reveal_node(Nibbles::default(), branch, Some(TrieMask::new(0b01))).unwrap(); - sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); + sparse.reveal_node(Nibbles::default(), branch).unwrap(); + sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf)).unwrap(); // Removing a non-existent leaf should be a noop let sparse_old = sparse.clone(); @@ -2227,7 +2200,7 @@ mod tests { account_rlp }; - let (hash_builder_root, hash_builder_updates, _, _) = run_hash_builder( + let (hash_builder_root, hash_builder_updates, _) = run_hash_builder( [(key1(), value()), (key2(), value())], Default::default(), [Nibbles::default()], diff --git a/crates/trie/trie/src/node_iter.rs b/crates/trie/trie/src/node_iter.rs index 16cecf23d6e8..60219eedd7cb 100644 --- a/crates/trie/trie/src/node_iter.rs +++ b/crates/trie/trie/src/node_iter.rs @@ -1,6 +1,5 @@ use crate::{hashed_cursor::HashedCursor, trie_cursor::TrieCursor, walker::TrieWalker, Nibbles}; use alloy_primitives::B256; -use alloy_trie::TrieMask; use reth_storage_errors::db::DatabaseError; /// Represents a branch node in the trie. @@ -12,19 +11,12 @@ pub struct TrieBranchNode { pub value: B256, /// Indicates whether children are in the trie. pub children_are_in_trie: bool, - /// The hash mask of the branch node. - pub hash_mask: TrieMask, } impl TrieBranchNode { /// Creates a new `TrieBranchNode`. - pub const fn new( - key: Nibbles, - value: B256, - children_are_in_trie: bool, - hash_mask: TrieMask, - ) -> Self { - Self { key, value, children_are_in_trie, hash_mask } + pub const fn new(key: Nibbles, value: B256, children_are_in_trie: bool) -> Self { + Self { key, value, children_are_in_trie } } } @@ -105,7 +97,6 @@ where key.clone(), self.walker.hash().unwrap(), self.walker.children_are_in_trie(), - self.walker.hash_mask(), )))) } } diff --git a/crates/trie/trie/src/proof/mod.rs b/crates/trie/trie/src/proof/mod.rs index 21026034bc13..71ca994f6912 100644 --- a/crates/trie/trie/src/proof/mod.rs +++ b/crates/trie/trie/src/proof/mod.rs @@ -104,21 +104,19 @@ where // Create a hash builder to rebuild the root node since it is not available in the database. let retainer = targets.keys().map(Nibbles::unpack).collect(); - let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); + let mut hash_builder = + HashBuilder::default().with_proof_retainer(retainer).with_updates(true); // Initialize all storage multiproofs as empty. // Storage multiproofs for non empty tries will be overwritten if necessary. let mut storages: HashMap<_, _> = targets.keys().map(|key| (*key, StorageMultiProof::empty())).collect(); - let mut branch_node_hash_masks = HashMap::default(); let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE); let mut account_node_iter = TrieNodeIter::new(walker, hashed_account_cursor); - let mut branch_node_hash_masks = HashMap::default(); while let Some(account_node) = account_node_iter.try_next()? { match account_node { TrieElement::Branch(node) => { - branch_node_hash_masks.insert(node.key.clone(), node.hash_mask); hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); } TrieElement::Leaf(hashed_address, account) => { @@ -153,11 +151,11 @@ where } } let _ = hash_builder.root(); - Ok(MultiProof { - account_subtree: hash_builder.take_proof_nodes(), - branch_node_hash_masks, - storages, - }) + let account_subtree = hash_builder.take_proof_nodes(); + let branch_node_hash_masks = + hash_builder.split().1.into_iter().map(|(path, node)| (path, node.hash_mask)).collect(); + + Ok(MultiProof { account_subtree, branch_node_hash_masks, storages }) } } @@ -250,16 +248,13 @@ where let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?; let walker = TrieWalker::new(trie_cursor, self.prefix_set.freeze()); - let mut branch_node_hash_masks = HashMap::default(); - let retainer = ProofRetainer::from_iter(target_nibbles); - let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer); + let mut hash_builder = + HashBuilder::default().with_proof_retainer(retainer).with_updates(true); let mut storage_node_iter = TrieNodeIter::new(walker, hashed_storage_cursor); - let mut branch_node_hash_masks = HashMap::default(); while let Some(node) = storage_node_iter.try_next()? { match node { TrieElement::Branch(node) => { - branch_node_hash_masks.insert(node.key.clone(), node.hash_mask); hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); } TrieElement::Leaf(hashed_slot, value) => { @@ -272,10 +267,10 @@ where } let root = hash_builder.root(); - Ok(StorageMultiProof { - root, - subtree: hash_builder.take_proof_nodes(), - branch_node_hash_masks, - }) + let subtree = hash_builder.take_proof_nodes(); + let branch_node_hash_masks = + hash_builder.split().1.into_iter().map(|(path, node)| (path, node.hash_mask)).collect(); + + Ok(StorageMultiProof { root, subtree, branch_node_hash_masks }) } } diff --git a/crates/trie/trie/src/trie_cursor/subnode.rs b/crates/trie/trie/src/trie_cursor/subnode.rs index 0843472e0999..457c1ba4685b 100644 --- a/crates/trie/trie/src/trie_cursor/subnode.rs +++ b/crates/trie/trie/src/trie_cursor/subnode.rs @@ -1,6 +1,5 @@ use crate::{BranchNodeCompact, Nibbles, StoredSubNode, CHILD_INDEX_RANGE}; use alloy_primitives::B256; -use alloy_trie::TrieMask; /// Cursor for iterating over a subtrie. #[derive(Clone)] @@ -88,12 +87,6 @@ impl CursorSubNode { .is_none_or(|node| self.nibble < 0 || node.tree_mask.is_bit_set(self.nibble as u8)) } - /// Returns the hash mask of the current node. - #[inline] - pub fn hash_mask(&self) -> TrieMask { - self.node.as_ref().map(|node| node.hash_mask).unwrap_or_default() - } - /// Returns `true` if the current nibble has a root hash. pub fn hash_flag(&self) -> bool { self.node.as_ref().is_some_and(|node| match self.nibble { diff --git a/crates/trie/trie/src/walker.rs b/crates/trie/trie/src/walker.rs index c3fd7d51197a..647c1486ef09 100644 --- a/crates/trie/trie/src/walker.rs +++ b/crates/trie/trie/src/walker.rs @@ -4,7 +4,6 @@ use crate::{ BranchNodeCompact, Nibbles, }; use alloy_primitives::{map::HashSet, B256}; -use alloy_trie::TrieMask; use reth_storage_errors::db::DatabaseError; #[cfg(feature = "metrics")] @@ -96,11 +95,6 @@ impl TrieWalker { self.stack.last().is_some_and(|n| n.tree_flag()) } - /// Returns the hash mask of the current node. - pub fn hash_mask(&self) -> TrieMask { - self.stack.last().map(|n| n.hash_mask()).unwrap_or_default() - } - /// Returns the next unprocessed key in the trie. pub fn next_unprocessed_key(&self) -> Option { self.key()