diff --git a/Cargo.lock b/Cargo.lock index 3f717b5f3..f1b011947 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5723,6 +5723,7 @@ dependencies = [ "anyhow", "async-stream", "axum", + "cfg-if", "clap", "compat", "directories", diff --git a/trace_decoder/Cargo.toml b/trace_decoder/Cargo.toml index de9b389a2..4b8d8e7bc 100644 --- a/trace_decoder/Cargo.toml +++ b/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] diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 31f065b3b..9a1ea9bbc 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -7,6 +7,7 @@ use std::{ 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,7 +20,7 @@ use mpt_trie::partial_trie::PartialTrie as _; 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, @@ -27,12 +28,22 @@ use crate::{ 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 pub fn entrypoint( trace: BlockTrace, other: OtherBlockData, batch_size_hint: usize, - observer: &mut impl Observer, + observer: &mut impl Observer>, + wire_disposition: WireDisposition, ) -> anyhow::Result> { ensure!(batch_size_hint != 0); @@ -41,7 +52,7 @@ pub fn entrypoint( 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 @@ pub fn entrypoint( 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 @@ pub fn entrypoint( /// /// 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, Hash2Code)> { + wire_disposition: WireDisposition, +) -> anyhow::Result<( + Either, + BTreeMap, + Hash2Code, +)> { Ok(match pre_images { // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/401 // refactor our convoluted input types @@ -177,17 +197,34 @@ fn start( .map(|v| (k, v)) }) .collect::>()?; - (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) } }) } diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index 049472c40..eea5ebe80 100644 --- a/trace_decoder/src/lib.rs +++ b/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; diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 8baf3cf29..8594ed11d 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/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 for HashedPartialTrie { } } +#[derive(Clone, Debug)] pub struct StateSmt { address2state: BTreeMap, hashed_out: BTreeMap, @@ -468,6 +470,48 @@ impl From 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 StateTrie for Either +where + L: StateTrie, + R: StateTrie, +{ + either! { + fn insert_by_address( + &mut self, + address: Address, + account: AccountRlp, + ) -> anyhow::Result>; + fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()>; + fn get_by_address(&self, address: Address) -> Option; + fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; + fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; + fn root(&self) -> H256; + } + fn iter(&self) -> impl Iterator + '_ { + 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. diff --git a/zero/Cargo.toml b/zero/Cargo.toml index 22c2a8bfb..7cbf2f351 100644 --- a/zero/Cargo.toml +++ b/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" diff --git a/zero/src/prover.rs b/zero/src/prover.rs index 6a194ddf1..e2207d993 100644 --- a/zero/src/prover.rs +++ b/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 {