0xaatif/messy smt2

trace_decoder/Cargo.toml

@@ -33,6 +33,7 @@ nunny = { workspace = true, features = ["serde"] }
 plonky2.workspace = true
 rlp.workspace = true
 serde.workspace = true
+smt_trie.workspace = true
 stackstack = "0.3.0"
 strum = { version = "0.26.3", features = ["derive"] }
 thiserror.workspace = true
@@ -52,7 +53,6 @@ libtest-mimic = "0.7.3"
 plonky2_maybe_rayon.workspace = true
 serde_json.workspace = true
 serde_path_to_error.workspace = true
-smt_trie.workspace = true
 zero.workspace = true
 
 [features]

trace_decoder/src/core.rs

@@ -7,6 +7,7 @@
 use alloy::primitives::address;
 use alloy_compat::Compat as _;
 use anyhow::{anyhow, bail, ensure, Context as _};
+use either::Either;
 use ethereum_types::{Address, H160, U256};
 use evm_arithmetization::{
     generation::{mpt::AccountRlp, TrieInputs},
@@ -19,20 +20,30 @@
 use nunny::NonEmpty;
 use zk_evm_common::gwei_to_wei;
 
-use crate::observer::Observer;
+use crate::{observer::Observer, typed_mpt::StateSmt};
 use crate::{
     typed_mpt::{ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie, TrieKey},
     BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage,
     OtherBlockData, SeparateStorageTriesPreImage, SeparateTriePreImage, SeparateTriePreImages,
     TxnInfo, TxnMeta, TxnTrace,
 };
 
+/// When parsing tries, which type to deserialize as.
+#[derive(Debug)]
+pub enum WireDisposition {
+    /// MPT
+    Type1,
+    /// SMT
+    Type2,
+}
+
 /// TODO(0xaatif): document this after <https://github.com/0xPolygonZero/zk_evm/issues/275>
 pub fn entrypoint(
     trace: BlockTrace,
     other: OtherBlockData,
     batch_size_hint: usize,
-    observer: &mut impl Observer<StateMpt>,
+    observer: &mut impl Observer<Either<StateMpt, StateSmt>>,
+    wire_disposition: WireDisposition,
 ) -> anyhow::Result<Vec<GenerationInputs>> {
     ensure!(batch_size_hint != 0);
 
@@ -41,7 +52,7 @@
         code_db,
         txn_info,
     } = trace;
-    let (state, storage, mut code) = start(trie_pre_images)?;
+    let (state, storage, mut code) = start(trie_pre_images, wire_disposition)?;
     code.extend(code_db);
 
     let OtherBlockData {
@@ -101,7 +112,10 @@
                 withdrawals,
                 ger_data,
                 tries: TrieInputs {
-                    state_trie: state.into(),
+                    state_trie: match state {
+                        Either::Left(mpt) => mpt.into(),
+                        Either::Right(_) => todo!("evm_arithmetization accepts an SMT"),
+                    },
                     transactions_trie: transaction.into(),
                     receipts_trie: receipt.into(),
                     storage_tries: storage.into_iter().map(|(k, v)| (k, v.into())).collect(),
@@ -127,9 +141,15 @@
 ///
 /// Turn either of those into our [`typed_mpt`](crate::typed_mpt)
 /// representations.
+#[allow(clippy::type_complexity)]
 fn start(
     pre_images: BlockTraceTriePreImages,
-) -> anyhow::Result<(StateMpt, BTreeMap<H256, StorageTrie>, Hash2Code)> {
+    wire_disposition: WireDisposition,
+) -> anyhow::Result<(
+    Either<StateMpt, StateSmt>,
+    BTreeMap<H256, StorageTrie>,
+    Hash2Code,
+)> {
     Ok(match pre_images {
         // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/401
         //                refactor our convoluted input types
@@ -177,17 +197,34 @@
                         .map(|v| (k, v))
                 })
                 .collect::<Result<_, _>>()?;
-            (state, storage, Hash2Code::new())
+            (Either::Left(state), storage, Hash2Code::new())
         }
         BlockTraceTriePreImages::Combined(CombinedPreImages { compact }) => {
             let instructions = crate::wire::parse(&compact)
                 .context("couldn't parse instructions from binary format")?;
-            let crate::type1::Frontend {
-                state,
-                storage,
-                code,
-            } = crate::type1::frontend(instructions)?;
-            (state, storage, code.into_iter().map(Into::into).collect())
+            let (state, storage, code) = match wire_disposition {
+                WireDisposition::Type1 => {
+                    let crate::type1::Frontend {
+                        state,
+                        storage,
+                        code,
+                    } = crate::type1::frontend(instructions)?;
+                    (
+                        Either::Left(state),
+                        storage,
+                        Hash2Code::from_iter(code.into_iter().map(NonEmpty::into_vec)),
+                    )
+                }
+                WireDisposition::Type2 => {
+                    let crate::type2::Frontend {
+                        trie,
+                        code,
+                        collation,
+                    } = crate::type2::frontend(instructions)?;
+                    todo!()
+                }
+            };
+            (state, storage, code)
         }
     })
 }

trace_decoder/src/lib.rs

