fix: StorageTrie, and full frontend support

trace_decoder/src/decoding.rs

@@ -224,15 +224,25 @@ impl ProcessedBlockTrace {
         other_data: OtherBlockData,
     ) -> TraceParsingResult<Vec<GenerationInputs>> {
         let mut curr_block_tries = PartialTrieState {
-            state: self.tries.state.clone(),
-            storage: self.tries.storage.clone(),
+            state: self.tries.state.as_hashed_partial_trie(),
+            storage: self
+                .tries
+                .storage
+                .iter()
+                .map(|(k, v)| (*k, v.as_hashed_partial_trie()))
+                .collect(),
             ..Default::default()
         };
 
         // This is just a copy of `curr_block_tries`.
         let initial_tries_for_dummies = PartialTrieState {
-            state: self.tries.state,
-            storage: self.tries.storage,
+            state: self.tries.state.as_hashed_partial_trie(),
+            storage: self
+                .tries
+                .storage
+                .iter()
+                .map(|(k, v)| (*k, v.as_hashed_partial_trie()))
+                .collect(),
             ..Default::default()
         };
 

trace_decoder/src/lib.rs

@@ -103,6 +103,7 @@ use keccak_hash::keccak as hash;
 use keccak_hash::H256;
 use mpt_trie::partial_trie::HashedPartialTrie;
 use serde::{Deserialize, Serialize};
+use typed_mpt::{StateTrie, StorageTrie, TriePath};
 
 /// Core payload needed to generate proof for a block.
 /// Additional data retrievable from the blockchain node (using standard ETH RPC
@@ -297,7 +298,6 @@ pub fn entrypoint(
     resolve: impl Fn(H256) -> Vec<u8>,
 ) -> anyhow::Result<Vec<GenerationInputs>> {
     use anyhow::Context as _;
-    use evm_arithmetization::generation::mpt::AccountRlp;
     use mpt_trie::partial_trie::PartialTrie as _;
 
     use crate::processed_block_trace::{
@@ -321,11 +321,42 @@ pub fn entrypoint(
             storage: SeparateStorageTriesPreImage::MultipleTries(storage),
         }) => ProcessedBlockTracePreImages {
             tries: PartialTriePreImages {
-                state,
+                state: state.items().try_fold(
+                    StateTrie::default(),
+                    |mut acc, (nibbles, hash_or_val)| {
+                        let path = TriePath::from_nibbles(nibbles);
+                        match hash_or_val {
+                            mpt_trie::trie_ops::ValOrHash::Val(bytes) => acc.insert_by_path(
+                                path,
+                                rlp::decode(&bytes)
+                                    .context("invalid AccountRlp in direct state trie")?,
+                            ),
+                            mpt_trie::trie_ops::ValOrHash::Hash(h) => {
+                                acc.insert_branch_by_path(path, h)
+                            }
+                        };
+                        anyhow::Ok(acc)
+                    },
+                )?,
                 storage: storage
                     .into_iter()
-                    .map(|(k, SeparateTriePreImage::Direct(v))| (k, v))
-                    .collect(),
+                    .map(|(k, SeparateTriePreImage::Direct(v))| {
+                        v.items()
+                            .try_fold(StorageTrie::default(), |mut acc, (nibbles, hash_or_val)| {
+                                let path = TriePath::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_branch(path, h)
+                                    }
+                                };
+                                anyhow::Ok(acc)
+                            })
+                            .map(|v| (k, v))
+                    })
+                    .collect::<Result<_, _>>()?,
             },
             extra_code_hash_mappings: None,
         },
