Feat/237 mpt trie ext to branch collapse error

mpt_trie/src/nibbles.rs

@@ -254,7 +254,7 @@ impl From<U256> for NibblesIntern {
     }
 }
 
-#[derive(Copy, Clone, Deserialize, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+#[derive(Copy, Clone, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
 /// A sequence of nibbles which is used as the key type into
 /// [`PartialTrie`][`crate::partial_trie::PartialTrie`].
 ///
@@ -323,6 +323,14 @@ impl Debug for Nibbles {
     }
 }
 
+/// While we could just derive `Default` and it would be correct, it's a bit
+/// cleaner to instead call [`Nibbles::new`] explicitly.
+impl Default for Nibbles {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl FromStr for Nibbles {
     type Err = StrToNibblesError;
 
@@ -355,6 +363,17 @@ impl UpperHex for Nibbles {
 }
 
 impl Nibbles {
+    /// Create `Nibbles` that is empty.
+    ///
+    /// Note that mean that the key size is `0` and does not mean that the key
+    /// contains the `0` [`Nibble`].
+    pub fn new() -> Self {
+        Self {
+            count: 0,
+            packed: NibblesIntern::default(),
+        }
+    }
+
     /// Creates `Nibbles` from big endian bytes.
     ///
     /// Returns an error if the byte slice is empty or is longer than `32`

mpt_trie/src/trie_ops.rs

@@ -31,11 +31,24 @@ 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 encontered an unexisting type of node during
-    /// an extension node collapse.
-    #[error("Extension managed to get an unexisting child node type! (child: {0})")]
+    /// 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})")]
     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),
@@ -360,6 +373,11 @@ impl<T: PartialTrie> Node<T> {
         }
     }
 
+    /// 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<K>(&mut self, k: K) -> TrieOpResult<Option<Vec<u8>>>
     where
         K: Into<Nibbles>,
@@ -368,8 +386,10 @@ impl<T: PartialTrie> Node<T> {
         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
-            let wrapped_node = try_collapse_if_extension(updated_root)?;
+            // 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::default())?;
             let node_ref: &Node<T> = &wrapped_node;
             *self = node_ref.clone();
 
@@ -521,12 +541,15 @@ fn delete_intern<N: PartialTrie>(
                     {
                         false => {
                             // Branch stays.
+
                             let mut updated_children = children.clone();
                             updated_children[nibble as usize] =
-                                try_collapse_if_extension(updated_child)?;
+                                try_collapse_if_extension(updated_child, &curr_k)?;
                             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,
@@ -559,7 +582,7 @@ fn delete_intern<N: PartialTrie>(
                     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)?;
+                                collapse_ext_node_if_needed(ext_nibbles, &updated_child, &curr_k)?;
                             Ok(Some((updated_node, value_deleted)))
                         })
                     })
@@ -576,16 +599,44 @@ fn delete_intern<N: PartialTrie>(
     }
 }
 
-fn try_collapse_if_extension<N: PartialTrie>(node: WrappedNode<N>) -> TrieOpResult<WrappedNode<N>> {
+fn try_collapse_if_extension<N: PartialTrie>(
+    node: WrappedNode<N>,
+    curr_key: &Nibbles,
+) -> TrieOpResult<WrappedNode<N>> {
     match node.as_ref() {
-        Node::Extension { nibbles, child } => collapse_ext_node_if_needed(nibbles, child),
+        Node::Extension { nibbles, child } => collapse_ext_node_if_needed(nibbles, child, curr_key),
         _ => 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<N: PartialTrie>(
     ext_nibbles: &Nibbles,
     child: &WrappedNode<N>,
+    curr_key: &Nibbles,
 ) -> TrieOpResult<WrappedNode<N>> {
     trace!(
         "Collapsing extension node ({:x}) with child {}...",
@@ -606,7 +657,11 @@ fn collapse_ext_node_if_needed<N: PartialTrie>(
             nibbles: leaf_nibbles,
             value,
         } => Ok(leaf(ext_nibbles.merge_nibbles(leaf_nibbles), value.clone())),
-        Node::Hash(_) => Ok(extension(*ext_nibbles, child.clone())),
+        Node::Hash(h) => Err(TrieOpError::ExtensionCollapsedIntoHashError(
+            curr_key.merge_nibbles(ext_nibbles),
+            *h,
+        )),
+        // Can never do this safely, so return an error.
         _ => Err(TrieOpError::HashNodeExtError(TrieNodeType::from(child))),
     }
 }