Fix bytecode hashing for Type2 SMT

Cargo.toml

@@ -45,7 +45,8 @@ bytes = "1.6.0"
 ciborium = "0.2.2"
 ciborium-io = "0.2.2"
 clap = { version = "4.5.7", features = ["derive", "env"] }
-compat = { path = "compat" }
+alloy-compat = "0.1.0"
+copyvec = "0.2.0"
 criterion = "0.5.1"
 dotenvy = "0.15.7"
 either = "1.12.0"
@@ -104,6 +105,7 @@ url = "2.5.2"
 winnow = "0.6.13"
 
 # local dependencies
+compat = { path = "compat" }
 evm_arithmetization = { path = "evm_arithmetization", version = "0.4.0", default-features = false }
 mpt_trie = { path = "mpt_trie", version = "0.4.1" }
 smt_trie = { path = "smt_trie", version = "0.1.1" }

evm_arithmetization/Cargo.toml

@@ -15,8 +15,13 @@ homepage.workspace = true
 keywords.workspace = true
 
 [dependencies]
+alloy.workspace = true
+alloy-compat.workspace = true
 anyhow.workspace = true
+bitvec.workspace = true
 bytes.workspace = true
+copyvec.workspace = true
+either.workspace = true
 env_logger.workspace = true
 ethereum-types.workspace = true
 hashbrown.workspace = true
@@ -43,14 +48,15 @@ serde = { workspace = true, features = ["derive"] }
 serde-big-array.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
-smt_trie = { workspace = true, optional = true }
+smt_trie = { workspace = true }
 starky = { workspace = true, features = ["parallel"] }
 static_assertions.workspace = true
 thiserror.workspace = true
 tiny-keccak.workspace = true
 tokio.workspace = true
 tower-lsp = "0.20.0"
 tracing.workspace = true
+u4.workspace = true
 url.workspace = true
 zk_evm_common.workspace = true
 zk_evm_proc_macro.workspace = true
@@ -64,7 +70,7 @@ ripemd.workspace = true
 default = ["eth_mainnet"]
 asmtools = ["hex"]
 polygon_pos = []
-cdk_erigon = ["smt_trie"]
+cdk_erigon = []
 eth_mainnet = []
 
 [[bin]]

evm_arithmetization/src/lib.rs

@@ -284,7 +284,9 @@ pub mod witness;
 pub mod curve_pairings;
 pub mod extension_tower;
 pub mod testing_utils;
+pub mod tries;
 pub mod util;
+pub mod world;
 
 // Public definitions and re-exports
 mod public_types;

trace_decoder/src/tries.rs → evm_arithmetization/src/tries.rs

@@ -7,10 +7,11 @@ use anyhow::ensure;
 use bitvec::{array::BitArray, slice::BitSlice};
 use copyvec::CopyVec;
 use ethereum_types::{Address, H256, U256};
-use evm_arithmetization::generation::mpt::AccountRlp;
 use mpt_trie::partial_trie::{HashedPartialTrie, Node, OnOrphanedHashNode, PartialTrie as _};
 use u4::{AsNibbles, U4};
 
+use crate::generation::mpt::AccountRlp;
+
 /// Bounded sequence of [`U4`],
 /// used as a key for [MPT](HashedPartialTrie) types in this module.
 ///

trace_decoder/src/world.rs → evm_arithmetization/src/world.rs

@@ -5,19 +5,43 @@ use anyhow::{ensure, Context as _};
 use either::Either;
 use ethereum_types::{Address, BigEndianHash as _, U256};
 use keccak_hash::H256;
+use smt_trie::code::hash_bytecode_h256;
+
+/// Utility trait to leverage a specific hash function across Type1 and Type2
+/// zkEVM variants.
+pub trait Hasher {
+    fn hash(bytes: &[u8]) -> H256;
+}
+
+pub struct PoseidonHash;
+pub struct KeccakHash;
+
+impl Hasher for PoseidonHash {
+    fn hash(bytes: &[u8]) -> H256 {
+        hash_bytecode_h256(bytes)
+    }
+}
+
+impl Hasher for KeccakHash {
+    fn hash(bytes: &[u8]) -> H256 {
+        keccak_hash::keccak(bytes)
+    }
+}
 
 use crate::tries::{MptKey, SmtKey, StateMpt, StorageTrie};
 
-/// The [core](crate::core) of this crate is agnostic over state and storage
-/// representations.
+/// The `core` module of the `trace_decoder` crate is agnostic over state and
+/// storage representations.
 ///
 /// This is the common interface to those data structures.
