diff --git a/mpt_trie/src/partial_trie.rs b/mpt_trie/src/partial_trie.rs index 27027eeae..e8f1ebea7 100644 --- a/mpt_trie/src/partial_trie.rs +++ b/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; + /// Creates a new partial trie from a node with a provided collapse + /// strategy. + fn new_with_strategy(n: Node, strategy: OnOrphanedHashNode) -> Self; + /// Inserts a node into the trie. fn insert(&mut self, k: K, v: V) -> TrieOpResult<()> where @@ -211,6 +215,10 @@ impl PartialTrie for StandardTrie { Self(n) } + fn new_with_strategy(n: Node, _strategy: OnOrphanedHashNode) -> Self { + Self(n) + } + fn insert(&mut self, k: K, v: V) -> TrieOpResult<()> where K: Into, @@ -240,7 +248,7 @@ impl PartialTrie for StandardTrie { where K: Into, { - 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, pub(crate) hash: Arc>>, + + 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, strategy: OnOrphanedHashNode) -> Self { + Self { + node, + hash: Arc::new(RwLock::new(None)), + strategy, } } @@ -364,7 +398,7 @@ impl PartialTrie for HashedPartialTrie { where K: Into, { - let res = self.node.trie_delete(k); + let res = self.node.trie_delete(k, self.strategy); self.set_hash(None); res diff --git a/mpt_trie/src/trie_ops.rs b/mpt_trie/src/trie_ops.rs index 1d6c0ab4d..c7cee9b09 100644 --- a/mpt_trie/src/trie_ops.rs +++ b/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 Node { /// 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>> + pub(crate) fn trie_delete( + &mut self, + k: K, + strategy: OnOrphanedHashNode, + ) -> TrieOpResult>> where K: Into, { 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 = &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 = &wrapped_node; + *self = node_ref.clone(); + + Ok(Some(deleted_val)) + }, + ) } pub(crate) fn trie_items(&self) -> impl Iterator { @@ -516,6 +524,7 @@ fn insert_into_trie_rec( fn delete_intern( node: &Node, mut curr_k: Nibbles, + strategy: OnOrphanedHashNode, ) -> TrieOpResult, Vec)>> { match node { Node::Empty => { @@ -532,7 +541,7 @@ fn delete_intern( 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( 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( .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( fn try_collapse_if_extension( node: WrappedNode, curr_key: &Nibbles, + strategy: OnOrphanedHashNode, ) -> 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, curr_key, strategy) + } _ => Ok(node), } } @@ -637,6 +653,7 @@ fn collapse_ext_node_if_needed( ext_nibbles: &Nibbles, child: &WrappedNode, curr_key: &Nibbles, + strategy: OnOrphanedHashNode, ) -> TrieOpResult> { trace!( "Collapsing extension node ({:x}) with child {}...", @@ -657,10 +674,13 @@ 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, - )), + 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))), } diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index 8ff70f340..468d4dc15 100644 --- a/trace_decoder/src/lib.rs +++ b/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::>()?, diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 305dffa9b..413163e34 100644 --- a/trace_decoder/src/type1.rs +++ b/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>>, @@ -22,6 +23,18 @@ pub struct Frontend { pub storage: BTreeMap, } +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) -> anyhow::Result { let executions = execute(instructions)?; ensure!( @@ -157,7 +170,7 @@ fn node2storagetrie(node: Node) -> anyhow::Result { Ok(()) } - let mut mpt = StorageTrie::default(); + let mut mpt = StorageTrie::new(OnOrphanedHashNode::CollapseToExtension); visit(&mut mpt, &stackstack::Stack::new(), node)?; Ok(mpt) } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 7671a06ea..ba0bb6259 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/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) -> Result>, Error> { let prev = self.untyped.get(key.into_nibbles()).map(Vec::from); self.untyped