@@ -57,15 +57,11 @@ mod interface;
 pub use interface::*;
 
 mod type1;
-// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275
-//                add backend/prod support for type 2
-#[cfg(test)]
-#[allow(dead_code)]
 mod type2;
 mod typed_mpt;
 mod wire;
 
-pub use core::entrypoint;
+pub use core::{entrypoint, WireDisposition};
 
 mod core;
 

trace_decoder/src/typed_mpt.rs

@@ -4,6 +4,7 @@ use core::fmt;
 use std::{collections::BTreeMap, marker::PhantomData};
 
 use copyvec::CopyVec;
+use either::Either;
 use ethereum_types::{Address, H256, U256};
 use evm_arithmetization::generation::mpt::AccountRlp;
 use mpt_trie::partial_trie::{HashedPartialTrie, Node, OnOrphanedHashNode, PartialTrie as _};
@@ -377,6 +378,7 @@ impl From<StateMpt> for HashedPartialTrie {
     }
 }
 
+#[derive(Clone, Debug)]
 pub struct StateSmt {
     address2state: BTreeMap<Address, AccountRlp>,
     hashed_out: BTreeMap<TrieKey, H256>,
@@ -468,6 +470,48 @@ impl From<StorageTrie> for HashedPartialTrie {
     }
 }
 
+macro_rules! either {
+    ($(fn $name:ident $params:tt -> $ret:ty);* $(;)?) => {
+        $(either!{ @ fn $name $params -> $ret })*
+    };
+    (@ fn $name:ident(&self $(, $var:ident : $ty:ty)* $(,)?) -> $ret:ty) => {
+        fn $name(&self $(, $var: $ty)*) -> $ret { match self {
+                Either::Left(it) => it.$name($($var),*),
+                Either::Right(it) => it.$name($($var),*),
+        }}
+    };
+    (@ fn $name:ident(&mut self $(, $var:ident : $ty:ty)* $(,)?) -> $ret:ty) => {
+        fn $name(&mut self $(, $var: $ty)*) -> $ret { match self {
+                Either::Left(it) => it.$name($($var),*),
+                Either::Right(it) => it.$name($($var),*),
+        }}
+    };
+}
+
+impl<L, R> StateTrie for Either<L, R>
+where
+    L: StateTrie,
+    R: StateTrie,
+{
+    either! {
+        fn insert_by_address(
+            &mut self,
+            address: Address,
+            account: AccountRlp,
+        ) -> anyhow::Result<Option<AccountRlp>>;
+        fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()>;
+        fn get_by_address(&self, address: Address) -> Option<AccountRlp>;
+        fn reporting_remove(&mut self, address: Address) -> anyhow::Result<Option<TrieKey>>;
+        fn mask(&mut self, address: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()>;
+        fn root(&self) -> H256;
+    }
+    fn iter(&self) -> impl Iterator<Item = (H256, AccountRlp)> + '_ {
+        self.as_ref()
+            .map_left(|it| it.iter())
+            .map_right(|it| it.iter())
+    }
+}
+
 /// If a branch collapse occurred after a delete, then we must ensure that
 /// the other single child that remains also is not hashed when passed into
 /// plonky2. Returns the key to the remaining child if a collapse occurred.

zero/Cargo.toml

@@ -15,6 +15,7 @@ alloy-compat = "0.1.0"
 anyhow.workspace = true
 async-stream.workspace = true
 axum.workspace = true
+cfg-if = "1.0.0"
 clap = { workspace = true, features = ["derive", "string"] }
 compat.workspace = true
 directories = "5.0.1"

zero/src/prover.rs

@@ -21,7 +21,7 @@ use tokio::io::AsyncWriteExt;
 use tokio::sync::mpsc::Receiver;
 use tokio::sync::{oneshot, Semaphore};
 use trace_decoder::observer::DummyObserver;
-use trace_decoder::{BlockTrace, OtherBlockData};
+use trace_decoder::{BlockTrace, OtherBlockData, WireDisposition};
 use tracing::{error, info};
 
 use crate::fs::generate_block_proof_file_name;
@@ -39,6 +39,14 @@ use crate::proof_types::GeneratedBlockProof;
 // batches as soon as they are generated.
 static PARALLEL_BLOCK_PROVING_PERMIT_POOL: Semaphore = Semaphore::const_new(0);
 
+const WIRE_DISPOSITION: WireDisposition = {
+    cfg_if::cfg_if!(if #[cfg(feature = "eth_mainnet")] {
+        WireDisposition::Type1
+    } else {
+        compile_error!("must select a feature");
+    })
+};
+
 #[derive(Debug, Clone)]
 pub struct ProverConfig {
     pub batch_size: usize,
@@ -88,6 +96,7 @@ impl BlockProverInput {
             self.other_data,
             batch_size,
             &mut DummyObserver::new(),
+            WIRE_DISPOSITION,
         )?;
 
         // Create segment proof.
@@ -181,6 +190,7 @@ impl BlockProverInput {
             self.other_data,
             batch_size,
             &mut DummyObserver::new(),
+            WIRE_DISPOSITION,
         )?;
 
         let seg_ops = ops::SegmentProofTestOnly {