From 9b3e1b5573cb01375f191602869ad39e1a7328d1 Mon Sep 17 00:00:00 2001 From: BGluth Date: Fri, 9 Aug 2024 11:19:20 -0600 Subject: [PATCH] Requested PR changes for #455 --- mpt_trie/src/debug_tools/diff.rs | 8 +-- mpt_trie/src/debug_tools/query.rs | 2 +- mpt_trie/src/nibbles.rs | 18 ++--- mpt_trie/src/special_query.rs | 4 +- mpt_trie/src/trie_ops.rs | 114 ++++-------------------------- mpt_trie/src/trie_subsets.rs | 13 ++-- mpt_trie/src/utils.rs | 2 +- 7 files changed, 39 insertions(+), 122 deletions(-) diff --git a/mpt_trie/src/debug_tools/diff.rs b/mpt_trie/src/debug_tools/diff.rs index b2ba6dd8e..c6a8458db 100644 --- a/mpt_trie/src/debug_tools/diff.rs +++ b/mpt_trie/src/debug_tools/diff.rs @@ -41,7 +41,7 @@ use crate::{ /// [branch][`Node::Branch`]s have no [`Nibble`] directly associated with them. fn get_key_piece_from_node(n: &Node) -> Nibbles { match n { - Node::Empty | Node::Hash(_) | Node::Branch { .. } => Nibbles::empty(), + Node::Empty | Node::Hash(_) | Node::Branch { .. } => Nibbles::default(), Node::Extension { nibbles, child: _ } | Node::Leaf { nibbles, value: _ } => *nibbles, } } @@ -187,7 +187,7 @@ fn find_latest_diff_point_between_tries( a: &HashedPartialTrie, b: &HashedPartialTrie, ) -> Option { - let state = DepthDiffPerCallState::new(a, b, Nibbles::empty(), 0); + let state = DepthDiffPerCallState::new(a, b, Nibbles::default(), 0); let mut longest_state = DepthNodeDiffState::default(); find_latest_diff_point_between_tries_rec(&state, &mut longest_state); @@ -327,7 +327,7 @@ fn find_latest_diff_point_between_tries_rec( create_diff_detection_state_based_from_hashes( a_hash, b_hash, - &state.new_from_parent(state.a, state.b, &Nibbles::empty()), + &state.new_from_parent(state.a, state.b, &Nibbles::default()), depth_state, ) } @@ -461,7 +461,7 @@ mod tests { let expected = DiffPoint { depth: 0, path: TriePath(vec![]), - key: Nibbles::empty(), + key: Nibbles::default(), a_info: expected_a, b_info: expected_b, }; diff --git a/mpt_trie/src/debug_tools/query.rs b/mpt_trie/src/debug_tools/query.rs index 402784523..753900262 100644 --- a/mpt_trie/src/debug_tools/query.rs +++ b/mpt_trie/src/debug_tools/query.rs @@ -23,7 +23,7 @@ fn get_key_piece_from_node_pulling_from_key_for_branches( curr_key: &Nibbles, ) -> Nibbles { match n { - Node::Empty | Node::Hash(_) => Nibbles::empty(), + Node::Empty | Node::Hash(_) => Nibbles::default(), Node::Branch { .. } => curr_key.get_next_nibbles(1), Node::Extension { nibbles, child: _ } | Node::Leaf { nibbles, value: _ } => *nibbles, } diff --git a/mpt_trie/src/nibbles.rs b/mpt_trie/src/nibbles.rs index bbb08bb22..d267b9082 100644 --- a/mpt_trie/src/nibbles.rs +++ b/mpt_trie/src/nibbles.rs @@ -324,10 +324,10 @@ impl Debug for Nibbles { } /// While we could just derive `Default` and it would be correct, it's a bit -/// cleaner to instead call [`Nibbles::empty`] explicitly. +/// cleaner to instead call [`Nibbles::defaultaultaultault`] explicitly. impl Default for Nibbles { fn default() -> Self { - Self::empty() + Self::new() } } @@ -367,7 +367,7 @@ impl Nibbles { /// /// Note that mean that the key size is `0` and does not mean that the key /// contains the `0` [`Nibble`]. - pub fn empty() -> Self { + pub fn new() -> Self { Self { count: 0, packed: NibblesIntern::default(), @@ -1145,7 +1145,7 @@ mod tests { #[test] fn push_nibble_front_works() { - test_and_assert_nib_push_func(Nibbles::empty(), 0x1, |n| n.push_nibble_front(0x1)); + test_and_assert_nib_push_func(Nibbles::default(), 0x1, |n| n.push_nibble_front(0x1)); test_and_assert_nib_push_func(0x1, 0x21, |n| n.push_nibble_front(0x2)); test_and_assert_nib_push_func( Nibbles::from_str(ZERO_NIBS_63).unwrap(), @@ -1156,7 +1156,7 @@ mod tests { #[test] fn push_nibble_back_works() { - test_and_assert_nib_push_func(Nibbles::empty(), 0x1, |n| n.push_nibble_back(0x1)); + test_and_assert_nib_push_func(Nibbles::default(), 0x1, |n| n.push_nibble_back(0x1)); test_and_assert_nib_push_func(0x1, 0x12, |n| n.push_nibble_back(0x2)); test_and_assert_nib_push_func( Nibbles::from_str(ZERO_NIBS_63).unwrap(), @@ -1167,7 +1167,7 @@ mod tests { #[test] fn push_nibbles_front_works() { - test_and_assert_nib_push_func(Nibbles::empty(), 0x1234, |n| { + test_and_assert_nib_push_func(Nibbles::default(), 0x1234, |n| { n.push_nibbles_front(&0x1234.into()) }); test_and_assert_nib_push_func(0x1234, 0x5671234, |n| n.push_nibbles_front(&0x567.into())); @@ -1180,7 +1180,7 @@ mod tests { #[test] fn push_nibbles_back_works() { - test_and_assert_nib_push_func(Nibbles::empty(), 0x1234, |n| { + test_and_assert_nib_push_func(Nibbles::default(), 0x1234, |n| { n.push_nibbles_back(&0x1234.into()) }); test_and_assert_nib_push_func(0x1234, 0x1234567, |n| n.push_nibbles_back(&0x567.into())); @@ -1206,13 +1206,13 @@ mod tests { fn get_next_nibbles_works() -> Result<(), StrToNibblesError> { let n: Nibbles = 0x1234.into(); - assert_eq!(n.get_next_nibbles(0), Nibbles::empty()); + assert_eq!(n.get_next_nibbles(0), Nibbles::default()); assert_eq!(n.get_next_nibbles(1), Nibbles::from(0x1)); assert_eq!(n.get_next_nibbles(2), Nibbles::from(0x12)); assert_eq!(n.get_next_nibbles(3), Nibbles::from(0x123)); assert_eq!(n.get_next_nibbles(4), Nibbles::from(0x1234)); - assert_eq!(Nibbles::from(0x0).get_next_nibbles(0), Nibbles::empty()); + assert_eq!(Nibbles::from(0x0).get_next_nibbles(0), Nibbles::default()); let n = Nibbles::from_str( "0x3ab76c381c0f8ea617ea96780ffd1e165c754b28a41a95922f9f70682c581353", diff --git a/mpt_trie/src/special_query.rs b/mpt_trie/src/special_query.rs index c0dbb4453..c133a6091 100644 --- a/mpt_trie/src/special_query.rs +++ b/mpt_trie/src/special_query.rs @@ -157,13 +157,13 @@ mod test { TrieSegment::Branch(2), TrieSegment::Extension(Nibbles::from_str("0x00").unwrap()), TrieSegment::Branch(0x1), - TrieSegment::Leaf(Nibbles::empty()), + TrieSegment::Leaf(Nibbles::default()), ], vec![ TrieSegment::Branch(2), TrieSegment::Extension(Nibbles::from_str("0x00").unwrap()), TrieSegment::Branch(0x2), - TrieSegment::Leaf(Nibbles::empty()), + TrieSegment::Leaf(Nibbles::default()), ], ]; diff --git a/mpt_trie/src/trie_ops.rs b/mpt_trie/src/trie_ops.rs index 918cfa319..5f94fc942 100644 --- a/mpt_trie/src/trie_ops.rs +++ b/mpt_trie/src/trie_ops.rs @@ -31,24 +31,11 @@ pub enum TrieOpError { #[error("Attempted to delete a value that ended up inside a hash node! (hash: {0})")] HashNodeDeleteError(H256), - /// An error that occurs when we encounter an non-existing type of node - /// during an extension node collapse. - #[error("Extension managed to get an non-existing child node type! (child: {0})")] + /// An error that occurs when encontered an unexisting type of node during + /// an extension node collapse. + #[error("Extension managed to get an unexisting child node type! (child: {0})")] HashNodeExtError(TrieNodeType), - /// An error that occurs when we attempted to collapse an extension node - /// into a hash node. - /// - /// If this occurs, then there is a chance that we can not collapse - /// correctly and will produce the incorrect trie (and also the incorrect - /// trie hash). If the hash node is a hash of a leaf, then we need to - /// collapse the extension key into the leaf. However, this information - /// is lost if the node is hashed, and we can not tell if the hash node - /// was made from a leaf. As such, it's the responsibility of whoever is - /// constructing & mutating the trie that this will never occur. - #[error("Attempted to collapse an extension node into a hash node! This is unsafe! (See https://github.com/0xPolygonZero/zk_evm/issues/237 for more info) (Extension key: {0:x}, child hash node: {1:x})")] - ExtensionCollapsedIntoHashError(Nibbles, H256), - /// Failed to insert a hash node into the trie. #[error("Attempted to place a hash node on an existing node! (hash: {0})")] ExistingHashNodeError(H256), @@ -253,7 +240,7 @@ impl Iterator for PartialTrieIter { next_iter_item = match stack_entry { IterStackEntry::Root(root) => { - self.advance_iter_to_next_empty_leaf_or_hash_node(&root, Nibbles::empty()) + self.advance_iter_to_next_empty_leaf_or_hash_node(&root, Nibbles::default()) } IterStackEntry::Extension(num_nibbles) => { // Drop nibbles that extension added since we are going back up the trie. @@ -373,11 +360,6 @@ impl Node { } } - /// Deletes a key if it exists in the trie. - /// - /// If the key exists, then the existing node value that was deleted is - /// returned. Otherwise, if the key is not present, then `None` is returned - /// instead. pub(crate) fn trie_delete(&mut self, k: K) -> TrieOpResult>> where K: Into, @@ -386,10 +368,8 @@ impl Node { trace!("Deleting a leaf node with key {} if it exists", k); delete_intern(&self.clone(), k)?.map_or(Ok(None), |(updated_root, deleted_val)| { - // Final check at the root if we have an extension node. While this check also - // exists as we recursively traverse down the trie, it can not perform this - // check on the root node. - let wrapped_node = try_collapse_if_extension(updated_root, &Nibbles::empty())?; + // Final check at the root if we have an extension node + let wrapped_node = try_collapse_if_extension(updated_root)?; let node_ref: &Node = &wrapped_node; *self = node_ref.clone(); @@ -399,7 +379,7 @@ impl Node { pub(crate) fn trie_items(&self) -> impl Iterator { PartialTrieIter { - curr_key_after_last_branch: Nibbles::empty(), + curr_key_after_last_branch: Nibbles::default(), trie_stack: vec![IterStackEntry::Root(self.clone().into())], } } @@ -541,15 +521,12 @@ fn delete_intern( { false => { // Branch stays. - let mut updated_children = children.clone(); updated_children[nibble as usize] = - try_collapse_if_extension(updated_child, &curr_k)?; + try_collapse_if_extension(updated_child)?; branch(updated_children, value.clone()) } true => { - // We need to collapse the branch into an extension/leaf node. - let (child_nibble, non_empty_node) = get_other_non_empty_child_and_nibble_in_two_elem_branch( children, nibble, @@ -582,7 +559,7 @@ fn delete_intern( delete_intern(child, curr_k).and_then(|res| { res.map_or(Ok(None), |(updated_child, value_deleted)| { let updated_node = - collapse_ext_node_if_needed(ext_nibbles, &updated_child, &curr_k)?; + collapse_ext_node_if_needed(ext_nibbles, &updated_child)?; Ok(Some((updated_node, value_deleted))) }) }) @@ -599,44 +576,16 @@ fn delete_intern( } } -fn try_collapse_if_extension( - node: WrappedNode, - curr_key: &Nibbles, -) -> TrieOpResult> { +fn try_collapse_if_extension(node: WrappedNode) -> TrieOpResult> { match node.as_ref() { - Node::Extension { nibbles, child } => collapse_ext_node_if_needed(nibbles, child, curr_key), + Node::Extension { nibbles, child } => collapse_ext_node_if_needed(nibbles, child), _ => Ok(node), } } -/// Attempt to collapse an extension node if we are required. -/// -/// The scenarios where we are required to do so are where the extension node is -/// pointing to a: -/// - Extension (the parent extension absorbs the child's key). -/// - Leaf (the leaf absorbs the extension's key). -/// -/// While an extension does not collapse when its child is a branch, we need to -/// still check that a specific edge case does not exist for the branch node. -/// Specifically, if all of the following holds true for the branch where: -/// - It has exactly two children. -/// - One child that is a hash node. -/// - One child that is a leaf node. -/// - The leaf child ends up getting deleted. -/// -/// Then we need to return an error, because if the hash node was created from a -/// leaf node, we are required to collapse the extension key into the leaf node. -/// However, since it's a hash node, we: -/// - Have no idea what the original node was. -/// - Can not access the underlying key if we needed to collapse it. -/// -/// Because of this, we need to rely on the user to not allow `mpt_trie` to -/// arrive at this state, as we can not ensure that we will be able to produce -/// the correct trie. fn collapse_ext_node_if_needed( ext_nibbles: &Nibbles, child: &WrappedNode, - curr_key: &Nibbles, ) -> TrieOpResult> { trace!( "Collapsing extension node ({:x}) with child {}...", @@ -657,11 +606,7 @@ fn collapse_ext_node_if_needed( nibbles: leaf_nibbles, value, } => Ok(leaf(ext_nibbles.merge_nibbles(leaf_nibbles), value.clone())), - Node::Hash(h) => Err(TrieOpError::ExtensionCollapsedIntoHashError( - curr_key.merge_nibbles(ext_nibbles), - *h, - )), - // Can never do this safely, so return an error. + Node::Hash(_) => Ok(extension(*ext_nibbles, child.clone())), _ => Err(TrieOpError::HashNodeExtError(TrieNodeType::from(child))), } } @@ -871,7 +816,6 @@ fn create_node_if_ins_val_not_hash) -> WrappedNode>( mod tests { use std::{collections::HashSet, iter::once}; - use keccak_hash::H256; use log::debug; use super::ValOrHash; @@ -885,7 +829,7 @@ mod tests { generate_n_random_variable_trie_value_entries, get_non_hash_values_in_trie, unwrap_iter_item_to_val, TestInsertValEntry, }, - trie_ops::{TrieOpError, TrieOpResult}, + trie_ops::TrieOpResult, utils::{create_mask_of_1s, TryFromIterator}, }; @@ -1202,38 +1146,6 @@ mod tests { Ok(()) } - #[test] - fn collapsing_an_extension_to_a_hash_node_returns_error() -> TrieOpResult<()> { - common_setup(); - - // We want to create a trie like this: - // ``` - // B - // / \ - // L H - // ``` - // If we delete the Leaf, then the Branch will collapse into an Extension and - // then the extension will attempt to collapse into it's child (the Hash node). - // An extension collapsing into a hash node is an error, so we should expect - // that an error is returned. - - // Branch at `0x1` - // Leaf at `0x12` - // Hash at `0x13` - let mut trie = HashedPartialTrie::default(); - - trie.insert(0x12, vec![10])?; - trie.insert(0x13, H256::zero())?; - - // Delete the leaf to trigger the collapse. - let res = trie.delete(0x12); - if let Err(TrieOpError::ExtensionCollapsedIntoHashError(_, _)) = res { - return Ok(()); - } - - panic!("Expected an extension collapsed into a hash error!"); - } - #[test] fn deletion_massive_trie() -> TrieOpResult<()> { common_setup(); diff --git a/mpt_trie/src/trie_subsets.rs b/mpt_trie/src/trie_subsets.rs index b8e9cbf23..74f50d769 100644 --- a/mpt_trie/src/trie_subsets.rs +++ b/mpt_trie/src/trie_subsets.rs @@ -439,7 +439,12 @@ mod tests { return_on_empty_or_hash: bool, ) -> Vec { let mut nodes = Vec::new(); - get_nodes_in_trie_intern_rec(trie, Nibbles::empty(), &mut nodes, return_on_empty_or_hash); + get_nodes_in_trie_intern_rec( + trie, + Nibbles::default(), + &mut nodes, + return_on_empty_or_hash, + ); nodes } @@ -561,7 +566,7 @@ mod tests { assert_node_exists( &all_non_empty_and_hash_nodes, TrieNodeType::Branch, - Nibbles::empty(), + Nibbles::default(), ); assert_node_exists(&all_non_empty_and_hash_nodes, TrieNodeType::Branch, 0x1); assert_node_exists(&all_non_empty_and_hash_nodes, TrieNodeType::Leaf, 0x1234); @@ -588,7 +593,7 @@ mod tests { assert_node_exists( &all_non_empty_and_hash_nodes_partial, TrieNodeType::Branch, - Nibbles::empty(), + Nibbles::default(), ); assert_node_exists( &all_non_empty_and_hash_nodes_partial, @@ -613,7 +618,7 @@ mod tests { assert_node_exists( &all_non_empty_and_hash_nodes_partial, TrieNodeType::Branch, - Nibbles::empty(), + Nibbles::default(), ); assert_node_exists( &all_non_empty_and_hash_nodes_partial, diff --git a/mpt_trie/src/utils.rs b/mpt_trie/src/utils.rs index 531d96d0e..5d530878e 100644 --- a/mpt_trie/src/utils.rs +++ b/mpt_trie/src/utils.rs @@ -109,7 +109,7 @@ pub trait IntoTrieKey { impl, T: Iterator> IntoTrieKey for T { fn into_key(self) -> Nibbles { - let mut key = Nibbles::empty(); + let mut key = Nibbles::default(); for seg in self { match seg.borrow() {