diff --git a/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs index ced00e9c39..c4f4e0ec36 100644 --- a/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/configured_sparse_trie.rs @@ -1,7 +1,7 @@ //! Configured sparse trie enum for switching between serial and parallel implementations. use alloy_primitives::B256; -use reth_trie::{Nibbles, ProofTrieNode, TrieMasks, TrieNode}; +use reth_trie::{BranchNodeMasks, Nibbles, ProofTrieNode, TrieNode}; use reth_trie_sparse::{ errors::SparseTrieResult, provider::TrieNodeProvider, LeafLookup, LeafLookupError, SerialSparseTrie, SparseTrieInterface, SparseTrieUpdates, @@ -44,7 +44,7 @@ impl SparseTrieInterface for ConfiguredSparseTrie { fn with_root( self, root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult { match self { @@ -75,7 +75,7 @@ impl SparseTrieInterface for ConfiguredSparseTrie { &mut self, path: Nibbles, node: TrieNode, - masks: TrieMasks, + masks: Option, ) -> SparseTrieResult<()> { match self { Self::Serial(trie) => trie.reveal_node(path, node, masks), diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs index 4a80be284c..e3ed7901ef 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::{BranchNodeMasks, BranchNodeMasksMap, ProofTrieNode, TrieMasks}; +pub use trie::{BranchNodeMasks, BranchNodeMasksMap, ProofTrieNode}; /// 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/trie.rs b/crates/trie/common/src/trie.rs index 1d2706c858..52f238e678 100644 --- a/crates/trie/common/src/trie.rs +++ b/crates/trie/common/src/trie.rs @@ -17,39 +17,23 @@ pub struct BranchNodeMasks { 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). -/// The masks represent which children are stored in different ways: -/// - `hash_mask`: Indicates which children are stored as hashes in the database -/// - `tree_mask`: Indicates which children are complete subtrees stored in the database -/// -/// These masks are essential for efficient trie traversal and serialization, as they -/// determine how nodes should be encoded and stored on disk. -#[derive(Debug, PartialEq, Eq, Clone, Copy)] -pub struct TrieMasks { - /// Branch node hash mask, if any. +impl BranchNodeMasks { + /// Creates masks from optional hash and tree masks. /// - /// When a bit is set, the corresponding child node's hash is stored in the trie. - /// - /// This mask enables selective hashing of child nodes. - pub hash_mask: Option, - /// Branch node tree mask, if any. - /// - /// When a bit is set, the corresponding child subtree is stored in the database. - pub tree_mask: Option, -} - -impl TrieMasks { - /// Helper function, returns both fields `hash_mask` and `tree_mask` as [`None`] - pub const fn none() -> Self { - Self { hash_mask: None, tree_mask: None } + /// Returns `None` if both masks are `None`, otherwise defaults missing masks to empty. + pub fn from_optional(hash_mask: Option, tree_mask: Option) -> Option { + match (hash_mask, tree_mask) { + (Some(h), Some(t)) => Some(Self { hash_mask: h, tree_mask: t }), + (Some(h), None) => Some(Self { hash_mask: h, tree_mask: TrieMask::default() }), + (None, Some(t)) => Some(Self { hash_mask: TrieMask::default(), tree_mask: t }), + (None, None) => None, + } } } +/// Map from nibble path to branch node masks. +pub type BranchNodeMasksMap = HashMap; + /// Carries all information needed by a sparse trie to reveal a particular node. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProofTrieNode { @@ -58,5 +42,6 @@ pub struct ProofTrieNode { /// The node itself. pub node: TrieNode, /// Tree and hash masks for the node, if known. - pub masks: TrieMasks, + /// Both masks are always set together (from database branch nodes). + pub masks: Option, } diff --git a/crates/trie/sparse-parallel/src/trie.rs b/crates/trie/sparse-parallel/src/trie.rs index c8e10cf2fe..6a1006fcbf 100644 --- a/crates/trie/sparse-parallel/src/trie.rs +++ b/crates/trie/sparse-parallel/src/trie.rs @@ -10,7 +10,7 @@ use reth_execution_errors::{SparseTrieErrorKind, SparseTrieResult}; use reth_trie_common::{ prefix_set::{PrefixSet, PrefixSetMut}, BranchNodeMasks, BranchNodeMasksMap, BranchNodeRef, ExtensionNodeRef, LeafNodeRef, Nibbles, - ProofTrieNode, RlpNode, TrieMasks, TrieNode, CHILD_INDEX_RANGE, + ProofTrieNode, RlpNode, TrieNode, CHILD_INDEX_RANGE, }; use reth_trie_sparse::{ provider::{RevealedNode, TrieNodeProvider}, @@ -151,7 +151,7 @@ impl SparseTrieInterface for ParallelSparseTrie { fn with_root( mut self, root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult { // A fresh/cleared `ParallelSparseTrie` has a `SparseNode::Empty` at its root in the upper @@ -188,14 +188,8 @@ impl SparseTrieInterface for ParallelSparseTrie { // Update the top-level branch node masks. This is simple and can't be done in parallel. for ProofTrieNode { path, masks, .. } in &nodes { - if masks.tree_mask.is_some() || masks.hash_mask.is_some() { - self.branch_node_masks.insert( - *path, - BranchNodeMasks { - tree_mask: masks.tree_mask.unwrap_or_default(), - hash_mask: masks.hash_mask.unwrap_or_default(), - }, - ); + if let Some(branch_masks) = masks { + self.branch_node_masks.insert(*path, *branch_masks); } } @@ -364,11 +358,8 @@ impl SparseTrieInterface for ParallelSparseTrie { ?hash_mask, "Revealing child (from upper)", ); - subtrie.reveal_node( - reveal_path, - &decoded, - TrieMasks { hash_mask, tree_mask }, - )?; + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + subtrie.reveal_node(reveal_path, &decoded, masks)?; } else { return Err(SparseTrieErrorKind::NodeNotFoundInProvider { path: reveal_path, @@ -937,7 +928,7 @@ impl ParallelSparseTrie { /// Self if successful, or an error if revealing fails. pub fn from_root( root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult { Self::default().with_root(root, masks, retain_updates) @@ -1241,42 +1232,42 @@ impl ParallelSparseTrie { ) -> SparseTrieResult { let remaining_child_subtrie = self.subtrie_for_path_mut(remaining_child_path); - let remaining_child_node = - match remaining_child_subtrie.nodes.get(remaining_child_path).unwrap() { - SparseNode::Hash(_) => { - debug!( + let remaining_child_node = match remaining_child_subtrie + .nodes + .get(remaining_child_path) + .unwrap() + { + SparseNode::Hash(_) => { + debug!( + target: "trie::parallel_sparse", + child_path = ?remaining_child_path, + leaf_full_path = ?full_path, + "Node child not revealed in remove_leaf, falling back to db", + ); + if let Some(RevealedNode { node, tree_mask, hash_mask }) = + provider.trie_node(remaining_child_path)? + { + let decoded = TrieNode::decode(&mut &node[..])?; + trace!( target: "trie::parallel_sparse", - child_path = ?remaining_child_path, - leaf_full_path = ?full_path, - "Node child not revealed in remove_leaf, falling back to db", + ?remaining_child_path, + ?decoded, + ?tree_mask, + ?hash_mask, + "Revealing remaining blinded branch child" ); - if let Some(RevealedNode { node, tree_mask, hash_mask }) = - provider.trie_node(remaining_child_path)? - { - let decoded = TrieNode::decode(&mut &node[..])?; - trace!( - target: "trie::parallel_sparse", - ?remaining_child_path, - ?decoded, - ?tree_mask, - ?hash_mask, - "Revealing remaining blinded branch child" - ); - remaining_child_subtrie.reveal_node( - *remaining_child_path, - &decoded, - TrieMasks { hash_mask, tree_mask }, - )?; - remaining_child_subtrie.nodes.get(remaining_child_path).unwrap().clone() - } else { - return Err(SparseTrieErrorKind::NodeNotFoundInProvider { - path: *remaining_child_path, - } - .into()) + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + remaining_child_subtrie.reveal_node(*remaining_child_path, &decoded, masks)?; + remaining_child_subtrie.nodes.get(remaining_child_path).unwrap().clone() + } else { + return Err(SparseTrieErrorKind::NodeNotFoundInProvider { + path: *remaining_child_path, } + .into()) } - node => node.clone(), - }; + } + node => node.clone(), + }; // If `recurse_into_extension` is true, and the remaining child is an extension node, then // its child will be ensured to be revealed as well. This is required for generation of @@ -1514,7 +1505,7 @@ impl ParallelSparseTrie { /// /// * `path` - The path where the node should be revealed /// * `node` - The trie node to reveal - /// * `masks` - Trie masks for branch nodes + /// * `masks` - Branch node masks if known /// /// # Returns /// @@ -1523,7 +1514,7 @@ impl ParallelSparseTrie { &mut self, path: Nibbles, node: &TrieNode, - masks: TrieMasks, + masks: Option, ) -> SparseTrieResult<()> { // If there is no subtrie for the path it means the path is UPPER_TRIE_MAX_DEPTH or less // nibbles, and so belongs to the upper trie. @@ -1692,11 +1683,8 @@ impl SparseSubtrie { ?hash_mask, "Revealing child (from lower)", ); - self.reveal_node( - reveal_path, - &decoded, - TrieMasks { hash_mask, tree_mask }, - )?; + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + self.reveal_node(reveal_path, &decoded, masks)?; } else { return Err(SparseTrieErrorKind::NodeNotFoundInProvider { path: reveal_path, @@ -1848,7 +1836,7 @@ impl SparseSubtrie { &mut self, path: Nibbles, node: &TrieNode, - masks: TrieMasks, + masks: Option, ) -> SparseTrieResult<()> { debug_assert!(path.starts_with(&self.path)); @@ -1890,10 +1878,9 @@ impl SparseSubtrie { // Memoize the hash of a previously blinded node in a new branch // node. hash: Some(*hash), - store_in_db_trie: Some( - masks.hash_mask.is_some_and(|mask| !mask.is_empty()) || - masks.tree_mask.is_some_and(|mask| !mask.is_empty()), - ), + store_in_db_trie: Some(masks.is_some_and(|m| { + !m.hash_mask.is_empty() || !m.tree_mask.is_empty() + })), }); } // Branch node already exists, or an extension node was placed where a @@ -2030,7 +2017,7 @@ impl SparseSubtrie { return Ok(()) } - self.reveal_node(path, &TrieNode::decode(&mut &child[..])?, TrieMasks::none()) + self.reveal_node(path, &TrieNode::decode(&mut &child[..])?, None) } /// Recalculates and updates the RLP hashes for the changed nodes in this subtrie. @@ -2658,8 +2645,8 @@ mod tests { prefix_set::PrefixSetMut, proof::{ProofNodes, ProofRetainer}, updates::TrieUpdates, - BranchNode, BranchNodeMasksMap, ExtensionNode, HashBuilder, LeafNode, ProofTrieNode, - RlpNode, TrieMask, TrieMasks, TrieNode, EMPTY_ROOT_HASH, + BranchNode, BranchNodeMasks, BranchNodeMasksMap, ExtensionNode, HashBuilder, LeafNode, + ProofTrieNode, RlpNode, TrieMask, TrieNode, EMPTY_ROOT_HASH, }; use reth_trie_db::DatabaseTrieCursorFactory; use reth_trie_sparse::{ @@ -3245,7 +3232,7 @@ mod tests { { let path = Nibbles::from_nibbles([0x1]); let node = create_leaf_node([0x2, 0x3], 42); - let masks = TrieMasks::none(); + let masks = None; trie.reveal_nodes(vec![ProofTrieNode { path, node, masks }]).unwrap(); @@ -3266,7 +3253,7 @@ mod tests { { let path = Nibbles::from_nibbles([0x1, 0x2]); let node = create_leaf_node([0x3, 0x4], 42); - let masks = TrieMasks::none(); + let masks = None; trie.reveal_nodes(vec![ProofTrieNode { path, node, masks }]).unwrap(); @@ -3290,7 +3277,7 @@ mod tests { { let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]); let node = create_leaf_node([0x4, 0x5], 42); - let masks = TrieMasks::none(); + let masks = None; trie.reveal_nodes(vec![ProofTrieNode { path, node, masks }]).unwrap(); @@ -3306,7 +3293,7 @@ mod tests { let path = Nibbles::new(); let child_hash = B256::repeat_byte(0xab); let node = create_extension_node([0x1], child_hash); - let masks = TrieMasks::none(); + let masks = None; let trie = ParallelSparseTrie::from_root(node, masks, true).unwrap(); assert_matches!( @@ -3325,7 +3312,7 @@ mod tests { let path = Nibbles::new(); let child_hash = B256::repeat_byte(0xcd); let node = create_extension_node([0x1, 0x2, 0x3], child_hash); - let masks = TrieMasks::none(); + let masks = None; let trie = ParallelSparseTrie::from_root(node, masks, true).unwrap(); // Extension node should be in upper trie @@ -3351,7 +3338,7 @@ mod tests { let path = Nibbles::from_nibbles([0x1]); let child_hash = B256::repeat_byte(0xcd); let node = create_extension_node([0x2], child_hash); - let masks = TrieMasks::none(); + let masks = None; trie.reveal_nodes(vec![ProofTrieNode { path, node, masks }]).unwrap(); @@ -3380,7 +3367,7 @@ mod tests { RlpNode::word_rlp(&B256::repeat_byte(0x22)), ]; let node = create_branch_node_with_children(&[0x0, 0x5], child_hashes.clone()); - let masks = TrieMasks::none(); + let masks = None; let trie = ParallelSparseTrie::from_root(node, masks, true).unwrap(); // Branch node should be in upper trie @@ -3413,7 +3400,7 @@ mod tests { RlpNode::word_rlp(&B256::repeat_byte(0x55)), ]; let node = create_branch_node_with_children(&[0x0, 0x7, 0xf], child_hashes.clone()); - let masks = TrieMasks::none(); + let masks = None; trie.reveal_nodes(vec![ProofTrieNode { path, node, masks }]).unwrap(); @@ -3472,10 +3459,10 @@ mod tests { // Reveal nodes using reveal_nodes trie.reveal_nodes(vec![ - ProofTrieNode { path: branch_path, node: branch_node, masks: TrieMasks::none() }, - ProofTrieNode { path: leaf_1_path, node: leaf_1, masks: TrieMasks::none() }, - ProofTrieNode { path: leaf_2_path, node: leaf_2, masks: TrieMasks::none() }, - ProofTrieNode { path: leaf_3_path, node: leaf_3, masks: TrieMasks::none() }, + ProofTrieNode { path: branch_path, node: branch_node, masks: None }, + ProofTrieNode { path: leaf_1_path, node: leaf_1, masks: None }, + ProofTrieNode { path: leaf_2_path, node: leaf_2, masks: None }, + ProofTrieNode { path: leaf_3_path, node: leaf_3, masks: None }, ]) .unwrap(); @@ -3568,12 +3555,12 @@ mod tests { ); // Reveal nodes - subtrie.reveal_node(branch_2_path, &branch_2, TrieMasks::none()).unwrap(); - subtrie.reveal_node(leaf_1_path, &leaf_1, TrieMasks::none()).unwrap(); - subtrie.reveal_node(extension_path, &extension, TrieMasks::none()).unwrap(); - subtrie.reveal_node(branch_1_path, &branch_1, TrieMasks::none()).unwrap(); - subtrie.reveal_node(leaf_2_path, &leaf_2, TrieMasks::none()).unwrap(); - subtrie.reveal_node(leaf_3_path, &leaf_3, TrieMasks::none()).unwrap(); + subtrie.reveal_node(branch_2_path, &branch_2, None).unwrap(); + subtrie.reveal_node(leaf_1_path, &leaf_1, None).unwrap(); + subtrie.reveal_node(extension_path, &extension, None).unwrap(); + subtrie.reveal_node(branch_1_path, &branch_1, None).unwrap(); + subtrie.reveal_node(leaf_2_path, &leaf_2, None).unwrap(); + subtrie.reveal_node(leaf_3_path, &leaf_3, None).unwrap(); // Run hash builder for two leaf nodes let (_, _, proof_nodes, _, _) = run_hash_builder( @@ -4233,10 +4220,10 @@ mod tests { TrieMask::new(0b1111111111111111), )) }, - masks: TrieMasks { - hash_mask: Some(TrieMask::new(0b1111111111111111)), - tree_mask: Some(TrieMask::new(0b0011110100100101)), - }, + masks: Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b1111111111111111), + tree_mask: TrieMask::new(0b0011110100100101), + }), }, // Branch at 0x4f88072 ProofTrieNode { @@ -4253,10 +4240,10 @@ mod tests { let branch_node = BranchNode::new(stack, TrieMask::new(0b0001000000000100)); TrieNode::Branch(branch_node) }, - masks: TrieMasks { - hash_mask: Some(TrieMask::new(0b0000000000000000)), - tree_mask: Some(TrieMask::new(0b0000000000000100)), - }, + masks: Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b0000000000000000), + tree_mask: TrieMask::new(0b0000000000000100), + }), }, // Extension at 0x4f880722 ProofTrieNode { @@ -4270,7 +4257,7 @@ mod tests { ); TrieNode::Extension(extension_node) }, - masks: TrieMasks { hash_mask: None, tree_mask: None }, + masks: None, }, // Leaf at 0x4f88072c ProofTrieNode { @@ -4287,7 +4274,7 @@ mod tests { ); TrieNode::Leaf(leaf_node) }, - masks: TrieMasks { hash_mask: None, tree_mask: None }, + masks: None, }, ]; @@ -4299,7 +4286,7 @@ mod tests { "56fab2b106a97eae9c7197f86d03bca292da6e0ac725b783082f7d950cc4e0fc" ))), )), - TrieMasks::none(), + None, true, ) .unwrap(); @@ -4390,11 +4377,11 @@ mod tests { ); // Step 2: Reveal nodes in the trie - let mut trie = ParallelSparseTrie::from_root(extension, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(extension, None, true).unwrap(); trie.reveal_nodes(vec![ - ProofTrieNode { path: branch_path, node: branch, masks: TrieMasks::none() }, - ProofTrieNode { path: leaf_1_path, node: leaf_1, masks: TrieMasks::none() }, - ProofTrieNode { path: leaf_2_path, node: leaf_2, masks: TrieMasks::none() }, + ProofTrieNode { path: branch_path, node: branch, masks: None }, + ProofTrieNode { path: leaf_1_path, node: leaf_1, masks: None }, + ProofTrieNode { path: leaf_2_path, node: leaf_2, masks: None }, ]) .unwrap(); @@ -4917,7 +4904,10 @@ mod tests { let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::from_root( branch.clone(), - TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None }, + Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b01), + tree_mask: TrieMask::default(), + }), false, ) .unwrap(); @@ -4932,12 +4922,15 @@ mod tests { ProofTrieNode { path: Nibbles::default(), node: branch, - masks: TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) }, + masks: Some(BranchNodeMasks { + hash_mask: TrieMask::default(), + tree_mask: TrieMask::new(0b01), + }), }, ProofTrieNode { path: Nibbles::from_nibbles([0x1]), node: TrieNode::Leaf(leaf), - masks: TrieMasks::none(), + masks: None, }, ]) .unwrap(); @@ -4966,7 +4959,10 @@ mod tests { let provider = DefaultTrieNodeProvider; let mut sparse = ParallelSparseTrie::from_root( branch.clone(), - TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None }, + Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b01), + tree_mask: TrieMask::default(), + }), false, ) .unwrap(); @@ -4981,12 +4977,15 @@ mod tests { ProofTrieNode { path: Nibbles::default(), node: branch, - masks: TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) }, + masks: Some(BranchNodeMasks { + hash_mask: TrieMask::default(), + tree_mask: TrieMask::new(0b01), + }), }, ProofTrieNode { path: Nibbles::from_nibbles([0x1]), node: TrieNode::Leaf(leaf), - masks: TrieMasks::none(), + masks: None, }, ]) .unwrap(); @@ -5312,12 +5311,22 @@ mod tests { ); let provider = DefaultTrieNodeProvider; + let masks = match ( + branch_node_hash_masks.get(&Nibbles::default()).copied(), + branch_node_tree_masks.get(&Nibbles::default()).copied(), + ) { + (Some(h), Some(t)) => Some(BranchNodeMasks { hash_mask: h, tree_mask: t }), + (Some(h), None) => { + Some(BranchNodeMasks { hash_mask: h, tree_mask: TrieMask::default() }) + } + (None, Some(t)) => { + Some(BranchNodeMasks { hash_mask: TrieMask::default(), tree_mask: t }) + } + (None, None) => None, + }; let mut sparse = ParallelSparseTrie::from_root( TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), - TrieMasks { - hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(), - tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(), - }, + masks, false, ) .unwrap(); @@ -5336,11 +5345,8 @@ mod tests { .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - ProofTrieNode { - path, - node: TrieNode::decode(&mut &node[..]).unwrap(), - masks: TrieMasks { hash_mask, tree_mask }, - } + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(revealed_nodes).unwrap(); @@ -5374,11 +5380,8 @@ mod tests { .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - ProofTrieNode { - path, - node: TrieNode::decode(&mut &node[..]).unwrap(), - masks: TrieMasks { hash_mask, tree_mask }, - } + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(revealed_nodes).unwrap(); @@ -5428,12 +5431,22 @@ mod tests { ); let provider = DefaultTrieNodeProvider; + let masks = match ( + branch_node_hash_masks.get(&Nibbles::default()).copied(), + branch_node_tree_masks.get(&Nibbles::default()).copied(), + ) { + (Some(h), Some(t)) => Some(BranchNodeMasks { hash_mask: h, tree_mask: t }), + (Some(h), None) => { + Some(BranchNodeMasks { hash_mask: h, tree_mask: TrieMask::default() }) + } + (None, Some(t)) => { + Some(BranchNodeMasks { hash_mask: TrieMask::default(), tree_mask: t }) + } + (None, None) => None, + }; let mut sparse = ParallelSparseTrie::from_root( TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), - TrieMasks { - hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(), - tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(), - }, + masks, false, ) .unwrap(); @@ -5453,11 +5466,8 @@ mod tests { .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - ProofTrieNode { - path, - node: TrieNode::decode(&mut &node[..]).unwrap(), - masks: TrieMasks { hash_mask, tree_mask }, - } + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(revealed_nodes).unwrap(); @@ -5491,11 +5501,8 @@ mod tests { .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - ProofTrieNode { - path, - node: TrieNode::decode(&mut &node[..]).unwrap(), - masks: TrieMasks { hash_mask, tree_mask }, - } + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(revealed_nodes).unwrap(); @@ -5537,12 +5544,22 @@ mod tests { ); let provider = DefaultTrieNodeProvider; + let masks = match ( + branch_node_hash_masks.get(&Nibbles::default()).copied(), + branch_node_tree_masks.get(&Nibbles::default()).copied(), + ) { + (Some(h), Some(t)) => Some(BranchNodeMasks { hash_mask: h, tree_mask: t }), + (Some(h), None) => { + Some(BranchNodeMasks { hash_mask: h, tree_mask: TrieMask::default() }) + } + (None, Some(t)) => { + Some(BranchNodeMasks { hash_mask: TrieMask::default(), tree_mask: t }) + } + (None, None) => None, + }; let mut sparse = ParallelSparseTrie::from_root( TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), - TrieMasks { - hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(), - tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(), - }, + masks, false, ) .unwrap(); @@ -5576,11 +5593,8 @@ mod tests { .map(|(path, node)| { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - ProofTrieNode { - path, - node: TrieNode::decode(&mut &node[..]).unwrap(), - masks: TrieMasks { hash_mask, tree_mask }, - } + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + ProofTrieNode { path, node: TrieNode::decode(&mut &node[..]).unwrap(), masks } }) .collect(); sparse.reveal_nodes(revealed_nodes).unwrap(); @@ -5595,8 +5609,7 @@ mod tests { #[test] fn test_update_leaf_cross_level() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test adding leaves that demonstrate the cross-level behavior // Based on the example: leaves 0x1234, 0x1245, 0x1334, 0x1345 @@ -5676,8 +5689,7 @@ mod tests { #[test] fn test_update_leaf_split_at_level_boundary() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // This test demonstrates what happens when we insert leaves that cause // splitting exactly at the upper/lower trie boundary (2 nibbles). @@ -5724,8 +5736,7 @@ mod tests { #[test] fn test_update_subtrie_with_multiple_leaves() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // First, add multiple leaves that will create a subtrie structure // All leaves share the prefix [0x1, 0x2] to ensure they create a subtrie @@ -5793,8 +5804,7 @@ mod tests { #[test] fn test_update_subtrie_extension_node_subtrie() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // All leaves share the prefix [0x1, 0x2] to ensure they create a subtrie // @@ -5823,8 +5833,7 @@ mod tests { #[test] fn update_subtrie_extension_node_cross_level() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // First, add multiple leaves that will create a subtrie structure // All leaves share the prefix [0x1, 0x2] to ensure they create a branch node and subtrie @@ -5856,8 +5865,7 @@ mod tests { #[test] fn test_update_single_nibble_paths() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: single nibble paths that create branches in upper trie // @@ -5901,8 +5909,7 @@ mod tests { #[test] fn test_update_deep_extension_chain() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: deep extension chains that span multiple levels // @@ -5946,8 +5953,7 @@ mod tests { #[test] fn test_update_branch_with_all_nibbles() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: branch node with all 16 possible nibble children // @@ -5996,8 +6002,7 @@ mod tests { #[test] fn test_update_creates_multiple_subtries() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: updates that create multiple subtries at once // @@ -6050,8 +6055,7 @@ mod tests { #[test] fn test_update_extension_to_branch_transformation() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: extension node transforms to branch when split // @@ -6167,8 +6171,7 @@ mod tests { #[test] fn test_update_long_shared_prefix_at_boundary() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: leaves with long shared prefix that ends exactly at 2-nibble boundary // @@ -6211,8 +6214,7 @@ mod tests { #[test] fn test_update_branch_to_extension_collapse() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test creating a trie with leaves that share a long common prefix // @@ -6257,8 +6259,7 @@ mod tests { let (new_leaf3_path, new_value3) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x6], 12); // Clear and add new leaves - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); trie.update_leaf(new_leaf1_path, new_value1.clone(), DefaultTrieNodeProvider).unwrap(); trie.update_leaf(new_leaf2_path, new_value2.clone(), DefaultTrieNodeProvider).unwrap(); trie.update_leaf(new_leaf3_path, new_value3.clone(), DefaultTrieNodeProvider).unwrap(); @@ -6284,8 +6285,7 @@ mod tests { #[test] fn test_update_shared_prefix_patterns() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: different patterns of shared prefixes // @@ -6328,8 +6328,7 @@ mod tests { #[test] fn test_progressive_branch_creation() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test starting with a single leaf and progressively adding leaves // that create branch nodes at shorter and shorter paths @@ -6437,8 +6436,7 @@ mod tests { #[test] fn test_update_max_depth_paths() { let ctx = ParallelSparseTrieTestContext; - let mut trie = - ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap(); + let mut trie = ParallelSparseTrie::from_root(TrieNode::EmptyRoot, None, true).unwrap(); // Test edge case: very long paths (64 nibbles - max for addresses/storage) // @@ -6506,10 +6504,10 @@ mod tests { TrieMask::new(0b1111111111111111), // state_mask: all 16 children present ); - let root_branch_masks = TrieMasks { - hash_mask: Some(TrieMask::new(0b1111111111111111)), - tree_mask: Some(TrieMask::new(0b1111111111111111)), - }; + let root_branch_masks = Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b1111111111111111), + tree_mask: TrieMask::new(0b1111111111111111), + }); let mut trie = ParallelSparseTrie::from_root( TrieNode::Branch(root_branch_node), @@ -6548,10 +6546,10 @@ mod tests { TrieMask::new(0b1111111111111111), // state_mask: all 16 children present ); - let branch_0x3_masks = TrieMasks { - hash_mask: Some(TrieMask::new(0b0100010000010101)), - tree_mask: Some(TrieMask::new(0b0100000000000000)), - }; + let branch_0x3_masks = Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b0100010000010101), + tree_mask: TrieMask::new(0b0100000000000000), + }); // Reveal node at path Nibbles(0x37) - leaf node let leaf_path = Nibbles::from_nibbles([0x3, 0x7]); @@ -6561,7 +6559,7 @@ mod tests { let leaf_value = hex!("f8440180a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0f57acd40259872606d76197ef052f3d35588dadf919ee1f0e3cb9b62d3f4b02c").to_vec(); let leaf_node = LeafNode::new(leaf_key, leaf_value); - let leaf_masks = TrieMasks::none(); + let leaf_masks = None; trie.reveal_nodes(vec![ ProofTrieNode { diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 1d68f19cd7..a1c6ac4e9f 100644 --- a/crates/trie/sparse/src/state.rs +++ b/crates/trie/sparse/src/state.rs @@ -15,8 +15,8 @@ use reth_primitives_traits::Account; use reth_trie_common::{ proof::ProofNodes, updates::{StorageTrieUpdates, TrieUpdates}, - BranchNodeMasksMap, DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles, - ProofTrieNode, RlpNode, StorageMultiProof, TrieAccount, TrieMasks, TrieNode, EMPTY_ROOT_HASH, + BranchNodeMasks, BranchNodeMasksMap, DecodedMultiProof, DecodedStorageMultiProof, MultiProof, + Nibbles, ProofTrieNode, RlpNode, StorageMultiProof, TrieAccount, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; use tracing::{instrument, trace}; @@ -497,17 +497,13 @@ where if path.is_empty() { // Handle special storage state root node case. - storage_trie_entry.reveal_root( - trie_node, - TrieMasks::none(), - retain_updates, - )?; + storage_trie_entry.reveal_root(trie_node, None, retain_updates)?; } else { // Reveal non-root storage trie node. storage_trie_entry .as_revealed_mut() .ok_or(SparseTrieErrorKind::Blind)? - .reveal_node(path, trie_node, TrieMasks::none())?; + .reveal_node(path, trie_node, None)?; } // Track the revealed path. @@ -518,14 +514,13 @@ where else if !self.revealed_account_paths.contains(&path) { if path.is_empty() { // Handle special state root node case. - self.state.reveal_root(trie_node, TrieMasks::none(), self.retain_updates)?; + self.state.reveal_root(trie_node, None, self.retain_updates)?; } else { // Reveal non-root state trie node. - self.state.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?.reveal_node( - path, - trie_node, - TrieMasks::none(), - )?; + self.state + .as_revealed_mut() + .ok_or(SparseTrieErrorKind::Blind)? + .reveal_node(path, trie_node, None)?; } // Track the revealed path. @@ -577,9 +572,8 @@ where }) .transpose()? .unwrap_or((TrieNode::EmptyRoot, None, None)); - self.state - .reveal_root(root_node, TrieMasks { hash_mask, tree_mask }, self.retain_updates) - .map_err(Into::into) + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + self.state.reveal_root(root_node, masks, self.retain_updates).map_err(Into::into) } SparseTrie::Revealed(ref mut trie) => Ok(trie), } @@ -973,21 +967,14 @@ 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; - 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() - } + branch_node_masks.get(&path).copied() } TrieNode::Extension(_) => { // There is always exactly one child of an extension node. result.new_nodes += 1; - TrieMasks::none() + None } - _ => TrieMasks::none(), + _ => None, }; let node = ProofTrieNode { path, node: proof_node, masks }; @@ -1369,12 +1356,12 @@ mod tests { root_node: Some(ProofTrieNode { path: Nibbles::default(), node: branch, - masks: TrieMasks::none(), + masks: None, }), nodes: vec![ProofTrieNode { path: Nibbles::from_nibbles([0x1]), node: leaf, - masks: TrieMasks::none(), + masks: None, }], // Branch, two of its children, one leaf new_nodes: 4, diff --git a/crates/trie/sparse/src/traits.rs b/crates/trie/sparse/src/traits.rs index 55a17a9a7f..1d652284a4 100644 --- a/crates/trie/sparse/src/traits.rs +++ b/crates/trie/sparse/src/traits.rs @@ -9,7 +9,7 @@ use alloy_primitives::{ }; use alloy_trie::BranchNodeCompact; use reth_execution_errors::SparseTrieResult; -use reth_trie_common::{Nibbles, ProofTrieNode, TrieMasks, TrieNode}; +use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNode, TrieNode}; use crate::provider::TrieNodeProvider; @@ -37,7 +37,7 @@ pub trait SparseTrieInterface: Sized + Debug + Send + Sync { fn with_root( self, root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult; @@ -72,7 +72,7 @@ pub trait SparseTrieInterface: Sized + Debug + Send + Sync { &mut self, path: Nibbles, node: TrieNode, - masks: TrieMasks, + masks: Option, ) -> SparseTrieResult<()> { self.reveal_nodes(vec![ProofTrieNode { path, node, masks }]) } diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index dd2f9cdacc..bb96a1d4c1 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -20,7 +20,7 @@ use reth_execution_errors::{SparseTrieErrorKind, SparseTrieResult}; use reth_trie_common::{ prefix_set::{PrefixSet, PrefixSetMut}, BranchNodeCompact, BranchNodeMasks, BranchNodeMasksMap, BranchNodeRef, ExtensionNodeRef, - LeafNodeRef, Nibbles, ProofTrieNode, RlpNode, TrieMask, TrieMasks, TrieNode, CHILD_INDEX_RANGE, + LeafNodeRef, Nibbles, ProofTrieNode, RlpNode, TrieMask, TrieNode, CHILD_INDEX_RANGE, EMPTY_ROOT_HASH, }; use smallvec::SmallVec; @@ -95,7 +95,7 @@ impl SparseTrie { pub fn reveal_root( &mut self, root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult<&mut T> { // if `Blind`, we initialize the revealed trie with the given root node, using a @@ -419,7 +419,7 @@ impl SparseTrieInterface for SerialSparseTrie { fn with_root( mut self, root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult { self = self.with_updates(retain_updates); @@ -448,7 +448,7 @@ impl SparseTrieInterface for SerialSparseTrie { &mut self, path: Nibbles, node: TrieNode, - masks: TrieMasks, + masks: Option, ) -> SparseTrieResult<()> { trace!(target: "trie::sparse", ?path, ?node, ?masks, "reveal_node called"); @@ -457,14 +457,8 @@ impl SparseTrieInterface for SerialSparseTrie { return Ok(()) } - if masks.tree_mask.is_some() || masks.hash_mask.is_some() { - self.branch_node_masks.insert( - path, - BranchNodeMasks { - tree_mask: masks.tree_mask.unwrap_or_default(), - hash_mask: masks.hash_mask.unwrap_or_default(), - }, - ); + if let Some(branch_masks) = masks { + self.branch_node_masks.insert(path, branch_masks); } match node { @@ -496,10 +490,9 @@ impl SparseTrieInterface for SerialSparseTrie { // Memoize the hash of a previously blinded node in a new branch // node. hash: Some(*hash), - store_in_db_trie: Some( - masks.hash_mask.is_some_and(|mask| !mask.is_empty()) || - masks.tree_mask.is_some_and(|mask| !mask.is_empty()), - ), + store_in_db_trie: Some(masks.is_some_and(|m| { + !m.hash_mask.is_empty() || !m.tree_mask.is_empty() + })), }); } // Branch node already exists, or an extension node was placed where a @@ -691,11 +684,9 @@ impl SparseTrieInterface for SerialSparseTrie { ?hash_mask, "Revealing extension node child", ); - self.reveal_node( - current, - decoded, - TrieMasks { hash_mask, tree_mask }, - )?; + let masks = + BranchNodeMasks::from_optional(hash_mask, tree_mask); + self.reveal_node(current, decoded, masks)?; } } } @@ -1115,7 +1106,7 @@ impl SerialSparseTrie { /// Self if successful, or an error if revealing fails. pub fn from_root( root: TrieNode, - masks: TrieMasks, + masks: Option, retain_updates: bool, ) -> SparseTrieResult { Self::default().with_root(root, masks, retain_updates) @@ -1173,7 +1164,7 @@ impl SerialSparseTrie { return Ok(()) } - self.reveal_node(path, TrieNode::decode(&mut &child[..])?, TrieMasks::none()) + self.reveal_node(path, TrieNode::decode(&mut &child[..])?, None) } /// Traverse the trie from the root down to the leaf at the given path, @@ -1308,11 +1299,8 @@ impl SerialSparseTrie { ?hash_mask, "Revealing remaining blinded branch child" ); - self.reveal_node( - *remaining_child_path, - decoded, - TrieMasks { hash_mask, tree_mask }, - )?; + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + self.reveal_node(*remaining_child_path, decoded, masks)?; self.nodes.get(remaining_child_path).unwrap().clone() } else { return Err(SparseTrieErrorKind::NodeNotFoundInProvider { @@ -2302,7 +2290,7 @@ mod find_leaf_tests { // 3. Initialize the sparse trie using from_root // This will internally create Hash nodes for paths "1" and "5" initially. - let mut sparse = SerialSparseTrie::from_root(root_trie_node, TrieMasks::none(), false) + let mut sparse = SerialSparseTrie::from_root(root_trie_node, None, false) .expect("Failed to create trie from root"); // Assertions before we reveal child5 @@ -2313,7 +2301,7 @@ mod find_leaf_tests { // 4. Explicitly reveal the leaf node for child 5 sparse - .reveal_node(revealed_leaf_prefix, TrieNode::Leaf(leaf_node_child5), TrieMasks::none()) + .reveal_node(revealed_leaf_prefix, TrieNode::Leaf(leaf_node_child5), None) .expect("Failed to reveal leaf node"); // Assertions after we reveal child 5 @@ -2950,7 +2938,10 @@ mod tests { let provider = DefaultTrieNodeProvider; let mut sparse = SerialSparseTrie::from_root( branch.clone(), - TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None }, + Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b01), + tree_mask: TrieMask::default(), + }), false, ) .unwrap(); @@ -2964,12 +2955,13 @@ mod tests { .reveal_node( Nibbles::default(), branch, - TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) }, + Some(BranchNodeMasks { + hash_mask: TrieMask::default(), + tree_mask: TrieMask::new(0b01), + }), ) .unwrap(); - sparse - .reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), TrieMasks::none()) - .unwrap(); + sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); // Removing a blinded leaf should result in an error assert_matches!( @@ -2995,7 +2987,10 @@ mod tests { let provider = DefaultTrieNodeProvider; let mut sparse = SerialSparseTrie::from_root( branch.clone(), - TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None }, + Some(BranchNodeMasks { + hash_mask: TrieMask::new(0b01), + tree_mask: TrieMask::default(), + }), false, ) .unwrap(); @@ -3009,12 +3004,13 @@ mod tests { .reveal_node( Nibbles::default(), branch, - TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) }, + Some(BranchNodeMasks { + hash_mask: TrieMask::default(), + tree_mask: TrieMask::new(0b01), + }), ) .unwrap(); - sparse - .reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), TrieMasks::none()) - .unwrap(); + sparse.reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), None).unwrap(); // Removing a non-existent leaf should be a noop let sparse_old = sparse.clone(); @@ -3200,12 +3196,13 @@ mod tests { ); let provider = DefaultTrieNodeProvider; + let masks = BranchNodeMasks::from_optional( + branch_node_hash_masks.get(&Nibbles::default()).copied(), + branch_node_tree_masks.get(&Nibbles::default()).copied(), + ); let mut sparse = SerialSparseTrie::from_root( TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), - TrieMasks { - hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(), - tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(), - }, + masks, false, ) .unwrap(); @@ -3221,13 +3218,8 @@ mod tests { for (path, node) in hash_builder_proof_nodes.nodes_sorted() { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - sparse - .reveal_node( - path, - TrieNode::decode(&mut &node[..]).unwrap(), - TrieMasks { hash_mask, tree_mask }, - ) - .unwrap(); + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + sparse.reveal_node(path, TrieNode::decode(&mut &node[..]).unwrap(), masks).unwrap(); } // Check that the branch node exists with only two nibbles set @@ -3256,13 +3248,8 @@ mod tests { for (path, node) in hash_builder_proof_nodes.nodes_sorted() { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - sparse - .reveal_node( - path, - TrieNode::decode(&mut &node[..]).unwrap(), - TrieMasks { hash_mask, tree_mask }, - ) - .unwrap(); + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + sparse.reveal_node(path, TrieNode::decode(&mut &node[..]).unwrap(), masks).unwrap(); } // Check that nothing changed in the branch node @@ -3310,12 +3297,13 @@ mod tests { ); let provider = DefaultTrieNodeProvider; + let masks = BranchNodeMasks::from_optional( + branch_node_hash_masks.get(&Nibbles::default()).copied(), + branch_node_tree_masks.get(&Nibbles::default()).copied(), + ); let mut sparse = SerialSparseTrie::from_root( TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), - TrieMasks { - hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(), - tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(), - }, + masks, false, ) .unwrap(); @@ -3332,13 +3320,8 @@ mod tests { for (path, node) in hash_builder_proof_nodes.nodes_sorted() { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - sparse - .reveal_node( - path, - TrieNode::decode(&mut &node[..]).unwrap(), - TrieMasks { hash_mask, tree_mask }, - ) - .unwrap(); + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + sparse.reveal_node(path, TrieNode::decode(&mut &node[..]).unwrap(), masks).unwrap(); } // Check that the branch node exists @@ -3367,13 +3350,8 @@ mod tests { for (path, node) in hash_builder_proof_nodes.nodes_sorted() { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - sparse - .reveal_node( - path, - TrieNode::decode(&mut &node[..]).unwrap(), - TrieMasks { hash_mask, tree_mask }, - ) - .unwrap(); + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + sparse.reveal_node(path, TrieNode::decode(&mut &node[..]).unwrap(), masks).unwrap(); } // Check that nothing changed in the extension node @@ -3413,12 +3391,13 @@ mod tests { ); let provider = DefaultTrieNodeProvider; + let masks = BranchNodeMasks::from_optional( + branch_node_hash_masks.get(&Nibbles::default()).copied(), + branch_node_tree_masks.get(&Nibbles::default()).copied(), + ); let mut sparse = SerialSparseTrie::from_root( TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(), - TrieMasks { - hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(), - tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(), - }, + masks, false, ) .unwrap(); @@ -3449,13 +3428,8 @@ mod tests { for (path, node) in hash_builder_proof_nodes.nodes_sorted() { let hash_mask = branch_node_hash_masks.get(&path).copied(); let tree_mask = branch_node_tree_masks.get(&path).copied(); - sparse - .reveal_node( - path, - TrieNode::decode(&mut &node[..]).unwrap(), - TrieMasks { hash_mask, tree_mask }, - ) - .unwrap(); + let masks = BranchNodeMasks::from_optional(hash_mask, tree_mask); + sparse.reveal_node(path, TrieNode::decode(&mut &node[..]).unwrap(), masks).unwrap(); } // Check that the branch node wasn't overwritten by the extension node in the proof diff --git a/crates/trie/trie/src/proof_v2/mod.rs b/crates/trie/trie/src/proof_v2/mod.rs index eb4ee5db7e..f2402d0591 100644 --- a/crates/trie/trie/src/proof_v2/mod.rs +++ b/crates/trie/trie/src/proof_v2/mod.rs @@ -15,7 +15,7 @@ use alloy_primitives::{B256, U256}; use alloy_rlp::Encodable; use alloy_trie::{BranchNodeCompact, TrieMask}; use reth_execution_errors::trie::StateProofError; -use reth_trie_common::{BranchNode, Nibbles, ProofTrieNode, RlpNode, TrieMasks, TrieNode}; +use reth_trie_common::{BranchNode, BranchNodeMasks, Nibbles, ProofTrieNode, RlpNode, TrieNode}; use std::cmp::Ordering; use tracing::{instrument, trace}; @@ -452,7 +452,7 @@ where self.branch_stack.push(ProofTrieBranch { ext_len: common_prefix_len as u8, state_mask: TrieMask::new(1 << first_child_nibble), - masks: TrieMasks::none(), + masks: None, }); trace!( @@ -696,10 +696,10 @@ where ProofTrieBranch { ext_len, state_mask: TrieMask::new(0), - masks: TrieMasks { - tree_mask: Some(cached_branch.tree_mask), - hash_mask: Some(cached_branch.hash_mask), - }, + masks: Some(BranchNodeMasks { + tree_mask: cached_branch.tree_mask, + hash_mask: cached_branch.hash_mask, + }), } } @@ -1226,7 +1226,7 @@ where self.retained_proofs.push(ProofTrieNode { path: Nibbles::new(), // root path node: TrieNode::EmptyRoot, - masks: TrieMasks::none(), + masks: None, }); } (true, Some(root_node)) => { @@ -1343,7 +1343,7 @@ where return Ok(vec![ProofTrieNode { path: Nibbles::default(), node: TrieNode::EmptyRoot, - masks: TrieMasks::none(), + masks: None, }]) } @@ -1647,16 +1647,9 @@ mod tests { // The legacy proof calculator will calculate masks for the root node, even // 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 if let Some(branch_masks) = - proof_legacy_result.branch_node_masks.get(path) - { - TrieMasks { - hash_mask: Some(branch_masks.hash_mask), - tree_mask: Some(branch_masks.tree_mask), - } + None } else { - TrieMasks::none() + proof_legacy_result.branch_node_masks.get(path).copied() }; ProofTrieNode { path: *path, node, masks } diff --git a/crates/trie/trie/src/proof_v2/node.rs b/crates/trie/trie/src/proof_v2/node.rs index 9300123fbe..6acdaae881 100644 --- a/crates/trie/trie/src/proof_v2/node.rs +++ b/crates/trie/trie/src/proof_v2/node.rs @@ -3,8 +3,8 @@ use alloy_rlp::Encodable; use alloy_trie::nodes::ExtensionNodeRef; use reth_execution_errors::trie::StateProofError; use reth_trie_common::{ - BranchNode, ExtensionNode, LeafNode, LeafNodeRef, Nibbles, ProofTrieNode, RlpNode, TrieMask, - TrieMasks, TrieNode, + BranchNode, BranchNodeMasks, ExtensionNode, LeafNode, LeafNodeRef, Nibbles, ProofTrieNode, + RlpNode, TrieMask, TrieNode, }; /// A trie node which is the child of a branch in the trie. @@ -29,7 +29,7 @@ pub(crate) enum ProofTrieBranchChild { /// The node itself, for use during RLP encoding. node: BranchNode, /// Bitmasks carried over from cached `BranchNodeCompact` values, if any. - masks: TrieMasks, + masks: Option, }, /// A node whose type is not known, as it has already been converted to an [`RlpNode`]. RlpNode(RlpNode), @@ -98,10 +98,10 @@ impl ProofTrieBranchChild { // this value, and the passed in buffer can remain with whatever large capacity it // already has. let rlp_val = buf.clone(); - (TrieNode::Leaf(LeafNode::new(short_key, rlp_val)), TrieMasks::none()) + (TrieNode::Leaf(LeafNode::new(short_key, rlp_val)), None) } Self::Extension { short_key, child } => { - (TrieNode::Extension(ExtensionNode { key: short_key, child }), TrieMasks::none()) + (TrieNode::Extension(ExtensionNode { key: short_key, child }), None) } Self::Branch { node, masks } => (TrieNode::Branch(node), masks), Self::RlpNode(_) => panic!("Cannot call `into_proof_trie_node` on RlpNode"), @@ -158,7 +158,7 @@ pub(crate) struct ProofTrieBranch { /// child on the stack for each set bit. pub(crate) state_mask: TrieMask, /// Bitmasks which are subsets of `state_mask`. - pub(crate) masks: TrieMasks, + pub(crate) masks: Option, } /// Trims the first `len` nibbles from the head of the given `Nibbles`.