-/// See also [crate::_DEVELOPER_DOCS].
-pub(crate) trait World {
+pub trait World {
     /// (State) subtries may be _hashed out.
     /// This type is a key which may identify a subtrie.
     type SubtriePath;
 
+    /// Hasher to use for contract bytecode.
+    type CodeHasher: Hasher;
+
     //////////////////////
     /// Account operations
     //////////////////////
@@ -44,8 +68,8 @@ pub(crate) trait World {
     /// Creates a new account at `address` if it does not exist.
     fn set_code(&mut self, address: Address, code: Either<&[u8], H256>) -> anyhow::Result<()>;
 
-    /// The [core](crate::core) of this crate tracks required subtries for
-    /// proving.
+    /// The `core` module of the `trace_decoder` crate tracks required subtries
+    /// for proving.
     ///
     /// In case of a state delete, it may be that certain parts of the subtrie
     /// must be retained. If so, it will be returned as [`Some`].
@@ -148,6 +172,8 @@ impl Type1World {
 
 impl World for Type1World {
     type SubtriePath = MptKey;
+    type CodeHasher = KeccakHash;
+
     fn contains(&mut self, address: Address) -> anyhow::Result<bool> {
         Ok(self.state.get(keccak_hash::keccak(address)).is_some())
     }
@@ -170,7 +196,7 @@ impl World for Type1World {
     fn set_code(&mut self, address: Address, code: Either<&[u8], H256>) -> anyhow::Result<()> {
         let key = keccak_hash::keccak(address);
         let mut acct = self.state.get(key).unwrap_or_default();
-        acct.code_hash = code.right_or_else(keccak_hash::keccak);
+        acct.code_hash = code.right_or_else(Self::CodeHasher::hash);
         self.state.insert(key, acct)
     }
     fn reporting_destroy(&mut self, address: Address) -> anyhow::Result<Option<Self::SubtriePath>> {
@@ -248,6 +274,8 @@ impl World for Type1World {
 
 impl World for Type2World {
     type SubtriePath = SmtKey;
+    type CodeHasher = PoseidonHash;
+
     fn contains(&mut self, address: Address) -> anyhow::Result<bool> {
         Ok(self.accounts.contains_key(&address))
     }
@@ -269,10 +297,14 @@ impl World for Type2World {
         let acct = self.accounts.entry(address).or_default();
         match code {
             Either::Left(bytes) => {
-                acct.code = Some(keccak_hash::keccak(bytes).into_uint());
-                acct.code_length = Some(U256::from(bytes.len()))
+                acct.code_length = Some(U256::from(bytes.len()));
+                if bytes.is_empty() {
+                    acct.code_hash = None;
+                } else {
+                    acct.code_hash = Some(Self::CodeHasher::hash(bytes).into_uint());
+                }
             }
-            Either::Right(hash) => acct.code = Some(hash.into_uint()),
+            Either::Right(hash) => acct.code_hash = Some(hash.into_uint()),
         };
         Ok(())
     }
@@ -349,7 +381,7 @@ impl World for Type2World {
 pub struct Type2Entry {
     pub balance: Option<U256>,
     pub nonce: Option<U256>,
-    pub code: Option<U256>,
+    pub code_hash: Option<U256>,
     pub code_length: Option<U256>,
     pub storage: BTreeMap<U256, U256>,
 }
@@ -383,7 +415,7 @@ impl Type2World {
             Type2Entry {
                 balance,
                 nonce,
-                code,
+                code_hash,
                 code_length,
                 storage,
             },
@@ -394,7 +426,7 @@ impl Type2World {
             for (value, key_fn) in [
                 (balance, key_balance as fn(_) -> _),
                 (nonce, key_nonce),
-                (code, key_code),
+                (code_hash, key_code),
                 (code_length, key_code_length),
             ] {
                 if let Some(value) = value {

smt_trie/src/code.rs

@@ -1,11 +1,11 @@
 /// Functions to hash contract bytecode using Poseidon.
 /// See `hashContractBytecode()` in https://github.com/0xPolygonHermez/zkevm-commonjs/blob/main/src/smt-utils.js for reference implementation.
-use ethereum_types::U256;
+use ethereum_types::H256;
 use plonky2::field::types::Field;
 use plonky2::hash::poseidon::{self, Poseidon};
 
 use crate::smt::{HashOut, F};
-use crate::utils::hashout2u;
+use crate::utils::hashout2h;
 
 pub fn hash_contract_bytecode(mut code: Vec<u8>) -> HashOut {
     poseidon_pad_byte_vec(&mut code);
@@ -43,8 +43,8 @@ pub fn poseidon_pad_byte_vec(bytes: &mut Vec<u8>) {
     *bytes.last_mut().unwrap() |= 0x80;
 }
 
-pub fn hash_bytecode_u256(code: Vec<u8>) -> U256 {
-    hashout2u(hash_contract_bytecode(code))
+pub fn hash_bytecode_h256(code: &[u8]) -> H256 {
+    hashout2h(hash_contract_bytecode(code.to_vec()))
 }
 
 #[cfg(test)]

smt_trie/src/utils.rs

@@ -1,4 +1,4 @@
-use ethereum_types::U256;
+use ethereum_types::{H256, U256};
 use plonky2::field::types::{Field, PrimeField64};
 use plonky2::hash::poseidon::Poseidon;
 
@@ -49,6 +49,13 @@ pub fn hashout2u(h: HashOut) -> U256 {
     key2u(Key(h.elements))
 }
 
+/// Convert a `HashOut` to a `H256`.
+pub fn hashout2h(h: HashOut) -> H256 {
+    let mut it = [0; 32];
+    hashout2u(h).to_big_endian(&mut it);
+    H256(it)
+}
+
 /// Convert a `Key` to a `U256`.
 pub fn key2u(key: Key) -> U256 {
     U256(key.0.map(|x| x.to_canonical_u64()))

trace_decoder/Cargo.toml

@@ -11,14 +11,14 @@ keywords.workspace = true
 
 [dependencies]
 alloy.workspace = true
-alloy-compat = "0.1.0"
+alloy-compat.workspace = true
 anyhow.workspace = true
 bitflags.workspace = true
 bitvec.workspace = true
 bytes.workspace = true
 ciborium.workspace = true
 ciborium-io.workspace = true
-copyvec = "0.2.0"
+copyvec.workspace = true
 either.workspace = true
 enum-as-inner.workspace = true
 ethereum-types.workspace = true
@@ -43,7 +43,7 @@ zk_evm_common.workspace = true
 
 [dev-dependencies]
 alloy.workspace = true
-alloy-compat = "0.1.0"
+alloy-compat.workspace = true
 assert2 = "0.3.15"
 camino = "1.1.9"
 clap.workspace = true