Add collapse strategy to PartialTrie variants (#501)

* Add collapse strategy to PartialTrie variants

* Fix typo

Co-authored-by: BGluth <[email protected]>

* Change terminology

* Rename

* Nit

---------

Co-authored-by: BGluth <[email protected]>

Author: Nashtare

mpt_trie/src/partial_trie.rs

@@ -49,6 +49,10 @@ pub trait PartialTrie:
     /// Creates a new partial trie from a node.
     fn new(n: Node<Self>) -> Self;
 
+    /// Creates a new partial trie from a node with a provided collapse
+    /// strategy.
+    fn new_with_strategy(n: Node<Self>, strategy: OnOrphanedHashNode) -> Self;
+
     /// Inserts a node into the trie.
     fn insert<K, V>(&mut self, k: K, v: V) -> TrieOpResult<()>
     where
@@ -211,6 +215,10 @@ impl PartialTrie for StandardTrie {
         Self(n)
     }
 
+    fn new_with_strategy(n: Node<Self>, _strategy: OnOrphanedHashNode) -> Self {
+        Self(n)
+    }
+
     fn insert<K, V>(&mut self, k: K, v: V) -> TrieOpResult<()>
     where
         K: Into<Nibbles>,
@@ -240,7 +248,7 @@ impl PartialTrie for StandardTrie {
     where
         K: Into<Nibbles>,
     {
-        self.0.trie_delete(k)
+        self.0.trie_delete(k, OnOrphanedHashNode::Reject)
     }
 
     fn hash(&self) -> H256 {
@@ -304,6 +312,23 @@ where
 pub struct HashedPartialTrie {
     pub(crate) node: Node<HashedPartialTrie>,
     pub(crate) hash: Arc<RwLock<Option<H256>>>,
+
+    pub(crate) strategy: OnOrphanedHashNode,
+}
+
+/// How to handle the following subtree on deletion of the indicated node.
+/// ```text
+///      BranchNode
+///        /    \
+/// DeleteMe    OrphanedHashNode
+/// ```
+#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
+pub enum OnOrphanedHashNode {
+    /// Replace `BranchNode` with an appropriate `ExtensionNode`
+    CollapseToExtension,
+    /// Return an error.
+    #[default]
+    Reject,
 }
 
 impl_from_for_trie_type!(HashedPartialTrie);
@@ -329,6 +354,15 @@ impl PartialTrie for HashedPartialTrie {
         Self {
             node,
             hash: Arc::new(RwLock::new(None)),
+            strategy: OnOrphanedHashNode::default(),
+        }
+    }
+
+    fn new_with_strategy(node: Node<Self>, strategy: OnOrphanedHashNode) -> Self {
+        Self {
+            node,
+            hash: Arc::new(RwLock::new(None)),
+            strategy,
         }
     }
 
@@ -364,7 +398,7 @@ impl PartialTrie for HashedPartialTrie {
     where
         K: Into<crate::nibbles::Nibbles>,
     {
-        let res = self.node.trie_delete(k);
+        let res = self.node.trie_delete(k, self.strategy);
         self.set_hash(None);
 
         res

mpt_trie/src/trie_ops.rs

@@ -10,7 +10,7 @@ use thiserror::Error;
 
 use crate::{
     nibbles::{Nibble, Nibbles},
-    partial_trie::{Node, PartialTrie, WrappedNode},
+    partial_trie::{Node, OnOrphanedHashNode, PartialTrie, WrappedNode},
     utils::TrieNodeType,
 };
 
@@ -378,23 +378,31 @@ impl<T: PartialTrie> Node<T> {
     /// 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>>>
+    pub(crate) fn trie_delete<K>(
+        &mut self,
+        k: K,
+        strategy: OnOrphanedHashNode,
+    ) -> TrieOpResult<Option<Vec<u8>>>
     where
         K: Into<Nibbles>,
     {
         let k: Nibbles = k.into();
         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::default())?;
-            let node_ref: &Node<T> = &wrapped_node;
-            *self = node_ref.clone();
-
-            Ok(Some(deleted_val))
-        })
+        delete_intern(&self.clone(), k, strategy)?.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::default(), strategy)?;
+                let node_ref: &Node<T> = &wrapped_node;
+                *self = node_ref.clone();
+
+                Ok(Some(deleted_val))
+            },
+        )
     }
 
     pub(crate) fn trie_items(&self) -> impl Iterator<Item = (Nibbles, ValOrHash)> {
@@ -516,6 +524,7 @@ fn insert_into_trie_rec<N: PartialTrie>(
 fn delete_intern<N: PartialTrie>(
     node: &Node<N>,
     mut curr_k: Nibbles,
+    strategy: OnOrphanedHashNode,
 ) -> TrieOpResult<Option<(WrappedNode<N>, Vec<u8>)>> {
     match node {
         Node::Empty => {
@@ -532,7 +541,7 @@ fn delete_intern<N: PartialTrie>(
             let nibble = curr_k.pop_next_nibble_front();
             trace!("Delete traversed Branch nibble {:x}", nibble);
 
-            delete_intern(&children[nibble as usize], curr_k)?.map_or(Ok(None),
+            delete_intern(&children[nibble as usize], curr_k, strategy)?.map_or(Ok(None),
                 |(updated_child, value_deleted)| {
                     // If the child we recursively called is deleted, then we may need to reduce
                     // this branch to an extension/leaf.
@@ -544,7 +553,7 @@ fn delete_intern<N: PartialTrie>(
 
                             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, &curr_k, strategy)?;
                             branch(updated_children, value.clone())
                         }
                         true => {
@@ -579,10 +588,14 @@ fn delete_intern<N: PartialTrie>(
                 .then(|| {
                     curr_k.truncate_n_nibbles_front_mut(ext_nibbles.count);
 
-                    delete_intern(child, curr_k).and_then(|res| {
+                    delete_intern(child, curr_k, strategy).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)?;
+                            let updated_node = collapse_ext_node_if_needed(
+                                ext_nibbles,
+                                &updated_child,
+                                &curr_k,
+                                strategy,
+                            )?;
                             Ok(Some((updated_node, value_deleted)))
                         })
                     })
@@ -602,9 +615,12 @@ fn delete_intern<N: PartialTrie>(
 fn try_collapse_if_extension<N: PartialTrie>(
     node: WrappedNode<N>,
     curr_key: &Nibbles,
+    strategy: OnOrphanedHashNode,
 ) -> TrieOpResult<WrappedNode<N>> {
     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, curr_key, strategy)
+        }
         _ => Ok(node),
     }
 }
@@ -637,6 +653,7 @@ fn collapse_ext_node_if_needed<N: PartialTrie>(
     ext_nibbles: &Nibbles,
     child: &WrappedNode<N>,
     curr_key: &Nibbles,
+    strategy: OnOrphanedHashNode,
 ) -> TrieOpResult<WrappedNode<N>> {
     trace!(
         "Collapsing extension node ({:x}) with child {}...",
@@ -657,10 +674,13 @@ fn collapse_ext_node_if_needed<N: PartialTrie>(
             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,
-        )),
+        Node::Hash(h) => match strategy {
+            OnOrphanedHashNode::CollapseToExtension => Ok(extension(*ext_nibbles, child.clone())),
+            OnOrphanedHashNode::Reject => 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))),
     }

trace_decoder/src/lib.rs

@@ -103,7 +103,7 @@ use evm_arithmetization::proof::{BlockHashes, BlockMetadata};
 use evm_arithmetization::GenerationInputs;
 use keccak_hash::keccak as hash;
 use keccak_hash::H256;
-use mpt_trie::partial_trie::HashedPartialTrie;
+use mpt_trie::partial_trie::{HashedPartialTrie, OnOrphanedHashNode};
 use processed_block_trace::ProcessedTxnInfo;
 use serde::{Deserialize, Serialize};
 use typed_mpt::{StateTrie, StorageTrie, TrieKey};
@@ -320,7 +320,7 @@ pub fn entrypoint(
         }) => ProcessedBlockTracePreImages {
             tries: PartialTriePreImages {
                 state: state.items().try_fold(
-                    StateTrie::default(),
+                    StateTrie::new(OnOrphanedHashNode::Reject),
                     |mut acc, (nibbles, hash_or_val)| {
                         let path = TrieKey::from_nibbles(nibbles);
                         match hash_or_val {
@@ -342,18 +342,21 @@ pub fn entrypoint(
                     .into_iter()
                     .map(|(k, SeparateTriePreImage::Direct(v))| {
                         v.items()
-                            .try_fold(StorageTrie::default(), |mut acc, (nibbles, hash_or_val)| {
-                                let path = TrieKey::from_nibbles(nibbles);
-                                match hash_or_val {
-                                    mpt_trie::trie_ops::ValOrHash::Val(value) => {
-                                        acc.insert(path, value)?;
-                                    }
-                                    mpt_trie::trie_ops::ValOrHash::Hash(h) => {
-                                        acc.insert_hash(path, h)?;
-                                    }
-                                };
-                                anyhow::Ok(acc)
-                            })
+                            .try_fold(
+                                StorageTrie::new(OnOrphanedHashNode::Reject),
+                                |mut acc, (nibbles, hash_or_val)| {
+                                    let path = TrieKey::from_nibbles(nibbles);
+                                    match hash_or_val {
+                                        mpt_trie::trie_ops::ValOrHash::Val(value) => {
+                                            acc.insert(path, value)?;
+                                        }
+                                        mpt_trie::trie_ops::ValOrHash::Hash(h) => {
+                                            acc.insert_hash(path, h)?;
+                                        }
+                                    };
+                                    anyhow::Ok(acc)
+                                },
+                            )
                             .map(|v| (k, v))
                     })
                     .collect::<Result<_, _>>()?,

trace_decoder/src/type1.rs

@@ -7,13 +7,14 @@ use std::collections::{BTreeMap, BTreeSet};
 use anyhow::{bail, ensure, Context as _};
 use either::Either;
 use evm_arithmetization::generation::mpt::AccountRlp;
+use mpt_trie::partial_trie::OnOrphanedHashNode;
 use nunny::NonEmpty;
 use u4::U4;
 
 use crate::typed_mpt::{StateTrie, StorageTrie, TrieKey};
 use crate::wire::{Instruction, SmtLeaf};
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
 pub struct Frontend {
     pub state: StateTrie,
     pub code: BTreeSet<NonEmpty<Vec<u8>>>,
@@ -22,6 +23,18 @@ pub struct Frontend {
     pub storage: BTreeMap<TrieKey, StorageTrie>,
 }
 
+impl Default for Frontend {
+    // This frontend is intended to be used with our custom `zeroTracer`,
+    // which covers branch-to-extension collapse edge cases.
+    fn default() -> Self {
+        Self {
+            state: StateTrie::new(OnOrphanedHashNode::CollapseToExtension),
+            code: BTreeSet::new(),
+            storage: BTreeMap::new(),
+        }
+    }
+}
+
 pub fn frontend(instructions: impl IntoIterator<Item = Instruction>) -> anyhow::Result<Frontend> {
     let executions = execute(instructions)?;
     ensure!(
@@ -157,7 +170,7 @@ fn node2storagetrie(node: Node) -> anyhow::Result<StorageTrie> {
         Ok(())
     }
 
-    let mut mpt = StorageTrie::default();
+    let mut mpt = StorageTrie::new(OnOrphanedHashNode::CollapseToExtension);
     visit(&mut mpt, &stackstack::Stack::new(), node)?;
     Ok(mpt)
 }

trace_decoder/src/typed_mpt.rs

@@ -7,7 +7,7 @@ use copyvec::CopyVec;
 use ethereum_types::{Address, H256};
 use evm_arithmetization::generation::mpt::AccountRlp;
 use mpt_trie::{
-    partial_trie::{HashedPartialTrie, Node, PartialTrie as _},
+    partial_trie::{HashedPartialTrie, Node, OnOrphanedHashNode, PartialTrie as _},
     trie_ops::TrieOpError,
 };
 use u4::{AsNibbles, U4};
@@ -243,6 +243,14 @@ pub struct StateTrie {
 }
 
 impl StateTrie {
+    pub fn new(strategy: OnOrphanedHashNode) -> Self {
+        Self {
+            typed: TypedMpt {
+                inner: HashedPartialTrie::new_with_strategy(Node::Empty, strategy),
+                _ty: PhantomData,
+            },
+        }
+    }
     pub fn insert_by_address(
         &mut self,
         address: Address,
@@ -319,6 +327,11 @@ pub struct StorageTrie {
     untyped: HashedPartialTrie,
 }
 impl StorageTrie {
+    pub fn new(strategy: OnOrphanedHashNode) -> Self {
+        Self {
+            untyped: HashedPartialTrie::new_with_strategy(Node::Empty, strategy),
+        }
+    }
     pub fn insert(&mut self, key: TrieKey, value: Vec<u8>) -> Result<Option<Vec<u8>>, Error> {
         let prev = self.untyped.get(key.into_nibbles()).map(Vec::from);
         self.untyped