@@ -339,12 +370,10 @@ pub fn entrypoint(
             } = type1::frontend(instructions)?;
             ProcessedBlockTracePreImages {
                 tries: PartialTriePreImages {
-                    state: state.as_hashed_partial_trie(),
+                    state,
                     storage: storage
                         .into_iter()
-                        .map(|(path, trie)| {
-                            (path.into_hash_left_padded(), trie.as_hashed_partial_trie())
-                        })
+                        .map(|(path, trie)| (path.into_hash_left_padded(), trie))
                         .collect(),
                 },
                 extra_code_hash_mappings: match code.is_empty() {
@@ -362,11 +391,8 @@ pub fn entrypoint(
     let all_accounts_in_pre_images = pre_images
         .tries
         .state
-        .items()
-        .filter_map(|(addr, data)| {
-            data.as_val()
-                .map(|data| (addr.into(), rlp::decode::<AccountRlp>(data).unwrap()))
-        })
+        .iter()
+        .filter_map(|(addr, data)| Some((addr.into_hash_left_padded(), data.right()?)))
         .collect::<Vec<_>>();
 
     let code_db = {
@@ -419,8 +445,8 @@ pub fn entrypoint(
 
 #[derive(Debug, Default)]
 struct PartialTriePreImages {
-    pub state: HashedPartialTrie,
-    pub storage: HashMap<H256, HashedPartialTrie>,
+    pub state: StateTrie,
+    pub storage: HashMap<H256, StorageTrie>,
 }
 
 /// Like `#[serde(with = "hex")`, but tolerates and emits leading `0x` prefixes

trace_decoder/src/type1.rs

@@ -54,7 +54,7 @@ fn visit(
         Node::Hash(Hash { raw_hash }) => {
             let clobbered = frontend
                 .state
-                .insert_branch(TriePath::new(path.iter().copied())?, raw_hash.into());
+                .insert_branch_by_path(TriePath::new(path.iter().copied())?, raw_hash.into());
             ensure!(clobbered.is_none(), "duplicate hash")
         }
         Node::Leaf(Leaf { key, value }) => {

trace_decoder/src/typed_mpt.rs

@@ -50,6 +50,17 @@ impl TriePath {
         }
         theirs
     }
+    pub fn from_nibbles(mut theirs: mpt_trie::nibbles::Nibbles) -> Self {
+        let mut ours = CopyVec::new();
+        while !theirs.is_empty() {
+            ours.try_push(
+                U4::new(theirs.pop_next_nibble_front())
+                    .expect("mpt_trie returned an invalid nibble"),
+            )
+            .expect("mpt_trie should not have more than 64 nibbles")
+        }
+        Self(ours)
+    }
 }
 
 /// Map where keys are [up to 64 nibbles](TriePath), and values are either an
@@ -201,7 +212,7 @@ impl StateTrie {
     ) -> Option<Either<H256, AccountRlp>> {
         self.typed.insert(path, account)
     }
-    pub fn insert_branch(
+    pub fn insert_branch_by_path(
         &mut self,
         path: TriePath,
         hash: H256,
@@ -240,22 +251,34 @@ impl IntoIterator for StateTrie {
 /// See <https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie/#storage-trie>
 #[derive(Debug, Clone, Default)]
 pub struct StorageTrie {
-    typed: TypedMpt<Vec<u8>>,
+    /// This does NOT use [`TypedMpt`] - T could be anything!
+    map: BTreeMap<TriePath, Either<H256, Vec<u8>>>,
 }
 impl StorageTrie {
     pub fn insert(&mut self, path: TriePath, value: Vec<u8>) -> Option<Either<H256, Vec<u8>>> {
-        self.typed.insert(path, value)
+        self.map.insert(path, Either::Right(value))
     }
     pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option<Either<H256, Vec<u8>>> {
-        self.typed.insert_branch(path, hash)
+        self.map.insert(path, Either::Left(hash))
     }
     pub fn root(&self) -> H256 {
-        self.typed.root()
+        self.as_hashed_partial_trie().hash()
     }
     pub fn remove(&mut self, path: TriePath) -> Option<Either<H256, Vec<u8>>> {
-        self.typed.map.remove(&path)
+        self.map.remove(&path)
     }
     pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie {
-        self.typed.as_hashed_partial_trie()
+        let mut theirs = mpt_trie::partial_trie::HashedPartialTrie::default();
+        for (k, v) in &self.map {
+            let nibbles = k.into_nibbles();
+            match v {
+                Either::Left(h) => theirs.insert(nibbles, *h),
+                // TODO(0xaatif): why is RLP-encoding here the right thing to do?
+                //                (Our tests fail without it).
+                Either::Right(v) => theirs.insert(nibbles, &*rlp::encode(v)),
+            }
+            .expect("error in mpt_trie")
+        }
+        theirs
     }
 }