refactor: use typed tries in trace_decoder

evm_arithmetization/src/generation/mpt.rs

@@ -15,7 +15,7 @@ use crate::util::h2u;
 use crate::witness::errors::{ProgramError, ProverInputError};
 use crate::Node;
 
-#[derive(RlpEncodable, RlpDecodable, Debug)]
+#[derive(RlpEncodable, RlpDecodable, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub struct AccountRlp {
     pub nonce: U256,
     pub balance: U256,

trace_decoder/Cargo.toml

@@ -16,6 +16,7 @@ bitvec = { workspace = true }
 bytes = { workspace = true }
 ciborium = { workspace = true }
 ciborium-io = { workspace = true }
+copyvec = "0.2.0"
 either = { workspace = true }
 enum-as-inner = { workspace = true }
 ethereum-types = { workspace = true }
@@ -31,6 +32,7 @@ plonky2 = { workspace = true }
 rlp = { workspace = true }
 serde = { workspace = true }
 smt_trie = { workspace = true }
+stackstack = "0.3.0"
 thiserror = { workspace = true }
 u4 = { workspace = true }
 winnow = { workspace = true }

trace_decoder/src/decoding.rs

@@ -67,7 +67,7 @@ impl std::fmt::Display for TraceParsingError {
         let h_slot = self.slot.map(|slot| {
             let mut buf = [0u8; 64];
             slot.to_big_endian(&mut buf);
-            hash(&buf)
+            hash(buf)
         });
         write!(
             f,
@@ -225,8 +225,13 @@ 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().clone(),
+            storage: self
+                .tries
+                .storage
+                .iter()
+                .map(|(k, v)| (*k, v.as_hashed_partial_trie().clone()))
+                .collect(),
             ..Default::default()
         };
 
@@ -319,7 +324,7 @@ impl ProcessedBlockTrace {
             .map(|(k, v)| {
                 (
                     Nibbles::from_h256_be(hash(
-                        &Nibbles::from_h256_be(H256::from_uint(k)).bytes_be(),
+                        Nibbles::from_h256_be(H256::from_uint(k)).bytes_be(),
                     )),
                     v,
                 )
@@ -480,7 +485,7 @@ impl ProcessedBlockTrace {
 
             for (slot, val) in storage_writes
                 .iter()
-                .map(|(k, v)| (Nibbles::from_h256_be(hash(&k.bytes_be())), v))
+                .map(|(k, v)| (Nibbles::from_h256_be(hash(k.bytes_be())), v))
             {
                 // If we are writing a zero, then we actually need to perform a delete.
                 match val == &ZERO_STORAGE_SLOT_VAL_RLPED {

trace_decoder/src/lib.rs

@@ -91,17 +91,20 @@ mod type1;
 #[cfg(test)]
 #[allow(dead_code)]
 mod type2;
+mod typed_mpt;
 mod wire;
 
 use std::collections::HashMap;
 
 use ethereum_types::{Address, U256};
 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 processed_block_trace::ProcessedTxnInfo;
 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
@@ -296,7 +299,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::{
@@ -320,11 +322,44 @@ 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_hash_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_hash(path, h)?;
+                                    }
+                                };
+                                anyhow::Ok(acc)
+                            })
+                            .map(|v| (k, v))
+                    })
+                    .collect::<Result<_, _>>()?,
             },
             extra_code_hash_mappings: None,
         },
@@ -339,7 +374,10 @@ pub fn entrypoint(
             ProcessedBlockTracePreImages {
                 tries: PartialTriePreImages {
                     state,
-                    storage: storage.into_iter().collect(),
+                    storage: storage
+                        .into_iter()
+                        .map(|(path, trie)| (path.into_hash_left_padded(), trie))
+                        .collect(),
                 },
                 extra_code_hash_mappings: match code.is_empty() {
                     true => None,
@@ -356,11 +394,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()
+        .map(|(addr, data)| (addr.into_hash_left_padded(), data))
         .collect::<Vec<_>>();
 
     let code_db = {
@@ -415,14 +450,10 @@ pub fn entrypoint(
     .into_txn_proof_gen_ir(other)?)
 }
 
-fn hash(bytes: &[u8]) -> ethereum_types::H256 {
-    keccak_hash::keccak(bytes).0.into()
-}
-
 #[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/processed_block_trace.rs

@@ -93,7 +93,7 @@ impl TxnInfo {
             nodes_used_by_txn.storage_accesses.push((
                 hashed_addr,
                 storage_access_keys
-                    .map(|k| Nibbles::from_h256_be(hash(&k.0)))
+                    .map(|H256(bytes)| Nibbles::from_h256_be(hash(bytes)))
                     .collect(),
             ));