From 81942e41be5de28b11f25fe4220b8effe4cfe46b Mon Sep 17 00:00:00 2001 From: kevaundray Date: Tue, 22 Apr 2025 20:52:54 +0100 Subject: [PATCH] chore: add a method on SparseTrie to check whether a leaf exists (#15758) Co-authored-by: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com> --- crates/trie/sparse/src/trie.rs | 504 +++++++++++++++++++++++++++++++++ 1 file changed, 504 insertions(+) diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 17bf19da9c..2f72479180 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -1057,7 +1057,189 @@ impl

RevealedSparseTrie

{ } } +/// Error type for a leaf lookup operation +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LeafLookupError { + /// The path leads to a blinded node, cannot determine if leaf exists. + /// This means the witness is not complete. + BlindedNode { + /// Path to the blinded node. + path: Nibbles, + /// Hash of the blinded node. + hash: B256, + }, + /// The path leads to a leaf with a different value than expected. + /// This means the witness is malformed. + ValueMismatch { + /// Path to the leaf. + path: Nibbles, + /// Expected value. + expected: Option>, + /// Actual value found. + actual: Vec, + }, +} + +/// Success value for a leaf lookup operation +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LeafLookup { + /// Leaf exists with expected value. + Exists, + /// Leaf does not exist (exclusion proof found). + NonExistent { + /// Path where the search diverged from the target path. + diverged_at: Nibbles, + }, +} + impl RevealedSparseTrie

