Missing changes in latest commit

smt_trie/src/code.rs

@@ -11,7 +11,8 @@ pub fn hash_contract_bytecode(mut code: Vec<u8>) -> HashOut {
     poseidon_pad_byte_vec(&mut code);
 
     // println!("code bytes = {:?}", code);
-    // println!("hash = {:?}", hashout2u(poseidon_hash_padded_byte_vec(code.clone())));
+    // println!("hash = {:?}",
+    // hashout2u(poseidon_hash_padded_byte_vec(code.clone())));
     poseidon_hash_padded_byte_vec(code)
 }
 

smt_trie/src/smt.rs

@@ -6,6 +6,7 @@ use std::collections::{BTreeMap, HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 
 use ethereum_types::{H256, U256};
+use mpt_trie::partial_trie::HashedPartialTrie;
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::types::{Field, PrimeField64};
 use plonky2::hash::poseidon::{Poseidon, PoseidonHash};
@@ -17,7 +18,6 @@ use crate::db::{Db, MemoryDb};
 use crate::utils::{
     f2limbs, get_unique_sibling, hash0, hash_key_hash, hashout2u, key2u, limbs2f, u2h, u2k,
 };
-use mpt_trie::partial_trie::HashedPartialTrie;
 
 pub(crate) const HASH_TYPE: u8 = 0;
 pub(crate) const INTERNAL_TYPE: u8 = 1;
@@ -521,7 +521,8 @@ fn serialize<D: Db>(
         let index = v.len();
         v.push(HASH_TYPE.into());
         v.push(key2u(key));
-        if index == 0 { // Empty hash node is at the beggining of the segment
+        if index == 0 {
+            // Empty hash node is at the beggining of the segment
             index
         } else {
             index + offset
@@ -558,7 +559,7 @@ fn serialize<D: Db>(
                 keys_to_include,
                 offset,
             )
-                .into();
+            .into();
             v[index + 1] = i_left;
             let i_right = serialize(
                 smt,
@@ -568,7 +569,7 @@ fn serialize<D: Db>(
                 keys_to_include,
                 offset,
             )
-                .into();
+            .into();
             v[index + 2] = i_right;
             index + offset
         }

trace_decoder/src/core.rs

@@ -10,21 +10,22 @@ use ethereum_types::{Address, BigEndianHash as _, U256};
 use evm_arithmetization::{
     generation::TrieInputs,
     proof::{BlockMetadata, TrieRoots},
+    world::{
+        tries::{MptKey, ReceiptTrie, StateMpt, StorageTrie, TransactionTrie},
+        type1, type2,
+        world::{StateWorld, Type1World, Type2World, World},
+    },
     GenerationInputs,
 };
 use itertools::Itertools as _;
 use keccak_hash::H256;
 use mpt_trie::partial_trie::PartialTrie as _;
 use nunny::NonEmpty;
+use smt_trie::code::hash_bytecode_u256;
 use zk_evm_common::gwei_to_wei;
 
+use crate::observer::{DummyObserver, Observer};
 use crate::{
-    observer::{DummyObserver, Observer},
-    world::Type2World,
-};
-use crate::{
-    tries::{MptKey, ReceiptTrie, StateMpt, StorageTrie, TransactionTrie},
-    world::{Type1World, World},
     BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage,
     OtherBlockData, SeparateStorageTriesPreImage, SeparateTriePreImage, SeparateTriePreImages,
     TxnInfo, TxnMeta, TxnTrace,
@@ -133,9 +134,17 @@ pub fn entrypoint(
                  after,
                  withdrawals,
              }| {
-                let (state, storage) = world
-                    .expect_left("TODO(0xaatif): evm_arithemetization accepts an SMT")
-                    .into_state_and_storage();
+                let contract_code_hash = |it: &[u8]| {
+                    if let Either::Right(_type2world) = &world {
+                        Either::Right(hash_bytecode_u256(it.to_vec()))
+                    } else {
+                        Either::Left(keccak_hash::keccak(&it))
+                    }
+                };
+                let contract_code = contract_code
+                    .into_iter()
+                    .map(|it| (contract_code_hash(&it), it))
+                    .collect();
                 GenerationInputs {
                     txn_number_before: first_txn_ix.into(),
                     gas_used_before: running_gas_used.into(),
@@ -147,18 +156,14 @@ pub fn entrypoint(
                     withdrawals,
                     ger_data,
                     tries: TrieInputs {
-                        state_trie: state.into(),
+                        state_trie: StateWorld { state: world },
                         transactions_trie: transaction.into(),
                         receipts_trie: receipt.into(),
-                        storage_tries: storage.into_iter().map(|(k, v)| (k, v.into())).collect(),
                     },
                     trie_roots_after: after,
                     checkpoint_state_trie_root,
                     checkpoint_consolidated_hash,
-                    contract_code: contract_code
-                        .into_iter()
-                        .map(|it| (keccak_hash::keccak(&it), it))
-                        .collect(),
+                    contract_code,
                     block_metadata: b_meta.clone(),
                     block_hashes: b_hashes.clone(),
                     burn_addr,
@@ -230,23 +235,22 @@ fn start(
             )
         }
         BlockTraceTriePreImages::Combined(CombinedPreImages { compact }) => {
-            let instructions = crate::wire::parse(&compact)
+            let instructions = evm_arithmetization::world::wire::parse(&compact)
                 .context("couldn't parse instructions from binary format")?;
             match wire_disposition {
                 WireDisposition::Type1 => {
-                    let crate::type1::Frontend {
+                    let type1::Frontend {
                         state,
                         storage,
                         code,
-                    } = crate::type1::frontend(instructions)?;
+                    } = type1::frontend(instructions)?;
                     (
                         Either::Left(Type1World::new(state, storage)?),
                         Hash2Code::from_iter(code.into_iter().map(NonEmpty::into_vec)),
                     )
                 }
                 WireDisposition::Type2 => {
-                    let crate::type2::Frontend { world: trie, code } =
-                        crate::type2::frontend(instructions)?;
+                    let type2::Frontend { world: trie, code } = type2::frontend(instructions)?;
                     (
                         Either::Right(trie),
                         Hash2Code::from_iter(code.into_iter().map(NonEmpty::into_vec)),

trace_decoder/src/interface.rs

@@ -67,7 +67,7 @@ pub enum SeparateTriePreImage {
 #[serde(rename_all = "snake_case")]
 pub struct CombinedPreImages {
     /// Compact combined state and storage tries.
-    #[serde(with = "crate::hex")]
+    #[serde(with = "evm_arithmetization::world::hex")]
     pub compact: Vec<u8>,
 }
 
@@ -100,13 +100,13 @@ pub struct TxnMeta {
     /// Txn byte code. This is also the raw RLP bytestring inserted into the txn
     /// trie by this txn. Note that the key is not included and this is only
     /// the rlped value of the node!
-    #[serde(with = "crate::hex")]
+    #[serde(with = "evm_arithmetization::world::hex")]
     pub byte_code: Vec<u8>,
 
     /// Rlped bytes of the new receipt value inserted into the receipt trie by
     /// this txn. Note that the key is not included and this is only the rlped
     /// value of the node!
-    #[serde(with = "crate::hex")]
+    #[serde(with = "evm_arithmetization::world::hex")]
     pub new_receipt_trie_node_byte: Vec<u8>,
 
     /// Gas used by this txn (Note: not cumulative gas used).
@@ -160,7 +160,7 @@ pub enum ContractCodeUsage {
 
     /// Contract was created (and these are the bytes). Note that this new
     /// contract code will not appear in the [`BlockTrace`] map.
-    Write(#[serde(with = "crate::hex")] Vec<u8>),
+    Write(#[serde(with = "evm_arithmetization::world::hex")] Vec<u8>),
 }
 
 /// Other data that is needed for proof gen.

trace_decoder/src/observer.rs

@@ -1,9 +1,9 @@
 use std::marker::PhantomData;
 
 use ethereum_types::U256;
+use evm_arithmetization::world::tries::{ReceiptTrie, TransactionTrie};
 
 use crate::core::IntraBlockTries;
-use crate::tries::{ReceiptTrie, TransactionTrie};
 
 /// Observer API for the trace decoder.
 /// Observer is used to collect various debugging and metadata info

trace_decoder/tests/consistent-with-header.rs

@@ -7,9 +7,13 @@ mod common;
 use alloy_compat::Compat as _;
 use assert2::check;
 use common::{cases, Case};
+use either::Either;
+use ethereum_types::BigEndianHash;
 use itertools::Itertools;
+use keccak_hash::H256;
 use libtest_mimic::{Arguments, Trial};
 use mpt_trie::partial_trie::PartialTrie as _;
+use smt_trie::utils::hashout2u;
 use trace_decoder::observer::DummyObserver;
 use zero::prover::WIRE_DISPOSITION;
 
@@ -36,12 +40,24 @@ fn main() -> anyhow::Result<()> {
                 check!(gen_inputs.len() >= 2);
                 check!(
                     Some(other.checkpoint_state_trie_root)
-                        == gen_inputs.first().map(|it| it.tries.state_trie.hash())
+                        == gen_inputs
+                            .first()
+                            .map(|it| match &it.tries.state_trie.state {
+                                Either::Left(type1world) => type1world.state_trie().hash(),
+                                Either::Right(type2world) =>
+                                    H256::from_uint(&hashout2u(type2world.as_smt().root)),
+                            })
                 );
                 let pairs = || gen_inputs.iter().tuple_windows::<(_, _)>();
                 check!(
                     pairs().position(|(before, after)| {
-                        before.trie_roots_after.state_root != after.tries.state_trie.hash()
+                        let after_hash = match &after.tries.state_trie.state {
+                            Either::Left(type1world) => type1world.state_trie().hash(),
+                            Either::Right(type2world) => {
+                                H256::from_uint(&hashout2u(type2world.as_smt().root))
+                            }
+                        };
+                        before.trie_roots_after.state_root != after_hash
                     }) == None
                 );
                 check!(

zero/Cargo.toml

@@ -60,7 +60,7 @@ mockall = "0.13.0"
 
 
 [features]
-default = ["cdk_erigon"]
+default = ["eth_mainnet"]
 eth_mainnet = ["evm_arithmetization/eth_mainnet", "trace_decoder/eth_mainnet"]
 cdk_erigon = ["evm_arithmetization/cdk_erigon", "trace_decoder/cdk_erigon"]
 polygon_pos = ["evm_arithmetization/polygon_pos", "trace_decoder/polygon_pos"]

zero/src/trie_diff/mod.rs

@@ -1,4 +1,4 @@
-use evm_arithmetization::generation::mpt::{AccountRlp, LegacyReceiptRlp};
+use evm_arithmetization::generation::mpt::{LegacyReceiptRlp, MptAccountRlp};
 use evm_arithmetization::generation::DebugOutputTries;
 use mpt_trie::debug_tools::diff::{create_full_diff_between_tries, DiffPoint};
 use mpt_trie::utils::TrieNodeType;
@@ -85,7 +85,7 @@ pub fn compare_tries(
     }
 
     let state_trie_diff = create_full_diff_between_tries(&left.state_trie, &right.state_trie);
-    compare_tries_and_output_results::<usize, AccountRlp>(
+    compare_tries_and_output_results::<usize, MptAccountRlp>(
         "state trie",
         state_trie_diff.diff_points,
         block_number,