{ + /// Attempts to find a leaf node at the specified path. + /// + /// This method traverses the trie from the root down to the given path, checking + /// if a leaf exists at that path. It can be used to verify the existence of a leaf + /// or to generate an exclusion proof (proof that a leaf does not exist). + /// + /// # Parameters + /// + /// - `path`: The path to search for. + /// - `expected_value`: Optional expected value. If provided, will verify the leaf value + /// matches. + /// + /// # Returns + /// + /// - `Ok(LeafLookup::Exists)` if the leaf exists with the expected value. + /// - `Ok(LeafLookup::NonExistent)` if the leaf definitely does not exist (exclusion proof). + /// - `Err(LeafLookupError)` if the search encountered a blinded node or found a different + /// value. + pub fn find_leaf( + &self, + path: &Nibbles, + expected_value: Option<&Vec>, + ) -> Result { + // Helper function to check if a value matches the expected value + fn check_value_match( + actual_value: &Vec, + expected_value: Option<&Vec>, + path: &Nibbles, + ) -> Result<(), LeafLookupError> { + if let Some(expected) = expected_value { + if actual_value != expected { + return Err(LeafLookupError::ValueMismatch { + path: path.clone(), + expected: Some(expected.clone()), + actual: actual_value.clone(), + }); + } + } + Ok(()) + } + + let mut current = Nibbles::default(); // Start at the root + + // Inclusion proof + // + // First, do a quick check if the value exists in our values map. + // We assume that if there exists a leaf node, then its value will + // be in the `values` map. + if let Some(actual_value) = self.values.get(path) { + // We found the leaf, check if the value matches (if expected value was provided) + check_value_match(actual_value, expected_value, path)?; + return Ok(LeafLookup::Exists); + } + + // If the value does not exist in the `values` map, then this means that the leaf either: + // - Does not exist in the trie + // - Is missing from the witness + // We traverse the trie to find the location where this leaf would have been, showing + // that it is not in the trie. Or we find a blinded node, showing that the witness is + // not complete. + while current.len() < path.len() { + match self.nodes.get(¤t) { + Some(SparseNode::Empty) | None => { + // None implies no node is at the current path (even in the full trie) + // Empty node means there is a node at this path and it is "Empty" + return Ok(LeafLookup::NonExistent { diverged_at: current }); + } + Some(&SparseNode::Hash(hash)) => { + // We hit a blinded node - cannot determine if leaf exists + return Err(LeafLookupError::BlindedNode { path: current.clone(), hash }); + } + Some(SparseNode::Leaf { key, .. }) => { + // We found a leaf node before reaching our target depth + + // Temporarily append the leaf key to `current` + let saved_len = current.len(); + current.extend_from_slice_unchecked(key); + + if ¤t == path { + // This should have been handled by our initial values map check + if let Some(value) = self.values.get(path) { + check_value_match(value, expected_value, path)?; + return Ok(LeafLookup::Exists); + } + } + + let diverged_at = current.slice(..saved_len); + + // The leaf node's path doesn't match our target path, + // providing an exclusion proof + return Ok(LeafLookup::NonExistent { diverged_at }); + } + Some(SparseNode::Extension { key, .. }) => { + // Temporarily append the extension key to `current` + let saved_len = current.len(); + current.extend_from_slice_unchecked(key); + + if path.len() < current.len() || !path.starts_with(¤t) { + let diverged_at = current.slice(..saved_len); + current.truncate(saved_len); // restore + return Ok(LeafLookup::NonExistent { diverged_at }); + } + // Prefix matched, so we keep walking with the longer `current`. + } + Some(SparseNode::Branch { state_mask, .. }) => { + // Check if branch has a child at the next nibble in our path + let nibble = path[current.len()]; + if !state_mask.is_bit_set(nibble) { + // No child at this nibble - exclusion proof + return Ok(LeafLookup::NonExistent { diverged_at: current }); + } + + // Continue down the branch + current.push_unchecked(nibble); + } + } + } + + // We've traversed to the end of the path and didn't find a leaf + // Check if there's a node exactly at our target path + match self.nodes.get(path) { + Some(SparseNode::Leaf { key, .. }) if key.is_empty() => { + // We found a leaf with an empty key (exact match) + // This should be handled by the values map check above + if let Some(value) = self.values.get(path) { + check_value_match(value, expected_value, path)?; + return Ok(LeafLookup::Exists); + } + } + Some(&SparseNode::Hash(hash)) => { + return Err(LeafLookupError::BlindedNode { path: path.clone(), hash }); + } + _ => { + // No leaf at exactly the target path + let parent_path = if path.is_empty() { + Nibbles::default() + } else { + path.slice(0..path.len() - 1) + }; + return Ok(LeafLookup::NonExistent { diverged_at: parent_path }); + } + } + + // If we get here, there's no leaf at the target path + Ok(LeafLookup::NonExistent { diverged_at: current }) + } + /// Update the leaf node with provided value. pub fn update_leaf(&mut self, path: Nibbles, value: Vec) -> SparseTrieResult<()> { self.prefix_set.insert(path.clone()); @@ -1578,6 +1760,328 @@ impl SparseTrieUpdates { } } +#[cfg(test)] +mod find_leaf_tests { + use super::*; + use crate::blinded::DefaultBlindedProvider; + use alloy_primitives::map::foldhash::fast::RandomState; + // Assuming this exists + use alloy_rlp::Encodable; + use assert_matches::assert_matches; + use reth_primitives_traits::Account; + use reth_trie_common::LeafNode; + + // Helper to create some test values + fn encode_value(nonce: u64) -> Vec { + let account = Account { nonce, ..Default::default() }; + let trie_account = account.into_trie_account(EMPTY_ROOT_HASH); + let mut buf = Vec::new(); + trie_account.encode(&mut buf); + buf + } + + const VALUE_A: fn() -> Vec = || encode_value(1); + const VALUE_B: fn() -> Vec = || encode_value(2); + + #[test] + fn find_leaf_existing_leaf() { + // Create a simple trie with one leaf + let mut sparse = RevealedSparseTrie::default(); + let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]); + let value = b"test_value".to_vec(); + + sparse.update_leaf(path.clone(), value.clone()).unwrap(); + + // Check that the leaf exists + let result = sparse.find_leaf(&path, None); + assert_matches!(result, Ok(LeafLookup::Exists)); + + // Check with expected value matching + let result = sparse.find_leaf(&path, Some(&value)); + assert_matches!(result, Ok(LeafLookup::Exists)); + } + + #[test] + fn find_leaf_value_mismatch() { + // Create a simple trie with one leaf + let mut sparse = RevealedSparseTrie::default(); + let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]); + let value = b"test_value".to_vec(); + let wrong_value = b"wrong_value".to_vec(); + + sparse.update_leaf(path.clone(), value).unwrap(); + + // Check with wrong expected value + let result = sparse.find_leaf(&path, Some(&wrong_value)); + assert_matches!( + result, + Err(LeafLookupError::ValueMismatch { path: p, expected: Some(e), actual: _a }) if p == path && e == wrong_value + ); + } + + #[test] + fn find_leaf_not_found_empty_trie() { + // Empty trie + let sparse = RevealedSparseTrie::default(); + let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]); + + // Leaf should not exist + let result = sparse.find_leaf(&path, None); + assert_matches!( + result, + Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == Nibbles::default() + ); + } + + #[test] + fn find_leaf_empty_trie() { + let sparse = RevealedSparseTrie::::default(); + let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); + + let result = sparse.find_leaf(&path, None); + + // In an empty trie, the search diverges immediately at the root. + assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == Nibbles::default()); + } + + #[test] + fn find_leaf_exists_no_value_check() { + let mut sparse = RevealedSparseTrie::::default(); + let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); + sparse.update_leaf(path.clone(), VALUE_A()).unwrap(); + + let result = sparse.find_leaf(&path, None); + assert_matches!(result, Ok(LeafLookup::Exists)); + } + + #[test] + fn find_leaf_exists_with_value_check_ok() { + let mut sparse = RevealedSparseTrie::::default(); + let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); + let value = VALUE_A(); + sparse.update_leaf(path.clone(), value.clone()).unwrap(); + + let result = sparse.find_leaf(&path, Some(&value)); + assert_matches!(result, Ok(LeafLookup::Exists)); + } + + #[test] + fn find_leaf_exclusion_branch_divergence() { + let mut sparse = RevealedSparseTrie::::default(); + let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Creates branch at 0x12 + let path2 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x5, 0x6]); // Belongs to same branch + let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x7, 0x8]); // Diverges at nibble 7 + + sparse.update_leaf(path1, VALUE_A()).unwrap(); + sparse.update_leaf(path2, VALUE_B()).unwrap(); + + let result = sparse.find_leaf(&search_path, None); + + // Diverged at the branch node because nibble '7' is not present. + let expected_divergence = Nibbles::from_nibbles_unchecked([0x1, 0x2]); + assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence); + } + + #[test] + fn find_leaf_exclusion_extension_divergence() { + let mut sparse = RevealedSparseTrie::::default(); + // This will create an extension node at root with key 0x12 + let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4, 0x5, 0x6]); + // This path diverges from the extension key + let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x7, 0x8]); + + sparse.update_leaf(path1, VALUE_A()).unwrap(); + + let result = sparse.find_leaf(&search_path, None); + + // Diverged where the extension node started because the path doesn't match its key prefix. + let expected_divergence = Nibbles::default(); + assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence); + } + + #[test] + fn find_leaf_exclusion_leaf_divergence() { + let mut sparse = RevealedSparseTrie::::default(); + let existing_leaf_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); + let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4, 0x5, 0x6]); + + sparse.update_leaf(existing_leaf_path, VALUE_A()).unwrap(); + + let result = sparse.find_leaf(&search_path, None); + + // Diverged when it hit the leaf node at the root, because the search path is longer + // than the leaf's key stored there. The code returns the path of the node (root) + // where the divergence occurred. + let expected_divergence = Nibbles::default(); + assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence); + } + + #[test] + fn find_leaf_exclusion_path_ends_at_branch() { + let mut sparse = RevealedSparseTrie::::default(); + let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Creates branch at 0x12 + let path2 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x5, 0x6]); + let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2]); // Path of the branch itself + + sparse.update_leaf(path1, VALUE_A()).unwrap(); + sparse.update_leaf(path2, VALUE_B()).unwrap(); + + let result = sparse.find_leaf(&search_path, None); + + // The path ends, but the node at the path is a branch, not a leaf. + // Diverged at the parent of the node found at the search path. + let expected_divergence = Nibbles::from_nibbles_unchecked([0x1]); + assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence); + } + + #[test] + fn find_leaf_error_blinded_node_at_leaf_path() { + // Scenario: The node *at* the leaf path is blinded. + let blinded_hash = B256::repeat_byte(0xBB); + let leaf_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); + + let mut nodes = alloy_primitives::map::HashMap::with_hasher(RandomState::default()); + // Create path to the blinded node + nodes.insert( + Nibbles::default(), + SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x1, 0x2])), + ); // Ext 0x12 + nodes.insert( + Nibbles::from_nibbles_unchecked([0x1, 0x2]), + SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x3])), + ); // Ext 0x123 + nodes.insert( + Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3]), + SparseNode::new_branch(TrieMask::new(0b10000)), + ); // Branch at 0x123, child 4 + nodes.insert(leaf_path.clone(), SparseNode::Hash(blinded_hash)); // Blinded node at 0x1234 + + let sparse = RevealedSparseTrie { + provider: DefaultBlindedProvider, + nodes, + branch_node_tree_masks: Default::default(), + branch_node_hash_masks: Default::default(), + /* The value is not in the values map, or else it would early return */ + values: Default::default(), + prefix_set: Default::default(), + updates: None, + rlp_buf: Vec::new(), + }; + + let result = sparse.find_leaf(&leaf_path, None); + + // Should error because it hit the blinded node exactly at the leaf path + assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash }) + if path == leaf_path && hash == blinded_hash + ); + } + + #[test] + fn find_leaf_error_blinded_node() { + let blinded_hash = B256::repeat_byte(0xAA); + let path_to_blind = Nibbles::from_nibbles_unchecked([0x1]); + let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); + + let mut nodes = HashMap::with_hasher(RandomState::default()); + + // Root is a branch with child 0x1 (blinded) and 0x5 (revealed leaf) + // So we set Bit 1 and Bit 5 in the state_mask + let state_mask = TrieMask::new(0b100010); + nodes.insert(Nibbles::default(), SparseNode::new_branch(state_mask)); + + nodes.insert(path_to_blind.clone(), SparseNode::Hash(blinded_hash)); + let path_revealed = Nibbles::from_nibbles_unchecked([0x5]); + let path_revealed_leaf = Nibbles::from_nibbles_unchecked([0x5, 0x6, 0x7, 0x8]); + nodes.insert( + path_revealed, + SparseNode::new_leaf(Nibbles::from_nibbles_unchecked([0x6, 0x7, 0x8])), + ); + + let mut values = HashMap::with_hasher(RandomState::default()); + values.insert(path_revealed_leaf, VALUE_A()); + + let sparse = RevealedSparseTrie { + provider: DefaultBlindedProvider, + nodes, + branch_node_tree_masks: Default::default(), + branch_node_hash_masks: Default::default(), + values, + prefix_set: Default::default(), + updates: None, + rlp_buf: Vec::new(), + }; + + let result = sparse.find_leaf(&search_path, None); + + // Should error because it hit the blinded node at path 0x1 + assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash }) + if path == path_to_blind && hash == blinded_hash + ); + } + + #[test] + fn find_leaf_error_blinded_node_via_reveal() { + let blinded_hash = B256::repeat_byte(0xAA); + let path_to_blind = Nibbles::from_nibbles_unchecked([0x1]); // Path of the blinded node itself + let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Path we will search for + + let revealed_leaf_prefix = Nibbles::from_nibbles_unchecked([0x5]); + let revealed_leaf_suffix = Nibbles::from_nibbles_unchecked([0x6, 0x7, 0x8]); + let revealed_leaf_full_path = Nibbles::from_nibbles_unchecked([0x5, 0x6, 0x7, 0x8]); + let revealed_value = VALUE_A(); + + // 1. Construct the RLP representation of the children for the root branch + let rlp_node_child1 = RlpNode::word_rlp(&blinded_hash); // Blinded node + + let leaf_node_child5 = LeafNode::new(revealed_leaf_suffix.clone(), revealed_value.clone()); + let leaf_node_child5_rlp_buf = alloy_rlp::encode(&leaf_node_child5); + let hash_of_child5 = keccak256(&leaf_node_child5_rlp_buf); + let rlp_node_child5 = RlpNode::word_rlp(&hash_of_child5); + + // 2. Construct the root BranchNode using the RLP of its children + // The stack order depends on the bit indices (1 and 5) + let root_branch_node = reth_trie_common::BranchNode::new( + vec![rlp_node_child1, rlp_node_child5], // Child 1 first, then Child 5 + TrieMask::new(0b100010), // Mask with bits 1 and 5 set + ); + let root_trie_node = TrieNode::Branch(root_branch_node); + + // 3. Initialize the sparse trie using from_root + // This will internally create Hash nodes for paths "1" and "5" initially. + let mut sparse = RevealedSparseTrie::from_root(root_trie_node, TrieMasks::none(), false) + .expect("Failed to create trie from root"); + + // Assertions before we reveal child5 + assert_matches!(sparse.nodes.get(&Nibbles::default()), Some(SparseNode::Branch { state_mask, .. }) if *state_mask == TrieMask::new(0b100010)); // Here we check that 1 and 5 are set in the state_mask + assert_matches!(sparse.nodes.get(&path_to_blind), Some(SparseNode::Hash(h)) if *h == blinded_hash ); + assert!(sparse.nodes.get(&revealed_leaf_prefix).unwrap().is_hash()); // Child 5 is initially a hash of its RLP + assert!(sparse.values.is_empty()); + + // 4. Explicitly reveal the leaf node for child 5 + sparse + .reveal_node( + revealed_leaf_prefix.clone(), + TrieNode::Leaf(leaf_node_child5), + TrieMasks::none(), + ) + .expect("Failed to reveal leaf node"); + + // Assertions after we reveal child 5 + assert_matches!(sparse.nodes.get(&Nibbles::default()), Some(SparseNode::Branch { state_mask, .. }) if *state_mask == TrieMask::new(0b100010)); + assert_matches!(sparse.nodes.get(&path_to_blind), Some(SparseNode::Hash(h)) if *h == blinded_hash ); + assert_matches!(sparse.nodes.get(&revealed_leaf_prefix), Some(SparseNode::Leaf { key, .. }) if *key == revealed_leaf_suffix); + assert_eq!(sparse.values.get(&revealed_leaf_full_path), Some(&revealed_value)); + + let result = sparse.find_leaf(&search_path, None); + + // 5. Assert the expected error + // Should error because it hit the blinded node at path "1" only node at "5" was revealed + assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash }) + if path == path_to_blind && hash == blinded_hash + ); + } +} + #[cfg(test)] mod tests { use super::*;