From da7fea8f4f57b14496123a6fd7b44a2130c0f6fe Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Wed, 2 Oct 2024 17:17:02 +0100 Subject: [PATCH 01/13] mark: 0xaatif/messy-smt2 From 5c57657cc4a847812f01d1cd7f6a915d21fc5f0f Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Thu, 3 Oct 2024 03:15:25 +0100 Subject: [PATCH 02/13] wip: initial implementation --- Cargo.lock | 1 + trace_decoder/Cargo.toml | 2 +- trace_decoder/src/core.rs | 61 +++++++++++++++++++++++++++------- trace_decoder/src/lib.rs | 6 +--- trace_decoder/src/typed_mpt.rs | 44 ++++++++++++++++++++++++ zero/Cargo.toml | 1 + zero/src/prover.rs | 12 ++++++- 7 files changed, 108 insertions(+), 19 deletions(-) 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 { From 81b3b450c393451d4300d645290f22186c6513c3 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Fri, 4 Oct 2024 06:51:46 +0100 Subject: [PATCH 03/13] wip --- Cargo.lock | 10 ++ Cargo.toml | 3 - trace_decoder/Cargo.toml | 1 + trace_decoder/benches/block_processing.rs | 2 + trace_decoder/src/core.rs | 1 + trace_decoder/src/type2.rs | 20 ++-- trace_decoder/src/typed_mpt.rs | 111 ++++++++++++++++-- trace_decoder/tests/consistent-with-header.rs | 2 + trace_decoder/tests/simulate-execution.rs | 16 ++- zero/src/bin/rpc.rs | 2 + zero/src/bin/trie_diff.rs | 5 + zero/src/prover.rs | 14 ++- 12 files changed, 154 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f1b011947..3288953cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1149,6 +1149,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "build-array" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67ef4e2687af237b2646687e19a0643bc369878216122e46c3f1a01c56baa9d5" +dependencies = [ + "arrayvec", +] + [[package]] name = "bumpalo" version = "3.16.0" @@ -5039,6 +5048,7 @@ dependencies = [ "assert2", "bitflags 2.6.0", "bitvec", + "build-array", "bytes", "camino", "ciborium", diff --git a/Cargo.toml b/Cargo.toml index 0bd15a3bd..4c8685f60 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,7 +41,6 @@ axum = "0.7.5" bitflags = "2.5.0" bitvec = "1.0.1" bytes = "1.6.0" -cargo_metadata = "0.18.1" ciborium = "0.2.2" ciborium-io = "0.2.2" clap = { version = "4.5.7", features = ["derive", "env"] } @@ -72,7 +71,6 @@ nunny = "0.2.1" once_cell = "1.19.0" paladin-core = "0.4.2" parking_lot = "0.12.3" -paste = "1.0.15" pest = "2.7.10" pest_derive = "2.7.10" pretty_env_logger = "0.5.0" @@ -94,7 +92,6 @@ syn = "2.0" thiserror = "1.0.61" tiny-keccak = "2.0.2" tokio = { version = "1.38.0", features = ["full"] } -toml = "0.8.14" tower = "0.4" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/trace_decoder/Cargo.toml b/trace_decoder/Cargo.toml index 4b8d8e7bc..5bdd24f12 100644 --- a/trace_decoder/Cargo.toml +++ b/trace_decoder/Cargo.toml @@ -15,6 +15,7 @@ alloy-compat = "0.1.0" anyhow.workspace = true bitflags.workspace = true bitvec.workspace = true +build-array = "0.1.2" bytes.workspace = true ciborium.workspace = true ciborium-io.workspace = true diff --git a/trace_decoder/benches/block_processing.rs b/trace_decoder/benches/block_processing.rs index adefdae3f..4e3582e98 100644 --- a/trace_decoder/benches/block_processing.rs +++ b/trace_decoder/benches/block_processing.rs @@ -8,6 +8,7 @@ use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; use trace_decoder::observer::DummyObserver; use trace_decoder::{BlockTrace, OtherBlockData}; +use zero::prover::WIRE_DISPOSITION; #[derive(Clone, Debug, serde::Deserialize)] pub struct ProverInput { @@ -39,6 +40,7 @@ fn criterion_benchmark(c: &mut Criterion) { other_data, batch_size, &mut DummyObserver::new(), + WIRE_DISPOSITION, ) .unwrap() }, diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 9a1ea9bbc..b006f970d 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -221,6 +221,7 @@ fn start( code, collation, } = crate::type2::frontend(instructions)?; + todo!() } }; diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index dd3e45c4b..eb70978f9 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -14,8 +14,10 @@ use itertools::{EitherOrBoth, Itertools as _}; use nunny::NonEmpty; use plonky2::field::types::Field; -use crate::wire::{Instruction, SmtLeaf, SmtLeafType}; - +use crate::{ + typed_mpt::StateSmt, + wire::{Instruction, SmtLeaf, SmtLeafType}, +}; type SmtTrie = smt_trie::smt::Smt; #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -138,7 +140,10 @@ fn node2trie( Either::Right(it) => Either::Right(it), }); + let mut lens = std::collections::BTreeMap::<_, usize>::new(); + for (path, hash) in hashes { + *lens.entry(path.len()).or_default() += 1; // needs to be called before `set`, below, "to avoid any issues" according // to the smt docs. trie.set_hash( @@ -151,6 +156,7 @@ fn node2trie( }, ) } + dbg!(lens); let mut collated = HashMap::::new(); for SmtLeaf { @@ -235,10 +241,10 @@ fn test_tries() { println!("case {}", ix); let instructions = crate::wire::parse(&case.bytes).unwrap(); let frontend = frontend(instructions).unwrap(); - assert_eq!(case.expected_state_root, { - let mut it = [0; 32]; - smt_trie::utils::hashout2u(frontend.trie.root).to_big_endian(&mut it); - ethereum_types::H256(it) - }); + // assert_eq!(case.expected_state_root, { + // let mut it = [0; 32]; + // smt_trie::utils::hashout2u(frontend.trie.root).to_big_endian(&mut + // it); ethereum_types::H256(it) + // }); } } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 8594ed11d..9d29d508b 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -3,9 +3,10 @@ use core::fmt; use std::{collections::BTreeMap, marker::PhantomData}; +use bitvec::{order::Msb0, view::BitView as _}; use copyvec::CopyVec; use either::Either; -use ethereum_types::{Address, H256, U256}; +use ethereum_types::{Address, BigEndianHash as _, H256, U256}; use evm_arithmetization::generation::mpt::AccountRlp; use mpt_trie::partial_trie::{HashedPartialTrie, Node, OnOrphanedHashNode, PartialTrie as _}; use u4::{AsNibbles, U4}; @@ -162,6 +163,17 @@ impl TrieKey { } Self(ours) } + fn into_bits(self) -> smt_trie::bits::Bits { + let mut bits = smt_trie::bits::Bits::default(); + for component in self.0 { + let byte = component as u8; + // the four high bits are zero + for bit in byte.view_bits::().into_iter().by_vals().skip(4) { + bits.push_bit(bit); + } + } + bits + } pub fn into_hash(self) -> Option { let Self(nibbles) = self; @@ -279,7 +291,6 @@ pub trait StateTrie { 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>; /// _Hash out_ parts of the trie that aren't in `txn_ixs`. @@ -305,6 +316,10 @@ impl StateMpt { }, } } + /// Insert a _hashed out_ part of the trie + pub fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> { + self.typed.insert_hash(key, hash) + } #[deprecated = "prefer operations on `Address` where possible, as SMT support requires this"] pub fn insert_by_hashed_address( &mut self, @@ -332,10 +347,6 @@ impl StateTrie for StateMpt { #[expect(deprecated)] self.insert_by_hashed_address(keccak_hash::keccak(address), account) } - /// Insert an _hashed out_ part of the trie - fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> { - self.typed.insert_hash(key, hash) - } fn get_by_address(&self, address: Address) -> Option { self.typed .get(TrieKey::from_hash(keccak_hash::keccak(address))) @@ -378,6 +389,10 @@ impl From for HashedPartialTrie { } } +// TODO(0xaatif): trackme +// We're covering for [`smt_trie`] in a couple of ways: +// - insertion operations aren't fallible, they just panic. +// - it documents a requirement that `set_hash` is called before `set`. #[derive(Clone, Debug)] pub struct StateSmt { address2state: BTreeMap, @@ -392,10 +407,6 @@ impl StateTrie for StateSmt { ) -> anyhow::Result> { Ok(self.address2state.insert(address, account)) } - fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> { - self.hashed_out.insert(key, hash); - Ok(()) - } fn get_by_address(&self, address: Address) -> Option { self.address2state.get(&address).copied() } @@ -413,7 +424,84 @@ impl StateTrie for StateSmt { .map(|(addr, acct)| (keccak_hash::keccak(addr), *acct)) } fn root(&self) -> H256 { - todo!() + conv_hash::smt2eth(self.as_smt().root) + } +} + +impl StateSmt { + fn as_smt(&self) -> smt_trie::smt::Smt { + let Self { + address2state, + hashed_out, + } = self; + let mut smt = smt_trie::smt::Smt::::default(); + for (k, v) in hashed_out { + smt.set_hash(k.into_bits(), conv_hash::eth2smt(*v)); + } + for ( + addr, + AccountRlp { + nonce, + balance, + storage_root, + code_hash, + }, + ) in address2state + { + smt.set(smt_trie::keys::key_nonce(*addr), *nonce); + smt.set(smt_trie::keys::key_balance(*addr), *balance); + smt.set(smt_trie::keys::key_code(*addr), code_hash.into_uint()); + smt.set( + // REVIEW(0xaatif): I don't know what to do here + smt_trie::keys::key_storage(*addr, U256::zero()), + storage_root.into_uint(), + ); + } + smt + } +} + +mod conv_hash { + use std::array; + + use ethereum_types::H256; + use itertools::Itertools as _; + use plonky2::{field::goldilocks_field::GoldilocksField, hash::hash_types::HashOut}; + + pub fn eth2smt(H256(bytes): H256) -> smt_trie::smt::HashOut { + let mut bytes = bytes.into_iter(); + // (no unsafe, no unstable) + let ret = HashOut { + elements: array::from_fn(|_ix| { + let (a, b, c, d, e, f, g, h) = bytes.next_tuple().unwrap(); + // REVIEW(0xaatif): what endianness? + GoldilocksField(u64::from_be_bytes([a, b, c, d, e, f, g, h])) + }), + }; + assert_eq!(bytes.len(), 0); + ret + } + pub fn smt2eth(HashOut { elements }: smt_trie::smt::HashOut) -> H256 { + H256( + build_array::ArrayBuilder::from_iter( + elements + .into_iter() + .flat_map(|GoldilocksField(u)| u.to_be_bytes()), + ) + .build_exact() + .unwrap(), + ) + } + + #[test] + fn test() { + for h in [ + H256::zero(), + H256(array::from_fn(|ix| ix as u8)), + H256([u8::MAX; 32]), + ] { + assert_eq!(smt2eth(eth2smt(h)), h); + } } } @@ -499,7 +587,6 @@ where 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<()>; diff --git a/trace_decoder/tests/consistent-with-header.rs b/trace_decoder/tests/consistent-with-header.rs index 609fd57bb..63df41b4f 100644 --- a/trace_decoder/tests/consistent-with-header.rs +++ b/trace_decoder/tests/consistent-with-header.rs @@ -11,6 +11,7 @@ use itertools::Itertools; use libtest_mimic::{Arguments, Trial}; use mpt_trie::partial_trie::PartialTrie as _; use trace_decoder::observer::DummyObserver; +use zero::prover::WIRE_DISPOSITION; fn main() -> anyhow::Result<()> { let mut trials = vec![]; @@ -29,6 +30,7 @@ fn main() -> anyhow::Result<()> { other.clone(), batch_size, &mut DummyObserver::new(), + WIRE_DISPOSITION, ) .map_err(|e| format!("{e:?}"))?; // get the full cause chain check!(gen_inputs.len() >= 2); diff --git a/trace_decoder/tests/simulate-execution.rs b/trace_decoder/tests/simulate-execution.rs index d0476c2b7..fc7136c34 100644 --- a/trace_decoder/tests/simulate-execution.rs +++ b/trace_decoder/tests/simulate-execution.rs @@ -9,6 +9,7 @@ use common::{cases, Case}; use libtest_mimic::{Arguments, Trial}; use plonky2::field::goldilocks_field::GoldilocksField; use trace_decoder::observer::DummyObserver; +use zero::prover::WIRE_DISPOSITION; fn main() -> anyhow::Result<()> { let mut trials = vec![]; @@ -20,11 +21,16 @@ fn main() -> anyhow::Result<()> { other, } in cases()? { - let gen_inputs = - trace_decoder::entrypoint(trace, other, batch_size, &mut DummyObserver::new()) - .context(format!( - "error in `trace_decoder` for {name} at batch size {batch_size}" - ))?; + let gen_inputs = trace_decoder::entrypoint( + trace, + other, + batch_size, + &mut DummyObserver::new(), + WIRE_DISPOSITION, + ) + .context(format!( + "error in `trace_decoder` for {name} at batch size {batch_size}" + ))?; for (ix, gi) in gen_inputs.into_iter().enumerate() { trials.push(Trial::test( format!("{name}@{batch_size}/{ix}"), diff --git a/zero/src/bin/rpc.rs b/zero/src/bin/rpc.rs index d49cdde5c..164751df2 100644 --- a/zero/src/bin/rpc.rs +++ b/zero/src/bin/rpc.rs @@ -14,6 +14,7 @@ use url::Url; use zero::block_interval::BlockInterval; use zero::block_interval::BlockIntervalStream; use zero::prover::BlockProverInput; +use zero::prover::WIRE_DISPOSITION; use zero::provider::CachedProvider; use zero::rpc; @@ -172,6 +173,7 @@ impl Cli { block_prover_input.other_data, batch_size, &mut DummyObserver::new(), + WIRE_DISPOSITION, )?; if let Some(index) = tx_info.transaction_index { diff --git a/zero/src/bin/trie_diff.rs b/zero/src/bin/trie_diff.rs index 312feceb1..b0005b819 100644 --- a/zero/src/bin/trie_diff.rs +++ b/zero/src/bin/trie_diff.rs @@ -26,6 +26,7 @@ use regex::Regex; use trace_decoder::observer::TriesObserver; use tracing::{error, info}; use zero::ops::register; +use zero::prover::WIRE_DISPOSITION; use zero::prover::{cli::CliProverConfig, BlockProverInput, ProverConfig}; #[derive(Parser)] @@ -89,6 +90,7 @@ async fn main() -> Result<()> { block_prover_input.other_data.clone(), prover_config.batch_size, &mut observer, + WIRE_DISPOSITION, )?; info!( "Number of collected batch tries for block {}: {}", @@ -138,6 +140,9 @@ async fn main() -> Result<()> { state_trie: observer.data[prover_tries.batch_index] .tries .state + .as_ref() + .left() + .unwrap() .as_hashed_partial_trie() .clone(), transaction_trie: observer.data[prover_tries.batch_index] diff --git a/zero/src/prover.rs b/zero/src/prover.rs index e2207d993..80612afc3 100644 --- a/zero/src/prover.rs +++ b/zero/src/prover.rs @@ -39,12 +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"); - }) +pub 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)] From 8fc28d57697459c09bd19f828f53b5eaed10f305 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sat, 5 Oct 2024 17:41:52 +0100 Subject: [PATCH 04/13] refactor: TrieKey -> MptKey --- trace_decoder/src/core.rs | 48 +++++++++---------- trace_decoder/src/type1.rs | 10 ++-- trace_decoder/src/typed_mpt.rs | 86 +++++++++++++++++----------------- 3 files changed, 72 insertions(+), 72 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index b006f970d..f07a25fbc 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -22,7 +22,7 @@ use zk_evm_common::gwei_to_wei; use crate::{observer::Observer, typed_mpt::StateSmt}; use crate::{ - typed_mpt::{ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie, TrieKey}, + typed_mpt::{MptKey, ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie}, BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage, OtherBlockData, SeparateStorageTriesPreImage, SeparateTriePreImage, SeparateTriePreImages, TxnInfo, TxnMeta, TxnTrace, @@ -160,7 +160,7 @@ fn start( let state = state.items().try_fold( StateMpt::default(), |mut acc, (nibbles, hash_or_val)| { - let path = TrieKey::from_nibbles(nibbles); + let path = MptKey::from_nibbles(nibbles); match hash_or_val { mpt_trie::trie_ops::ValOrHash::Val(bytes) => { #[expect(deprecated)] // this is MPT specific @@ -183,7 +183,7 @@ fn start( .map(|(k, SeparateTriePreImage::Direct(v))| { v.items() .try_fold(StorageTrie::default(), |mut acc, (nibbles, hash_or_val)| { - let path = TrieKey::from_nibbles(nibbles); + let path = MptKey::from_nibbles(nibbles); match hash_or_val { mpt_trie::trie_ops::ValOrHash::Val(value) => { acc.insert(path, value)?; @@ -331,7 +331,7 @@ fn middle( for (haddr, acct) in state_trie.iter() { let storage = storage_tries.entry(haddr).or_insert({ let mut it = StorageTrie::default(); - it.insert_hash(TrieKey::default(), acct.storage_root) + it.insert_hash(MptKey::default(), acct.storage_root) .expect("empty trie insert cannot fail"); it }); @@ -366,7 +366,7 @@ fn middle( // We want to perform mask the TrieInputs above, // but won't know the bounds until after the loop below, // so store that information here. - let mut storage_masks = BTreeMap::<_, BTreeSet>::new(); + let mut storage_masks = BTreeMap::<_, BTreeSet>::new(); let mut state_mask = BTreeSet::new(); if txn_ix == 0 { @@ -455,7 +455,7 @@ fn middle( storage_written .keys() .chain(&storage_read) - .map(|it| TrieKey::from_hash(keccak_hash::keccak(it))), + .map(|it| MptKey::from_hash(keccak_hash::keccak(it))), ); if do_writes { @@ -486,7 +486,7 @@ fn middle( }; for (k, v) in storage_written { - let slot = TrieKey::from_hash(keccak_hash::keccak(k)); + let slot = MptKey::from_hash(keccak_hash::keccak(k)); match v.is_zero() { // this is actually a delete true => storage_mask.extend(storage.reporting_remove(slot)?), @@ -499,7 +499,7 @@ fn middle( } state_trie.insert_by_address(addr, acct)?; - state_mask.insert(TrieKey::from_address(addr)); + state_mask.insert(MptKey::from_address(addr)); } else { // Simple state access @@ -524,7 +524,7 @@ fn middle( // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/pull/613 // masking like this SHOULD be a space-saving optimization, // BUT if it's omitted, we actually get state root mismatches - state_mask.insert(TrieKey::from_address(addr)); + state_mask.insert(MptKey::from_address(addr)); } } @@ -548,7 +548,7 @@ fn middle( withdrawals: match loop_ix == loop_len { true => { for (addr, amt) in &withdrawals { - state_mask.insert(TrieKey::from_address(*addr)); + state_mask.insert(MptKey::from_address(*addr)); let mut acct = state_trie .get_by_address(*addr) .context("missing address for withdrawal")?; @@ -606,8 +606,8 @@ fn do_pre_execution( block: &BlockMetadata, ger_data: Option<(H256, H256)>, storage: &mut BTreeMap, - trim_storage: &mut BTreeMap>, - trim_state: &mut BTreeSet, + trim_storage: &mut BTreeMap>, + trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, ) -> anyhow::Result<()> { // Ethereum mainnet: EIP-4788 @@ -645,8 +645,8 @@ fn do_scalable_hook( block: &BlockMetadata, ger_data: Option<(H256, H256)>, storage: &mut BTreeMap, - trim_storage: &mut BTreeMap>, - trim_state: &mut BTreeSet, + trim_storage: &mut BTreeMap>, + trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, ) -> anyhow::Result<()> { use evm_arithmetization::testing_utils::{ @@ -663,7 +663,7 @@ fn do_scalable_hook( .context("missing scalable contract storage trie")?; let scalable_trim = trim_storage.entry(ADDRESS_SCALABLE_L2).or_default(); - let timestamp_slot_key = TrieKey::from_slot_position(U256::from(TIMESTAMP_STORAGE_POS.1)); + let timestamp_slot_key = MptKey::from_slot_position(U256::from(TIMESTAMP_STORAGE_POS.1)); let timestamp = scalable_storage .get(×tamp_slot_key) @@ -677,7 +677,7 @@ fn do_scalable_hook( (U256::from(LAST_BLOCK_STORAGE_POS.1), block.block_number), (U256::from(TIMESTAMP_STORAGE_POS.1), timestamp), ] { - let slot = TrieKey::from_slot_position(ix); + let slot = MptKey::from_slot_position(ix); // These values are never 0. scalable_storage.insert(slot, alloy::rlp::encode(u.compat()))?; @@ -690,12 +690,12 @@ fn do_scalable_hook( let mut arr = [0; 64]; (block.block_number - 1).to_big_endian(&mut arr[0..32]); U256::from(STATE_ROOT_STORAGE_POS.1).to_big_endian(&mut arr[32..64]); - let slot = TrieKey::from_hash(keccak_hash::keccak(arr)); + let slot = MptKey::from_hash(keccak_hash::keccak(arr)); scalable_storage.insert(slot, alloy::rlp::encode(prev_block_root_hash.compat()))?; scalable_trim.insert(slot); - trim_state.insert(TrieKey::from_address(ADDRESS_SCALABLE_L2)); + trim_state.insert(MptKey::from_address(ADDRESS_SCALABLE_L2)); let mut scalable_acct = state_trie .get_by_address(ADDRESS_SCALABLE_L2) .context("missing scalable contract address")?; @@ -716,12 +716,12 @@ fn do_scalable_hook( let mut arr = [0; 64]; arr[0..32].copy_from_slice(&root.0); U256::from(GLOBAL_EXIT_ROOT_STORAGE_POS.1).to_big_endian(&mut arr[32..64]); - let slot = TrieKey::from_hash(keccak_hash::keccak(arr)); + let slot = MptKey::from_hash(keccak_hash::keccak(arr)); ger_storage.insert(slot, alloy::rlp::encode(l1blockhash.compat()))?; ger_trim.insert(slot); - trim_state.insert(TrieKey::from_address(GLOBAL_EXIT_ROOT_ADDRESS)); + trim_state.insert(MptKey::from_address(GLOBAL_EXIT_ROOT_ADDRESS)); let mut ger_acct = state_trie .get_by_address(GLOBAL_EXIT_ROOT_ADDRESS) .context("missing GER contract address")?; @@ -744,9 +744,9 @@ fn do_scalable_hook( fn do_beacon_hook( block_timestamp: U256, storage: &mut BTreeMap, - trim_storage: &mut BTreeMap>, + trim_storage: &mut BTreeMap>, parent_beacon_block_root: H256, - trim_state: &mut BTreeSet, + trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, ) -> anyhow::Result<()> { use evm_arithmetization::testing_utils::{ @@ -769,7 +769,7 @@ fn do_beacon_hook( U256::from_big_endian(parent_beacon_block_root.as_bytes()), ), ] { - let slot = TrieKey::from_slot_position(ix); + let slot = MptKey::from_slot_position(ix); beacon_trim.insert(slot); match u.is_zero() { @@ -780,7 +780,7 @@ fn do_beacon_hook( } } } - trim_state.insert(TrieKey::from_address(BEACON_ROOTS_CONTRACT_ADDRESS)); + trim_state.insert(MptKey::from_address(BEACON_ROOTS_CONTRACT_ADDRESS)); let mut beacon_acct = state_trie .get_by_address(BEACON_ROOTS_CONTRACT_ADDRESS) .context("missing beacon contract address")?; diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index aeea0dbb6..cfa0ed615 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -12,7 +12,7 @@ use mpt_trie::partial_trie::OnOrphanedHashNode; use nunny::NonEmpty; use u4::U4; -use crate::typed_mpt::{StateMpt, StateTrie as _, StorageTrie, TrieKey}; +use crate::typed_mpt::{StateMpt, StateTrie as _, StorageTrie, MptKey}; use crate::wire::{Instruction, SmtLeaf}; #[derive(Debug, Clone)] @@ -66,10 +66,10 @@ fn visit( Node::Hash(Hash { raw_hash }) => { frontend .state - .insert_hash_by_key(TrieKey::new(path.iter().copied())?, raw_hash.into())?; + .insert_hash_by_key(MptKey::new(path.iter().copied())?, raw_hash.into())?; } Node::Leaf(Leaf { key, value }) => { - let path = TrieKey::new(path.iter().copied().chain(key))? + let path = MptKey::new(path.iter().copied().chain(key))? .into_hash() .context("invalid depth for leaf of state trie")?; match value { @@ -141,12 +141,12 @@ fn node2storagetrie(node: Node) -> anyhow::Result { ) -> anyhow::Result<()> { match node { Node::Hash(Hash { raw_hash }) => { - mpt.insert_hash(TrieKey::new(path.iter().copied())?, raw_hash.into())?; + mpt.insert_hash(MptKey::new(path.iter().copied())?, raw_hash.into())?; } Node::Leaf(Leaf { key, value }) => { match value { Either::Left(Value { raw_value }) => mpt.insert( - TrieKey::new(path.iter().copied().chain(key))?, + MptKey::new(path.iter().copied().chain(key))?, rlp::encode(&raw_value.as_slice()).to_vec(), )?, Either::Right(_) => bail!("unexpected account node in storage trie"), diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 9d29d508b..cbce06536 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -32,13 +32,13 @@ impl TypedMpt { /// Insert a node which represents an out-of-band sub-trie. /// /// See [module documentation](super) for more. - fn insert_hash(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> { + fn insert_hash(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { self.inner.insert(key.into_nibbles(), hash)?; Ok(()) } /// Returns an [`Error`] if the `key` crosses into a part of the trie that /// isn't hydrated. - fn insert(&mut self, key: TrieKey, value: T) -> anyhow::Result> + fn insert(&mut self, key: MptKey, value: T) -> anyhow::Result> where T: rlp::Encodable + rlp::Decodable, { @@ -52,7 +52,7 @@ impl TypedMpt { /// /// # Panics /// - If [`rlp::decode`]-ing for `T` doesn't round-trip. - fn get(&self, key: TrieKey) -> Option + fn get(&self, key: MptKey) -> Option where T: rlp::Decodable, { @@ -69,12 +69,12 @@ impl TypedMpt { self.inner.hash() } /// Note that this returns owned paths and items. - fn iter(&self) -> impl Iterator + '_ + fn iter(&self) -> impl Iterator + '_ where T: rlp::Decodable, { self.inner.keys().filter_map(|nib| { - let path = TrieKey::from_nibbles(nib); + let path = MptKey::from_nibbles(nib); Some((path, self.get(path)?)) }) } @@ -90,7 +90,7 @@ impl<'a, T> IntoIterator for &'a TypedMpt where T: rlp::Decodable, { - type Item = (TrieKey, T); + type Item = (MptKey, T); type IntoIter = Box + 'a>; fn into_iter(self) -> Self::IntoIter { Box::new(self.iter()) @@ -102,9 +102,9 @@ where /// /// Semantically equivalent to [`mpt_trie::nibbles::Nibbles`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] -pub struct TrieKey(CopyVec); +pub struct MptKey(CopyVec); -impl fmt::Display for TrieKey { +impl fmt::Display for MptKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { for u in self.0 { f.write_fmt(format_args!("{:x}", u))? @@ -113,9 +113,9 @@ impl fmt::Display for TrieKey { } } -impl TrieKey { +impl MptKey { pub fn new(components: impl IntoIterator) -> anyhow::Result { - Ok(TrieKey(CopyVec::try_from_iter(components)?)) + Ok(MptKey(CopyVec::try_from_iter(components)?)) } pub fn into_hash_left_padded(mut self) -> H256 { for _ in 0..self.0.spare_capacity_mut().len() { @@ -138,7 +138,7 @@ impl TrieKey { } pub fn from_txn_ix(txn_ix: usize) -> Self { - TrieKey::new(AsNibbles(rlp::encode(&txn_ix))).expect( + MptKey::new(AsNibbles(rlp::encode(&txn_ix))).expect( "\ rlp of an usize goes through a u64, which is 8 bytes, which will be 9 bytes RLP'ed. @@ -185,9 +185,9 @@ impl TrieKey { #[test] fn key_into_hash() { - assert_eq!(TrieKey::new([]).unwrap().into_hash(), None); + assert_eq!(MptKey::new([]).unwrap().into_hash(), None); assert_eq!( - TrieKey::new(itertools::repeat_n(u4::u4!(0), 64)) + MptKey::new(itertools::repeat_n(u4::u4!(0), 64)) .unwrap() .into_hash(), Some(H256::zero()) @@ -209,10 +209,10 @@ impl TransactionTrie { pub fn insert(&mut self, txn_ix: usize, val: Vec) -> anyhow::Result>> { let prev = self .untyped - .get(TrieKey::from_txn_ix(txn_ix).into_nibbles()) + .get(MptKey::from_txn_ix(txn_ix).into_nibbles()) .map(Vec::from); self.untyped - .insert(TrieKey::from_txn_ix(txn_ix).into_nibbles(), val)?; + .insert(MptKey::from_txn_ix(txn_ix).into_nibbles(), val)?; Ok(prev) } pub fn root(&self) -> H256 { @@ -227,7 +227,7 @@ impl TransactionTrie { &self.untyped, txn_ixs .into_iter() - .map(|it| TrieKey::from_txn_ix(it).into_nibbles()), + .map(|it| MptKey::from_txn_ix(it).into_nibbles()), )?; Ok(()) } @@ -254,10 +254,10 @@ impl ReceiptTrie { pub fn insert(&mut self, txn_ix: usize, val: Vec) -> anyhow::Result>> { let prev = self .untyped - .get(TrieKey::from_txn_ix(txn_ix).into_nibbles()) + .get(MptKey::from_txn_ix(txn_ix).into_nibbles()) .map(Vec::from); self.untyped - .insert(TrieKey::from_txn_ix(txn_ix).into_nibbles(), val)?; + .insert(MptKey::from_txn_ix(txn_ix).into_nibbles(), val)?; Ok(prev) } pub fn root(&self) -> H256 { @@ -272,7 +272,7 @@ impl ReceiptTrie { &self.untyped, txn_ixs .into_iter() - .map(|it| TrieKey::from_txn_ix(it).into_nibbles()), + .map(|it| MptKey::from_txn_ix(it).into_nibbles()), )?; Ok(()) } @@ -292,9 +292,9 @@ pub trait StateTrie { account: AccountRlp, ) -> anyhow::Result>; fn get_by_address(&self, address: Address) -> Option; - fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; + fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; /// _Hash out_ parts of the trie that aren't in `txn_ixs`. - fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; + fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; fn iter(&self) -> impl Iterator + '_; fn root(&self) -> H256; } @@ -317,7 +317,7 @@ impl StateMpt { } } /// Insert a _hashed out_ part of the trie - pub fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> { + pub fn insert_hash_by_key(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { self.typed.insert_hash(key, hash) } #[deprecated = "prefer operations on `Address` where possible, as SMT support requires this"] @@ -326,7 +326,7 @@ impl StateMpt { key: H256, account: AccountRlp, ) -> anyhow::Result> { - self.typed.insert(TrieKey::from_hash(key), account) + self.typed.insert(MptKey::from_hash(key), account) } pub fn iter(&self) -> impl Iterator + '_ { self.typed @@ -349,20 +349,20 @@ impl StateTrie for StateMpt { } fn get_by_address(&self, address: Address) -> Option { self.typed - .get(TrieKey::from_hash(keccak_hash::keccak(address))) + .get(MptKey::from_hash(keccak_hash::keccak(address))) } /// Delete the account at `address`, returning any remaining branch on /// collapse - fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { + fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { delete_node_and_report_remaining_key_if_branch_collapsed( self.typed.as_mut_hashed_partial_trie_unchecked(), - TrieKey::from_address(address), + MptKey::from_address(address), ) } - fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()> { + fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()> { let inner = mpt_trie::trie_subsets::create_trie_subset( self.typed.as_hashed_partial_trie(), - addresses.into_iter().map(TrieKey::into_nibbles), + addresses.into_iter().map(MptKey::into_nibbles), )?; self.typed = TypedMpt { inner, @@ -396,7 +396,7 @@ impl From for HashedPartialTrie { #[derive(Clone, Debug)] pub struct StateSmt { address2state: BTreeMap, - hashed_out: BTreeMap, + hashed_out: BTreeMap, } impl StateTrie for StateSmt { @@ -410,11 +410,11 @@ impl StateTrie for StateSmt { fn get_by_address(&self, address: Address) -> Option { self.address2state.get(&address).copied() } - fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { + fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { self.address2state.remove(&address); Ok(None) } - fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()> { + fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()> { let _ = address; Ok(()) } @@ -518,15 +518,15 @@ impl StorageTrie { untyped: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), } } - pub fn get(&mut self, key: &TrieKey) -> Option<&[u8]> { + pub fn get(&mut self, key: &MptKey) -> Option<&[u8]> { self.untyped.get(key.into_nibbles()) } - pub fn insert(&mut self, key: TrieKey, value: Vec) -> anyhow::Result>> { + pub fn insert(&mut self, key: MptKey, value: Vec) -> anyhow::Result>> { let prev = self.get(&key).map(Vec::from); self.untyped.insert(key.into_nibbles(), value)?; Ok(prev) } - pub fn insert_hash(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> { + pub fn insert_hash(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { self.untyped.insert(key.into_nibbles(), hash)?; Ok(()) } @@ -536,17 +536,17 @@ impl StorageTrie { pub const fn as_hashed_partial_trie(&self) -> &HashedPartialTrie { &self.untyped } - pub fn reporting_remove(&mut self, key: TrieKey) -> anyhow::Result> { + pub fn reporting_remove(&mut self, key: MptKey) -> anyhow::Result> { delete_node_and_report_remaining_key_if_branch_collapsed(&mut self.untyped, key) } pub fn as_mut_hashed_partial_trie_unchecked(&mut self) -> &mut HashedPartialTrie { &mut self.untyped } /// _Hash out_ the parts of the trie that aren't in `paths`. - pub fn mask(&mut self, paths: impl IntoIterator) -> anyhow::Result<()> { + pub fn mask(&mut self, paths: impl IntoIterator) -> anyhow::Result<()> { self.untyped = mpt_trie::trie_subsets::create_trie_subset( &self.untyped, - paths.into_iter().map(TrieKey::into_nibbles), + paths.into_iter().map(MptKey::into_nibbles), )?; Ok(()) } @@ -588,8 +588,8 @@ where account: AccountRlp, ) -> 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 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 + '_ { @@ -604,18 +604,18 @@ where /// plonky2. Returns the key to the remaining child if a collapse occurred. fn delete_node_and_report_remaining_key_if_branch_collapsed( trie: &mut HashedPartialTrie, - key: TrieKey, -) -> anyhow::Result> { + key: MptKey, +) -> anyhow::Result> { let old_trace = get_trie_trace(trie, key); trie.delete(key.into_nibbles())?; let new_trace = get_trie_trace(trie, key); Ok( node_deletion_resulted_in_a_branch_collapse(&old_trace, &new_trace) - .map(TrieKey::from_nibbles), + .map(MptKey::from_nibbles), ) } -fn get_trie_trace(trie: &HashedPartialTrie, k: TrieKey) -> mpt_trie::utils::TriePath { +fn get_trie_trace(trie: &HashedPartialTrie, k: MptKey) -> mpt_trie::utils::TriePath { mpt_trie::special_query::path_for_query(trie, k.into_nibbles(), true).collect() } From 98690ccee9509ddd47c30a0d1dae0a86ca6f4803 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sat, 5 Oct 2024 19:49:01 +0100 Subject: [PATCH 05/13] wip --- trace_decoder/src/core.rs | 125 +++++++++++++++++++++------- trace_decoder/src/typed_mpt.rs | 144 ++++++++++++++++++++++----------- 2 files changed, 194 insertions(+), 75 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index f07a25fbc..701ed43f1 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -20,7 +20,10 @@ use mpt_trie::partial_trie::PartialTrie as _; use nunny::NonEmpty; use zk_evm_common::gwei_to_wei; -use crate::{observer::Observer, typed_mpt::StateSmt}; +use crate::{ + observer::{DummyObserver, Observer}, + typed_mpt::StateSmt, +}; use crate::{ typed_mpt::{MptKey, ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie}, BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage, @@ -28,7 +31,7 @@ use crate::{ TxnInfo, TxnMeta, TxnTrace, }; -/// When parsing tries, which type to deserialize as. +/// When parsing tries from binary format, which type to deserialize as. #[derive(Debug)] pub enum WireDisposition { /// MPT @@ -42,7 +45,7 @@ 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); @@ -72,16 +75,32 @@ pub fn entrypoint( *amt = gwei_to_wei(*amt) } - let batches = middle( - state, - storage, - batch(txn_info, batch_size_hint), - &mut code, - &b_meta, - ger_data, - withdrawals, - observer, - )?; + match state { + Either::Left(mpt) => { + let batches = middle( + mpt, + storage, + batch(txn_info, batch_size_hint), + &mut code, + &b_meta, + ger_data, + withdrawals, + observer, + )?; + } + Either::Right(smt) => { + let batches = middle( + smt, + storage, + batch(txn_info, batch_size_hint), + &mut code, + &b_meta, + ger_data, + withdrawals, + &mut DummyObserver::new(), // TODO(0xaatif) + )?; + } + } let mut running_gas_used = 0; Ok(batches @@ -299,6 +318,29 @@ struct Batch { pub withdrawals: Vec<(Address, U256)>, } +impl Batch { + fn map(self, f: impl FnMut(T) -> U) -> Batch { + let Self { + first_txn_ix, + gas_used, + contract_code, + byte_code, + before, + after, + withdrawals, + } = self; + Batch { + first_txn_ix, + gas_used, + contract_code, + byte_code, + before: before.map(f), + after, + withdrawals, + } + } +} + /// [`evm_arithmetization::generation::TrieInputs`], /// generic over state trie representation. #[derive(Debug)] @@ -309,6 +351,23 @@ pub struct IntraBlockTries { pub receipt: ReceiptTrie, } +impl IntraBlockTries { + fn map(self, mut f: impl FnMut(T) -> U) -> IntraBlockTries { + let Self { + state, + storage, + transaction, + receipt, + } = self; + IntraBlockTries { + state: f(state), + storage, + transaction, + receipt, + } + } +} + /// Does the main work mentioned in the [module documentation](super). #[allow(clippy::too_many_arguments)] fn middle( @@ -326,7 +385,10 @@ fn middle( mut withdrawals: Vec<(Address, U256)>, // called with the untrimmed tries after each batch observer: &mut impl Observer, -) -> anyhow::Result>> { +) -> anyhow::Result>> +where + StateTrieT::Key: Ord + From
, +{ // Initialise the storage tries. for (haddr, acct) in state_trie.iter() { let storage = storage_tries.entry(haddr).or_insert({ @@ -367,7 +429,7 @@ fn middle( // but won't know the bounds until after the loop below, // so store that information here. let mut storage_masks = BTreeMap::<_, BTreeSet>::new(); - let mut state_mask = BTreeSet::new(); + let mut state_mask = BTreeSet::::new(); if txn_ix == 0 { do_pre_execution( @@ -499,7 +561,7 @@ fn middle( } state_trie.insert_by_address(addr, acct)?; - state_mask.insert(MptKey::from_address(addr)); + state_mask.insert(::from(addr)); } else { // Simple state access @@ -524,7 +586,7 @@ fn middle( // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/pull/613 // masking like this SHOULD be a space-saving optimization, // BUT if it's omitted, we actually get state root mismatches - state_mask.insert(MptKey::from_address(addr)); + state_mask.insert(::from(addr)); } } @@ -548,7 +610,7 @@ fn middle( withdrawals: match loop_ix == loop_len { true => { for (addr, amt) in &withdrawals { - state_mask.insert(MptKey::from_address(*addr)); + state_mask.insert(::from(*addr)); let mut acct = state_trie .get_by_address(*addr) .context("missing address for withdrawal")?; @@ -607,9 +669,12 @@ fn do_pre_execution( ger_data: Option<(H256, H256)>, storage: &mut BTreeMap, trim_storage: &mut BTreeMap>, - trim_state: &mut BTreeSet, + trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, -) -> anyhow::Result<()> { +) -> anyhow::Result<()> +where + StateTrieT::Key: From
+ Ord, +{ // Ethereum mainnet: EIP-4788 if cfg!(feature = "eth_mainnet") { return do_beacon_hook( @@ -646,9 +711,12 @@ fn do_scalable_hook( ger_data: Option<(H256, H256)>, storage: &mut BTreeMap, trim_storage: &mut BTreeMap>, - trim_state: &mut BTreeSet, + trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, -) -> anyhow::Result<()> { +) -> anyhow::Result<()> +where + StateTrieT::Key: From
+ Ord, +{ use evm_arithmetization::testing_utils::{ ADDRESS_SCALABLE_L2, ADDRESS_SCALABLE_L2_ADDRESS_HASHED, GLOBAL_EXIT_ROOT_ADDRESS, GLOBAL_EXIT_ROOT_ADDRESS_HASHED, GLOBAL_EXIT_ROOT_STORAGE_POS, LAST_BLOCK_STORAGE_POS, @@ -695,7 +763,7 @@ fn do_scalable_hook( scalable_storage.insert(slot, alloy::rlp::encode(prev_block_root_hash.compat()))?; scalable_trim.insert(slot); - trim_state.insert(MptKey::from_address(ADDRESS_SCALABLE_L2)); + trim_state.insert(::from(ADDRESS_SCALABLE_L2)); let mut scalable_acct = state_trie .get_by_address(ADDRESS_SCALABLE_L2) .context("missing scalable contract address")?; @@ -721,7 +789,7 @@ fn do_scalable_hook( ger_storage.insert(slot, alloy::rlp::encode(l1blockhash.compat()))?; ger_trim.insert(slot); - trim_state.insert(MptKey::from_address(GLOBAL_EXIT_ROOT_ADDRESS)); + trim_state.insert(::from(GLOBAL_EXIT_ROOT_ADDRESS)); let mut ger_acct = state_trie .get_by_address(GLOBAL_EXIT_ROOT_ADDRESS) .context("missing GER contract address")?; @@ -746,9 +814,12 @@ fn do_beacon_hook( storage: &mut BTreeMap, trim_storage: &mut BTreeMap>, parent_beacon_block_root: H256, - trim_state: &mut BTreeSet, + trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, -) -> anyhow::Result<()> { +) -> anyhow::Result<()> +where + StateTrieT::Key: From
+ Ord, +{ use evm_arithmetization::testing_utils::{ BEACON_ROOTS_CONTRACT_ADDRESS, BEACON_ROOTS_CONTRACT_ADDRESS_HASHED, HISTORY_BUFFER_LENGTH, }; @@ -780,7 +851,7 @@ fn do_beacon_hook( } } } - trim_state.insert(MptKey::from_address(BEACON_ROOTS_CONTRACT_ADDRESS)); + trim_state.insert(::from(BEACON_ROOTS_CONTRACT_ADDRESS)); let mut beacon_acct = state_trie .get_by_address(BEACON_ROOTS_CONTRACT_ADDRESS) .context("missing beacon contract address")?; diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index cbce06536..6a24a6d4c 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -1,9 +1,10 @@ //! Principled MPT types used in this library. use core::fmt; -use std::{collections::BTreeMap, marker::PhantomData}; +use std::{cmp, collections::BTreeMap, marker::PhantomData}; -use bitvec::{order::Msb0, view::BitView as _}; +use anyhow::ensure; +use bitvec::{order::Msb0, slice::BitSlice, view::BitView as _}; use copyvec::CopyVec; use either::Either; use ethereum_types::{Address, BigEndianHash as _, H256, U256}; @@ -183,8 +184,14 @@ impl MptKey { } } +impl From
for MptKey { + fn from(value: Address) -> Self { + Self::from_hash(keccak_hash::keccak(value)) + } +} + #[test] -fn key_into_hash() { +fn mpt_key_into_hash() { assert_eq!(MptKey::new([]).unwrap().into_hash(), None); assert_eq!( MptKey::new(itertools::repeat_n(u4::u4!(0), 64)) @@ -194,6 +201,85 @@ fn key_into_hash() { ) } +/// Bounded sequence of bits, +/// used as a key for [`StateSmt`]. +/// +/// Semantically equivalent to +#[derive(Clone, Copy)] +pub struct SmtKey { + bits: bitvec::array::BitArray<[u8; 32]>, + len: usize, +} + +impl SmtKey { + fn as_bitslice(&self) -> &BitSlice { + self.bits.as_bitslice().get(..self.len).unwrap() + } +} + +impl fmt::Debug for SmtKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list() + .entries(self.as_bitslice().iter().map(|it| match *it { + true => 1, + false => 0, + })) + .finish() + } +} + +impl fmt::Display for SmtKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for bit in self.as_bitslice() { + f.write_str(match *bit { + true => "1", + false => "0", + })? + } + Ok(()) + } +} + +impl SmtKey { + pub fn new(components: impl IntoIterator) -> anyhow::Result { + let mut bits = bitvec::array::BitArray::default(); + let mut len = 0; + for (ix, bit) in components.into_iter().enumerate() { + ensure!( + bits.get(ix).is_some(), + "expected at most {} components", + bits.len() + ); + bits.set(ix, bit); + len += 1 + } + Ok(Self { bits, len }) + } +} + +impl From
for SmtKey { + fn from(value: Address) -> Self { + todo!() + } +} + +impl Ord for SmtKey { + fn cmp(&self, other: &Self) -> cmp::Ordering { + self.as_bitslice().cmp(other.as_bitslice()) + } +} +impl PartialOrd for SmtKey { + fn partial_cmp(&self, other: &Self) -> Option { + self.as_bitslice().partial_cmp(other.as_bitslice()) + } +} +impl Eq for SmtKey {} +impl PartialEq for SmtKey { + fn eq(&self, other: &Self) -> bool { + self.as_bitslice().eq(other.as_bitslice()) + } +} + /// Per-block, `txn_ix -> [u8]`. /// /// See @@ -286,15 +372,16 @@ impl From for HashedPartialTrie { /// TODO(0xaatif): document this after refactoring is done https://github.com/0xPolygonZero/zk_evm/issues/275 pub trait StateTrie { + type Key; fn insert_by_address( &mut self, address: Address, account: AccountRlp, ) -> anyhow::Result>; fn get_by_address(&self, address: Address) -> Option; - fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; + fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; /// _Hash out_ parts of the trie that aren't in `txn_ixs`. - fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; + fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; fn iter(&self) -> impl Iterator + '_; fn root(&self) -> H256; } @@ -339,6 +426,7 @@ impl StateMpt { } impl StateTrie for StateMpt { + type Key = MptKey; fn insert_by_address( &mut self, address: Address, @@ -400,6 +488,7 @@ pub struct StateSmt { } impl StateTrie for StateSmt { + type Key = SmtKey; fn insert_by_address( &mut self, address: Address, @@ -410,11 +499,11 @@ impl StateTrie for StateSmt { fn get_by_address(&self, address: Address) -> Option { self.address2state.get(&address).copied() } - fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { + fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { self.address2state.remove(&address); Ok(None) } - fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()> { + fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()> { let _ = address; Ok(()) } @@ -558,47 +647,6 @@ 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 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. From d5d2aedf4f27cedf15e9fcc15a3644ad0fca41ad Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 6 Oct 2024 00:16:28 +0100 Subject: [PATCH 06/13] getting happier --- trace_decoder/src/core.rs | 38 ++++++++++++++++++++-------------- trace_decoder/src/typed_mpt.rs | 33 ++++++++++++++--------------- 2 files changed, 37 insertions(+), 34 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 701ed43f1..3fbed8189 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -75,9 +75,9 @@ pub fn entrypoint( *amt = gwei_to_wei(*amt) } - match state { - Either::Left(mpt) => { - let batches = middle( + let batches = match state { + Either::Left(mpt) => Either::Left( + middle( mpt, storage, batch(txn_info, batch_size_hint), @@ -86,21 +86,27 @@ pub fn entrypoint( ger_data, withdrawals, observer, - )?; - } + )? + .into_iter() + .map(|it| it.map(Either::Left)), + ), Either::Right(smt) => { - let batches = middle( - smt, - storage, - batch(txn_info, batch_size_hint), - &mut code, - &b_meta, - ger_data, - withdrawals, - &mut DummyObserver::new(), // TODO(0xaatif) - )?; + Either::Right( + middle( + smt, + storage, + batch(txn_info, batch_size_hint), + &mut code, + &b_meta, + ger_data, + withdrawals, + &mut DummyObserver::new(), // TODO(0xaatif) + )? + .into_iter() + .map(|it| it.map(Either::Right)), + ) } - } + }; let mut running_gas_used = 0; Ok(batches diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 6a24a6d4c..4df34891a 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -4,9 +4,8 @@ use core::fmt; use std::{cmp, collections::BTreeMap, marker::PhantomData}; use anyhow::ensure; -use bitvec::{order::Msb0, slice::BitSlice, view::BitView as _}; +use bitvec::{array::BitArray, slice::BitSlice}; use copyvec::CopyVec; -use either::Either; use ethereum_types::{Address, BigEndianHash as _, H256, U256}; use evm_arithmetization::generation::mpt::AccountRlp; use mpt_trie::partial_trie::{HashedPartialTrie, Node, OnOrphanedHashNode, PartialTrie as _}; @@ -164,17 +163,6 @@ impl MptKey { } Self(ours) } - fn into_bits(self) -> smt_trie::bits::Bits { - let mut bits = smt_trie::bits::Bits::default(); - for component in self.0 { - let byte = component as u8; - // the four high bits are zero - for bit in byte.view_bits::().into_iter().by_vals().skip(4) { - bits.push_bit(bit); - } - } - bits - } pub fn into_hash(self) -> Option { let Self(nibbles) = self; @@ -255,11 +243,20 @@ impl SmtKey { } Ok(Self { bits, len }) } + + fn into_bits(self) -> smt_trie::bits::Bits { + let mut bits = smt_trie::bits::Bits::default(); + for bit in self.as_bitslice() { + bits.push_bit(*bit) + } + bits + } } impl From
for SmtKey { - fn from(value: Address) -> Self { - todo!() + fn from(addr: Address) -> Self { + let H256(bytes) = keccak_hash::keccak(addr); + Self::new(BitArray::<_>::new(bytes)).unwrap() } } @@ -270,7 +267,7 @@ impl Ord for SmtKey { } impl PartialOrd for SmtKey { fn partial_cmp(&self, other: &Self) -> Option { - self.as_bitslice().partial_cmp(other.as_bitslice()) + Some(self.cmp(other)) } } impl Eq for SmtKey {} @@ -380,7 +377,7 @@ pub trait StateTrie { ) -> anyhow::Result>; fn get_by_address(&self, address: Address) -> Option; fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; - /// _Hash out_ parts of the trie that aren't in `txn_ixs`. + /// _Hash out_ parts of the trie that aren't in `addresses`. fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; fn iter(&self) -> impl Iterator + '_; fn root(&self) -> H256; @@ -484,7 +481,7 @@ impl From for HashedPartialTrie { #[derive(Clone, Debug)] pub struct StateSmt { address2state: BTreeMap, - hashed_out: BTreeMap, + hashed_out: BTreeMap, } impl StateTrie for StateSmt { From f27bbdf2fac784439a2db0801e62128fecb23ba9 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 6 Oct 2024 01:58:00 +0100 Subject: [PATCH 07/13] rewrite type2 --- trace_decoder/src/core.rs | 7 +- trace_decoder/src/type1.rs | 4 +- trace_decoder/src/type2.rs | 232 ++++++++++++++++----------------- trace_decoder/src/typed_mpt.rs | 4 +- trace_decoder/src/wire.rs | 6 +- 5 files changed, 122 insertions(+), 131 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 3fbed8189..181fa054f 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -241,11 +241,8 @@ fn start( ) } WireDisposition::Type2 => { - let crate::type2::Frontend { - trie, - code, - collation, - } = crate::type2::frontend(instructions)?; + let crate::type2::Frontend { trie, code } = + crate::type2::frontend(instructions)?; todo!() } diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index cfa0ed615..bce0c9134 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -12,7 +12,7 @@ use mpt_trie::partial_trie::OnOrphanedHashNode; use nunny::NonEmpty; use u4::U4; -use crate::typed_mpt::{StateMpt, StateTrie as _, StorageTrie, MptKey}; +use crate::typed_mpt::{MptKey, StateMpt, StorageTrie}; use crate::wire::{Instruction, SmtLeaf}; #[derive(Debug, Clone)] @@ -380,6 +380,8 @@ fn finish_stack(v: &mut Vec) -> anyhow::Result { #[test] fn test_tries() { + use crate::typed_mpt::StateTrie as _; + for (ix, case) in serde_json::from_str::>(include_str!("cases/zero_jerigon.json")) .unwrap() diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index eb70978f9..5d88be6b4 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -1,37 +1,36 @@ //! Frontend for the witness format emitted by e.g [`0xPolygonHermez/cdk-erigon`](https://github.com/0xPolygonHermez/cdk-erigon/) //! Ethereum node. -use std::{ - collections::{HashMap, HashSet}, - iter, -}; +use std::collections::{BTreeMap, HashSet}; use anyhow::{bail, ensure, Context as _}; -use bitvec::vec::BitVec; -use either::Either; -use ethereum_types::BigEndianHash as _; -use itertools::{EitherOrBoth, Itertools as _}; +use ethereum_types::{Address, BigEndianHash as _, U256}; +use itertools::EitherOrBoth; +use keccak_hash::H256; use nunny::NonEmpty; -use plonky2::field::types::Field; +use plonky2::field::types::{Field, Field64 as _}; +use smt_trie::keys::{key_balance, key_code, key_code_length, key_nonce, key_storage}; +use stackstack::Stack; use crate::{ - typed_mpt::StateSmt, + typed_mpt::SmtKey, wire::{Instruction, SmtLeaf, SmtLeafType}, }; type SmtTrie = smt_trie::smt::Smt; -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] +/// Combination of all the [`SmtLeaf::node_type`]s +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] pub struct CollatedLeaf { pub balance: Option, pub nonce: Option, - pub code_hash: Option, - pub storage_root: Option, + pub code: Option, + pub code_length: Option, + pub storage: BTreeMap, } pub struct Frontend { pub trie: SmtTrie, pub code: HashSet>>, - pub collation: HashMap, } /// # Panics @@ -39,13 +38,8 @@ pub struct Frontend { /// NOT call this function on untrusted inputs. pub fn frontend(instructions: impl IntoIterator) -> anyhow::Result { let (node, code) = fold(instructions).context("couldn't fold smt from instructions")?; - let (trie, collation) = - node2trie(node).context("couldn't construct trie and collation from folded node")?; - Ok(Frontend { - trie, - code, - collation, - }) + let trie = node2trie(node).context("couldn't construct trie and collation from folded node")?; + Ok(Frontend { trie, code }) } /// Node in a binary (SMT) tree. @@ -107,9 +101,9 @@ fn fold1(instructions: impl IntoIterator) -> anyhow::Result< Ok(Some(match mask { // note that the single-child bits are reversed... - 0b0001 => Node::Branch(EitherOrBoth::Left(get_child()?)), - 0b0010 => Node::Branch(EitherOrBoth::Right(get_child()?)), - 0b0011 => Node::Branch(EitherOrBoth::Both(get_child()?, get_child()?)), + 0b_01 => Node::Branch(EitherOrBoth::Left(get_child()?)), + 0b_10 => Node::Branch(EitherOrBoth::Right(get_child()?)), + 0b_11 => Node::Branch(EitherOrBoth::Both(get_child()?, get_child()?)), other => bail!("unexpected bit pattern in Branch mask: {:#b}", other), })) } @@ -121,113 +115,111 @@ fn fold1(instructions: impl IntoIterator) -> anyhow::Result< } } -/// Pack a [`Node`] tree into an [`SmtTrie`]. -/// Also summarizes the [`Node::Leaf`]s out-of-band. -/// -/// # Panics -/// - if the tree is too deep. -/// - if [`SmtLeaf::address`] or [`SmtLeaf::value`] are the wrong length. -/// - if [`SmtLeafType::Storage`] is the wrong length. -/// - [`SmtTrie`] panics internally. -fn node2trie( - node: Node, -) -> anyhow::Result<(SmtTrie, HashMap)> { +fn node2trie(node: Node) -> anyhow::Result { let mut trie = SmtTrie::default(); - - let (hashes, leaves) = - iter_leaves(node).partition_map::, Vec<_>, _, _, _>(|(path, leaf)| match leaf { - Either::Left(it) => Either::Left((path, it)), - Either::Right(it) => Either::Right(it), - }); - - let mut lens = std::collections::BTreeMap::<_, usize>::new(); - - for (path, hash) in hashes { - *lens.entry(path.len()).or_default() += 1; - // needs to be called before `set`, below, "to avoid any issues" according - // to the smt docs. + let mut hashes = BTreeMap::new(); + let mut leaves = BTreeMap::new(); + visit(&mut hashes, &mut leaves, Stack::new(), node)?; + for (key, hash) in hashes { trie.set_hash( - bits2bits(path), + key.into_smt_bits(), smt_trie::smt::HashOut { elements: { - let ethereum_types::U256(arr) = ethereum_types::H256(hash).into_uint(); + let ethereum_types::U256(arr) = hash.into_uint(); + for u in arr { + ensure!(u < smt_trie::smt::F::ORDER); + } arr.map(smt_trie::smt::F::from_canonical_u64) }, }, - ) + ); } - dbg!(lens); - - let mut collated = HashMap::::new(); - for SmtLeaf { - node_type, - address, - value, - } in leaves + for ( + addr, + CollatedLeaf { + balance, + nonce, + code, + code_length, + storage, + }, + ) in leaves { - let address = ethereum_types::Address::from_slice(&address); - let collated = collated.entry(address).or_default(); - let value = ethereum_types::U256::from_big_endian(&value); - let key = match node_type { - SmtLeafType::Balance => { - ensure!(collated.balance.is_none(), "double write of field"); - collated.balance = Some(value); - smt_trie::keys::key_balance(address) - } - SmtLeafType::Nonce => { - ensure!(collated.nonce.is_none(), "double write of field"); - collated.nonce = Some(value); - smt_trie::keys::key_nonce(address) - } - SmtLeafType::Code => { - ensure!(collated.code_hash.is_none(), "double write of field"); - collated.code_hash = Some({ - let mut it = ethereum_types::H256::zero(); - value.to_big_endian(it.as_bytes_mut()); - it - }); - smt_trie::keys::key_code(address) - } - SmtLeafType::Storage(it) => { - ensure!(collated.storage_root.is_none(), "double write of field"); - // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 - // do we not do anything with the storage here? - smt_trie::keys::key_storage(address, ethereum_types::U256::from_big_endian(&it)) + for (value, key_fn) in [ + (balance, key_balance as fn(_) -> _), + (nonce, key_nonce), + (code, key_code), + (code_length, key_code_length), + ] { + if let Some(value) = value { + trie.set(key_fn(addr), value); } - SmtLeafType::CodeLength => smt_trie::keys::key_code_length(address), - }; - trie.set(key, value) - } - Ok((trie, collated)) -} - -/// # Panics -/// - on overcapacity -fn bits2bits(ours: BitVec) -> smt_trie::bits::Bits { - let mut theirs = smt_trie::bits::Bits::empty(); - for it in ours { - theirs.push_bit(it) + } + for (slot, value) in storage { + trie.set(key_storage(addr, slot), value); + } } - theirs + Ok(trie) } -/// Simple, inefficient visitor of all leaves of the [`Node`] tree. -#[allow(clippy::type_complexity)] -fn iter_leaves(node: Node) -> Box)>> { +fn visit( + hashes: &mut BTreeMap, + leaves: &mut BTreeMap, + path: Stack, + node: Node, +) -> anyhow::Result<()> { match node { - Node::Hash(it) => Box::new(iter::once((BitVec::new(), Either::Left(it)))), - Node::Branch(it) => { - let (left, right) = it.left_and_right(); - let left = left - .into_iter() - .flat_map(|it| iter_leaves(*it).update(|(path, _)| path.insert(0, false))); - let right = right - .into_iter() - .flat_map(|it| iter_leaves(*it).update(|(path, _)| path.insert(0, true))); - Box::new(left.chain(right)) + Node::Branch(children) => { + let (left, right) = children.left_and_right(); + if let Some(left) = left { + visit(hashes, leaves, path.pushed(false), *left)?; + } + if let Some(right) = right { + visit(hashes, leaves, path.pushed(true), *right)?; + } + } + Node::Hash(hash) => { + hashes.insert(SmtKey::new(path.iter().copied())?, H256(hash)); + } + Node::Leaf(SmtLeaf { + node_type, + address, // TODO(0xaatif): field should be fixed length + value, // TODO(0xaatif): field should be fixed length + }) => { + let address = Address::from_slice(&address); + let collated = leaves.entry(address).or_default(); + let value = U256::from_big_endian(&value); + macro_rules! ensure { + ($expr:expr) => { + ::anyhow::ensure!($expr, "double write of field for address {}", address) + }; + } + match node_type { + SmtLeafType::Balance => { + ensure!(collated.balance.is_none()); + collated.balance = Some(value) + } + SmtLeafType::Nonce => { + ensure!(collated.nonce.is_none()); + collated.nonce = Some(value) + } + SmtLeafType::Code => { + ensure!(collated.code.is_none()); + collated.code = Some(value) + } + SmtLeafType::Storage(slot) => { + // TODO(0xaatif): ^ field should be fixed length + let clobbered = collated.storage.insert(U256::from_big_endian(&slot), value); + ensure!(clobbered.is_none()) + } + SmtLeafType::CodeLength => { + ensure!(collated.code_length.is_none()); + collated.code_length = Some(value) + } + }; } - Node::Leaf(it) => Box::new(iter::once((BitVec::new(), Either::Right(it)))), } + Ok(()) } #[test] @@ -241,10 +233,10 @@ fn test_tries() { println!("case {}", ix); let instructions = crate::wire::parse(&case.bytes).unwrap(); let frontend = frontend(instructions).unwrap(); - // assert_eq!(case.expected_state_root, { - // let mut it = [0; 32]; - // smt_trie::utils::hashout2u(frontend.trie.root).to_big_endian(&mut - // it); ethereum_types::H256(it) - // }); + assert_eq!(case.expected_state_root, { + let mut it = [0; 32]; + smt_trie::utils::hashout2u(frontend.trie.root).to_big_endian(&mut it); + ethereum_types::H256(it) + }); } } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 4df34891a..fa14f2ffc 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -244,7 +244,7 @@ impl SmtKey { Ok(Self { bits, len }) } - fn into_bits(self) -> smt_trie::bits::Bits { + pub fn into_smt_bits(self) -> smt_trie::bits::Bits { let mut bits = smt_trie::bits::Bits::default(); for bit in self.as_bitslice() { bits.push_bit(*bit) @@ -522,7 +522,7 @@ impl StateSmt { } = self; let mut smt = smt_trie::smt::Smt::::default(); for (k, v) in hashed_out { - smt.set_hash(k.into_bits(), conv_hash::eth2smt(*v)); + smt.set_hash(k.into_smt_bits(), conv_hash::eth2smt(*v)); } for ( addr, diff --git a/trace_decoder/src/wire.rs b/trace_decoder/src/wire.rs index 6f56f1e44..9d1a7fb10 100644 --- a/trace_decoder/src/wire.rs +++ b/trace_decoder/src/wire.rs @@ -82,8 +82,8 @@ pub enum Instruction { #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct SmtLeaf { pub node_type: SmtLeafType, - pub address: NonEmpty>, - pub value: NonEmpty>, + pub address: NonEmpty>, // TODO(0xaatif): this should be a fixed length + pub value: NonEmpty>, // TODO(0xaatif): this should be a fixed length } #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -91,7 +91,7 @@ pub enum SmtLeafType { Balance, Nonce, Code, - Storage(NonEmpty>), + Storage(NonEmpty>), // TODO(0xaatif): this should be a fixed length CodeLength, } From 29aaa1eca6935ca83e46689ebeda69d6a2cad37e Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 6 Oct 2024 02:13:45 +0100 Subject: [PATCH 08/13] refactor: insert_by_address -> () --- trace_decoder/src/type1.rs | 3 +-- trace_decoder/src/typed_mpt.rs | 28 ++++++++-------------------- 2 files changed, 9 insertions(+), 22 deletions(-) diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index bce0c9134..96f49d94d 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -106,8 +106,7 @@ fn visit( }, }; #[expect(deprecated)] // this is MPT-specific code - let clobbered = frontend.state.insert_by_hashed_address(path, account)?; - ensure!(clobbered.is_none(), "duplicate account"); + frontend.state.insert_by_hashed_address(path, account)?; } } } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index fa14f2ffc..a40d39db8 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -38,14 +38,13 @@ impl TypedMpt { } /// Returns an [`Error`] if the `key` crosses into a part of the trie that /// isn't hydrated. - fn insert(&mut self, key: MptKey, value: T) -> anyhow::Result> + fn insert(&mut self, key: MptKey, value: T) -> anyhow::Result<()> where T: rlp::Encodable + rlp::Decodable, { - let prev = self.get(key); self.inner .insert(key.into_nibbles(), rlp::encode(&value).to_vec())?; - Ok(prev) + Ok(()) } /// Note that this returns [`None`] if `key` crosses into a part of the /// trie that isn't hydrated. @@ -370,11 +369,7 @@ impl From for HashedPartialTrie { /// TODO(0xaatif): document this after refactoring is done https://github.com/0xPolygonZero/zk_evm/issues/275 pub trait StateTrie { type Key; - fn insert_by_address( - &mut self, - address: Address, - account: AccountRlp, - ) -> anyhow::Result>; + fn insert_by_address(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()>; fn get_by_address(&self, address: Address) -> Option; fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; /// _Hash out_ parts of the trie that aren't in `addresses`. @@ -409,7 +404,7 @@ impl StateMpt { &mut self, key: H256, account: AccountRlp, - ) -> anyhow::Result> { + ) -> anyhow::Result<()> { self.typed.insert(MptKey::from_hash(key), account) } pub fn iter(&self) -> impl Iterator + '_ { @@ -424,11 +419,7 @@ impl StateMpt { impl StateTrie for StateMpt { type Key = MptKey; - fn insert_by_address( - &mut self, - address: Address, - account: AccountRlp, - ) -> anyhow::Result> { + fn insert_by_address(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { #[expect(deprecated)] self.insert_by_hashed_address(keccak_hash::keccak(address), account) } @@ -486,12 +477,9 @@ pub struct StateSmt { impl StateTrie for StateSmt { type Key = SmtKey; - fn insert_by_address( - &mut self, - address: Address, - account: AccountRlp, - ) -> anyhow::Result> { - Ok(self.address2state.insert(address, account)) + fn insert_by_address(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { + self.address2state.insert(address, account); + Ok(()) } fn get_by_address(&self, address: Address) -> Option { self.address2state.get(&address).copied() From 847c37dde8a30ea9d5b14847bbdf13d43da4179e Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 7 Oct 2024 14:47:57 +0100 Subject: [PATCH 09/13] progress --- trace_decoder/src/core.rs | 7 +- trace_decoder/src/type2.rs | 155 ++++++++++++++++++++++----------- trace_decoder/src/typed_mpt.rs | 21 ++++- trace_decoder/src/wire.rs | 10 ++- zero/src/bin/trie_diff.rs | 3 - zero/src/prover.rs | 2 + 6 files changed, 135 insertions(+), 63 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 181fa054f..cbc45b6e6 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -243,8 +243,11 @@ fn start( WireDisposition::Type2 => { let crate::type2::Frontend { trie, code } = crate::type2::frontend(instructions)?; - - todo!() + ( + Either::Right(trie), + BTreeMap::new(), + Hash2Code::from_iter(code.into_iter().map(NonEmpty::into_vec)), + ) } }; (state, storage, code) diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index 5d88be6b4..7cb952888 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -4,19 +4,17 @@ use std::collections::{BTreeMap, HashSet}; use anyhow::{bail, ensure, Context as _}; -use ethereum_types::{Address, BigEndianHash as _, U256}; +use ethereum_types::{Address, U256}; +use evm_arithmetization::generation::mpt::AccountRlp; use itertools::EitherOrBoth; use keccak_hash::H256; use nunny::NonEmpty; -use plonky2::field::types::{Field, Field64 as _}; -use smt_trie::keys::{key_balance, key_code, key_code_length, key_nonce, key_storage}; use stackstack::Stack; use crate::{ - typed_mpt::SmtKey, + typed_mpt::{SmtKey, StateSmt}, wire::{Instruction, SmtLeaf, SmtLeafType}, }; -type SmtTrie = smt_trie::smt::Smt; /// Combination of all the [`SmtLeaf::node_type`]s #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -29,7 +27,7 @@ pub struct CollatedLeaf { } pub struct Frontend { - pub trie: SmtTrie, + pub trie: StateSmt, pub code: HashSet>>, } @@ -115,51 +113,43 @@ fn fold1(instructions: impl IntoIterator) -> anyhow::Result< } } -fn node2trie(node: Node) -> anyhow::Result { - let mut trie = SmtTrie::default(); +fn node2trie(node: Node) -> anyhow::Result { let mut hashes = BTreeMap::new(); let mut leaves = BTreeMap::new(); visit(&mut hashes, &mut leaves, Stack::new(), node)?; - for (key, hash) in hashes { - trie.set_hash( - key.into_smt_bits(), - smt_trie::smt::HashOut { - elements: { - let ethereum_types::U256(arr) = hash.into_uint(); - for u in arr { - ensure!(u < smt_trie::smt::F::ORDER); - } - arr.map(smt_trie::smt::F::from_canonical_u64) - }, - }, - ); - } - for ( - addr, - CollatedLeaf { - balance, - nonce, - code, - code_length, - storage, - }, - ) in leaves - { - for (value, key_fn) in [ - (balance, key_balance as fn(_) -> _), - (nonce, key_nonce), - (code, key_code), - (code_length, key_code_length), - ] { - if let Some(value) = value { - trie.set(key_fn(addr), value); - } - } - for (slot, value) in storage { - trie.set(key_storage(addr, slot), value); - } - } - Ok(trie) + Ok( + #[expect(deprecated, reason = "this is the frontend")] + StateSmt::new_unchecked( + leaves + .into_iter() + .map( + |( + addr, + CollatedLeaf { + balance, + nonce, + // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/707 + // we shouldn't ignore these fields + code: _, + code_length: _, + storage: _, + }, + )| { + ( + addr, + AccountRlp { + nonce: nonce.unwrap_or_default(), + balance: balance.unwrap_or_default(), + storage_root: H256::zero(), + code_hash: H256::zero(), + }, + ) + }, + ) + .collect(), + hashes, + ), + ) } fn visit( @@ -183,8 +173,8 @@ fn visit( } Node::Leaf(SmtLeaf { node_type, - address, // TODO(0xaatif): field should be fixed length - value, // TODO(0xaatif): field should be fixed length + address, + value, }) => { let address = Address::from_slice(&address); let collated = leaves.entry(address).or_default(); @@ -208,7 +198,6 @@ fn visit( collated.code = Some(value) } SmtLeafType::Storage(slot) => { - // TODO(0xaatif): ^ field should be fixed length let clobbered = collated.storage.insert(U256::from_big_endian(&slot), value); ensure!(clobbered.is_none()) } @@ -224,6 +213,65 @@ fn visit( #[test] fn test_tries() { + type Smt = smt_trie::smt::Smt; + use ethereum_types::BigEndianHash as _; + use plonky2::field::types::{Field, Field64 as _}; + + // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/707 + // this logic should live in StateSmt, but we need to + // - abstract over state and storage tries + // - parameterize the account types + // we preserve this code as a tested record of how it _should_ + // be done. + fn node2trie(node: Node) -> anyhow::Result { + let mut trie = Smt::default(); + let mut hashes = BTreeMap::new(); + let mut leaves = BTreeMap::new(); + visit(&mut hashes, &mut leaves, Stack::new(), node)?; + for (key, hash) in hashes { + trie.set_hash( + key.into_smt_bits(), + smt_trie::smt::HashOut { + elements: { + let ethereum_types::U256(arr) = hash.into_uint(); + for u in arr { + ensure!(u < smt_trie::smt::F::ORDER); + } + arr.map(smt_trie::smt::F::from_canonical_u64) + }, + }, + ); + } + for ( + addr, + CollatedLeaf { + balance, + nonce, + code, + code_length, + storage, + }, + ) in leaves + { + use smt_trie::keys::{key_balance, key_code, key_code_length, key_nonce, key_storage}; + + for (value, key_fn) in [ + (balance, key_balance as fn(_) -> _), + (nonce, key_nonce), + (code, key_code), + (code_length, key_code_length), + ] { + if let Some(value) = value { + trie.set(key_fn(addr), value); + } + } + for (slot, value) in storage { + trie.set(key_storage(addr, slot), value); + } + } + Ok(trie) + } + for (ix, case) in serde_json::from_str::>(include_str!("cases/hermez_cdk_erigon.json")) .unwrap() @@ -232,10 +280,11 @@ fn test_tries() { { println!("case {}", ix); let instructions = crate::wire::parse(&case.bytes).unwrap(); - let frontend = frontend(instructions).unwrap(); + let (node, _code) = fold(instructions).unwrap(); + let trie = node2trie(node).unwrap(); assert_eq!(case.expected_state_root, { let mut it = [0; 32]; - smt_trie::utils::hashout2u(frontend.trie.root).to_big_endian(&mut it); + smt_trie::utils::hashout2u(trie.root).to_big_endian(&mut it); ethereum_types::H256(it) }); } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index a40d39db8..4c8e5a6c8 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -465,7 +465,7 @@ impl From for HashedPartialTrie { } } -// TODO(0xaatif): trackme +// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/706 // We're covering for [`smt_trie`] in a couple of ways: // - insertion operations aren't fallible, they just panic. // - it documents a requirement that `set_hash` is called before `set`. @@ -503,6 +503,17 @@ impl StateTrie for StateSmt { } impl StateSmt { + #[deprecated = "this should only be called from the frontend parsing"] + pub(crate) fn new_unchecked( + address2state: BTreeMap, + hashed_out: BTreeMap, + ) -> Self { + Self { + address2state, + hashed_out, + } + } + fn as_smt(&self) -> smt_trie::smt::Smt { let Self { address2state, @@ -526,7 +537,8 @@ impl StateSmt { smt.set(smt_trie::keys::key_balance(*addr), *balance); smt.set(smt_trie::keys::key_code(*addr), code_hash.into_uint()); smt.set( - // REVIEW(0xaatif): I don't know what to do here + // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/707 + // combined abstraction for state and storage smt_trie::keys::key_storage(*addr, U256::zero()), storage_root.into_uint(), ); @@ -549,6 +561,8 @@ mod conv_hash { elements: array::from_fn(|_ix| { let (a, b, c, d, e, f, g, h) = bytes.next_tuple().unwrap(); // REVIEW(0xaatif): what endianness? + // do we want the `canonical_u64` methods like + // the frontend uses? GoldilocksField(u64::from_be_bytes([a, b, c, d, e, f, g, h])) }), }; @@ -560,6 +574,9 @@ mod conv_hash { build_array::ArrayBuilder::from_iter( elements .into_iter() + // REVIEW(0xaatif): what endianness? + // do we want the `canonical_u64` methods + // like the frontend uses? .flat_map(|GoldilocksField(u)| u.to_be_bytes()), ) .build_exact() diff --git a/trace_decoder/src/wire.rs b/trace_decoder/src/wire.rs index 9d1a7fb10..52a6a9b40 100644 --- a/trace_decoder/src/wire.rs +++ b/trace_decoder/src/wire.rs @@ -80,18 +80,22 @@ pub enum Instruction { } #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/705 +// `address` and `value` should be fixed length fields pub struct SmtLeaf { pub node_type: SmtLeafType, - pub address: NonEmpty>, // TODO(0xaatif): this should be a fixed length - pub value: NonEmpty>, // TODO(0xaatif): this should be a fixed length + pub address: NonEmpty>, + pub value: NonEmpty>, } #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/705 +// `Storage` should contain a fixed length field pub enum SmtLeafType { Balance, Nonce, Code, - Storage(NonEmpty>), // TODO(0xaatif): this should be a fixed length + Storage(NonEmpty>), CodeLength, } diff --git a/zero/src/bin/trie_diff.rs b/zero/src/bin/trie_diff.rs index b0005b819..7d4630404 100644 --- a/zero/src/bin/trie_diff.rs +++ b/zero/src/bin/trie_diff.rs @@ -140,9 +140,6 @@ async fn main() -> Result<()> { state_trie: observer.data[prover_tries.batch_index] .tries .state - .as_ref() - .left() - .unwrap() .as_hashed_partial_trie() .clone(), transaction_trie: observer.data[prover_tries.batch_index] diff --git a/zero/src/prover.rs b/zero/src/prover.rs index 80612afc3..e9397726a 100644 --- a/zero/src/prover.rs +++ b/zero/src/prover.rs @@ -43,6 +43,8 @@ pub const WIRE_DISPOSITION: WireDisposition = { cfg_if::cfg_if! { if #[cfg(feature = "eth_mainnet")] { WireDisposition::Type1 + } else if #[cfg(feature = "cdk_erigon")] { + WireDisposition::Type2 } else { compile_error!("must select a feature"); } From 13f03ad4a6a49dc2b1e9010d641e0202b6f238d0 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 7 Oct 2024 16:42:00 +0100 Subject: [PATCH 10/13] trace_decoder::typed_mpt -> trace_decoder::tries --- trace_decoder/src/core.rs | 4 ++-- trace_decoder/src/lib.rs | 2 +- trace_decoder/src/observer.rs | 2 +- trace_decoder/src/{typed_mpt.rs => tries.rs} | 2 +- trace_decoder/src/type1.rs | 4 ++-- trace_decoder/src/type2.rs | 2 +- 6 files changed, 8 insertions(+), 8 deletions(-) rename trace_decoder/src/{typed_mpt.rs => tries.rs} (99%) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 8a07d98f4..ad2784e7b 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -21,10 +21,10 @@ use zk_evm_common::gwei_to_wei; use crate::{ observer::{DummyObserver, Observer}, - typed_mpt::StateSmt, + tries::StateSmt, }; use crate::{ - typed_mpt::{MptKey, ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie}, + tries::{MptKey, ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie}, BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage, OtherBlockData, SeparateStorageTriesPreImage, SeparateTriePreImage, SeparateTriePreImages, TxnInfo, TxnMeta, TxnTrace, diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index eea5ebe80..057d11e89 100644 --- a/trace_decoder/src/lib.rs +++ b/trace_decoder/src/lib.rs @@ -56,9 +56,9 @@ mod interface; pub use interface::*; +mod tries; mod type1; mod type2; -mod typed_mpt; mod wire; pub use core::{entrypoint, WireDisposition}; diff --git a/trace_decoder/src/observer.rs b/trace_decoder/src/observer.rs index 320019e55..f9811e87c 100644 --- a/trace_decoder/src/observer.rs +++ b/trace_decoder/src/observer.rs @@ -4,7 +4,7 @@ use std::marker::PhantomData; use ethereum_types::{H256, U256}; use crate::core::IntraBlockTries; -use crate::typed_mpt::{ReceiptTrie, StorageTrie, TransactionTrie}; +use crate::tries::{ReceiptTrie, StorageTrie, TransactionTrie}; /// Observer API for the trace decoder. /// Observer is used to collect various debugging and metadata info diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/tries.rs similarity index 99% rename from trace_decoder/src/typed_mpt.rs rename to trace_decoder/src/tries.rs index 4c8e5a6c8..fb7f37f83 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/tries.rs @@ -1,4 +1,4 @@ -//! Principled MPT types used in this library. +//! Principled trie types and abstractions used in this library. use core::fmt; use std::{cmp, collections::BTreeMap, marker::PhantomData}; diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 96f49d94d..c44beaec7 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -12,7 +12,7 @@ use mpt_trie::partial_trie::OnOrphanedHashNode; use nunny::NonEmpty; use u4::U4; -use crate::typed_mpt::{MptKey, StateMpt, StorageTrie}; +use crate::tries::{MptKey, StateMpt, StorageTrie}; use crate::wire::{Instruction, SmtLeaf}; #[derive(Debug, Clone)] @@ -379,7 +379,7 @@ fn finish_stack(v: &mut Vec) -> anyhow::Result { #[test] fn test_tries() { - use crate::typed_mpt::StateTrie as _; + use crate::tries::StateTrie as _; for (ix, case) in serde_json::from_str::>(include_str!("cases/zero_jerigon.json")) diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index 7cb952888..d03831e43 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -12,7 +12,7 @@ use nunny::NonEmpty; use stackstack::Stack; use crate::{ - typed_mpt::{SmtKey, StateSmt}, + tries::{SmtKey, StateSmt}, wire::{Instruction, SmtLeaf, SmtLeafType}, }; From adeaa4e7d7d994c80880a295aab5a4296bf1a9b6 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Fri, 11 Oct 2024 08:00:53 +0100 Subject: [PATCH 11/13] ci: private rustdoc --- .github/workflows/lint.yml | 12 +++++++++--- trace_decoder/src/core.rs | 9 +++++---- trace_decoder/src/tries.rs | 8 ++++---- trace_decoder/src/type2.rs | 2 +- trace_decoder/src/wire.rs | 4 ++-- 5 files changed, 21 insertions(+), 14 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 446d6a314..13089b3d0 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -13,8 +13,6 @@ concurrency: env: CARGO_TERM_COLOR: always - BINSTALL_NO_CONFIRM: true - RUSTDOCFLAGS: "-D warnings" jobs: clippy: @@ -42,7 +40,15 @@ jobs: steps: - uses: actions/checkout@v3 - uses: ./.github/actions/rust - - run: cargo doc --all --no-deps + - run: RUSTDOCFLAGS='-D warnings -A rustdoc::private_intra_doc_links' cargo doc --all --no-deps + # TODO(zero): https://github.com/0xPolygonZero/zk_evm/issues/718 + - run: > + RUSTDOCFLAGS='-D warnings -A rustdoc::private_intra_doc_links' cargo doc --no-deps --document-private-items + --package trace_decoder + --package compat + --package smt_trie + --package zk_evm_proc_macro + --package zk_evm_common cargo-fmt: runs-on: ubuntu-latest timeout-minutes: 5 diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index ad2784e7b..46495030f 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -30,7 +30,9 @@ use crate::{ TxnInfo, TxnMeta, TxnTrace, }; -/// When parsing tries from binary format, which type to deserialize as. +/// Expected trie type when parsing from binary in a [`BlockTrace`]. +/// +/// See [`crate::wire`] and [`CombinedPreImages`] for more. #[derive(Debug)] pub enum WireDisposition { /// MPT @@ -171,8 +173,7 @@ pub fn entrypoint( /// [`HashedPartialTrie`](mpt_trie::partial_trie::HashedPartialTrie), /// or a [`wire`](crate::wire)-encoded representation of one. /// -/// Turn either of those into our [`typed_mpt`](crate::typed_mpt) -/// representations. +/// Turn either of those into our [internal representations](crate::tries). #[allow(clippy::type_complexity)] fn start( pre_images: BlockTraceTriePreImages, @@ -900,7 +901,7 @@ fn map_receipt_bytes(bytes: Vec) -> anyhow::Result> { /// If there are any txns that create contracts, then they will also /// get added here as we process the deltas. struct Hash2Code { - /// Key must always be [`hash`] of value. + /// Key must always be [`hash`](keccak_hash) of value. inner: HashMap>, } diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index fb7f37f83..60f0e0271 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -36,8 +36,8 @@ impl TypedMpt { self.inner.insert(key.into_nibbles(), hash)?; Ok(()) } - /// Returns an [`Error`] if the `key` crosses into a part of the trie that - /// isn't hydrated. + /// Returns [`Err`] if the `key` crosses into a part of the trie that + /// is hashed out. fn insert(&mut self, key: MptKey, value: T) -> anyhow::Result<()> where T: rlp::Encodable + rlp::Decodable, @@ -47,7 +47,7 @@ impl TypedMpt { Ok(()) } /// Note that this returns [`None`] if `key` crosses into a part of the - /// trie that isn't hydrated. + /// trie that is hashed out. /// /// # Panics /// - If [`rlp::decode`]-ing for `T` doesn't round-trip. @@ -366,7 +366,7 @@ impl From for HashedPartialTrie { } } -/// TODO(0xaatif): document this after refactoring is done https://github.com/0xPolygonZero/zk_evm/issues/275 +/// TODO(0xaatif): document this after refactoring is done pub trait StateTrie { type Key; fn insert_by_address(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()>; diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index d03831e43..fc5203067 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -42,7 +42,7 @@ pub fn frontend(instructions: impl IntoIterator) -> anyhow:: /// Node in a binary (SMT) tree. /// -/// This is an intermediary type on the way to [`SmtTrie`]. +/// This is an intermediary type on the way to [`StateSmt`]. enum Node { Branch(EitherOrBoth>), Hash([u8; 32]), diff --git a/trace_decoder/src/wire.rs b/trace_decoder/src/wire.rs index 52a6a9b40..63dee6040 100644 --- a/trace_decoder/src/wire.rs +++ b/trace_decoder/src/wire.rs @@ -1,6 +1,6 @@ //! We support two wire formats: -//! - Type 1, based on [this specification](https://gist.github.com/mandrigin/ff7eccf30d0ef9c572bafcb0ab665cff#the-bytes-layout). -//! - Type 2, loosely based on [this specification](https://github.com/0xPolygonHermez/cdk-erigon/blob/d1d6b3c7a4c81c46fd995c1baa5c1f8069ff0348/turbo/trie/WITNESS.md) +//! - Type 1 (AKA MPT), based on [this specification](https://gist.github.com/mandrigin/ff7eccf30d0ef9c572bafcb0ab665cff#the-bytes-layout). +//! - Type 2 (AKA SMT), loosely based on [this specification](https://github.com/0xPolygonHermez/cdk-erigon/blob/d1d6b3c7a4c81c46fd995c1baa5c1f8069ff0348/turbo/trie/WITNESS.md) //! //! Fortunately, their opcodes don't conflict, so we can have a single //! [`Instruction`] type, with shared parsing logic in this module, and bail on From 512b6937ea807750bc9f7c35c3acce6d9a2006eb Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Fri, 11 Oct 2024 12:43:03 +0100 Subject: [PATCH 12/13] review: markups --- trace_decoder/src/tries.rs | 32 ++++++++++++--------- trace_decoder/src/type2.rs | 59 ++++++++++++++++++-------------------- 2 files changed, 47 insertions(+), 44 deletions(-) diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 60f0e0271..f61dd9800 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -191,7 +191,7 @@ fn mpt_key_into_hash() { /// Bounded sequence of bits, /// used as a key for [`StateSmt`]. /// -/// Semantically equivalent to +/// Semantically equivalent to [`smt_trie::bits::Bits`]. #[derive(Clone, Copy)] pub struct SmtKey { bits: bitvec::array::BitArray<[u8; 32]>, @@ -255,7 +255,7 @@ impl SmtKey { impl From
for SmtKey { fn from(addr: Address) -> Self { let H256(bytes) = keccak_hash::keccak(addr); - Self::new(BitArray::<_>::new(bytes)).unwrap() + Self::new(BitArray::<_>::new(bytes)).expect("SmtKey has room for 256 bits") } } @@ -503,7 +503,6 @@ impl StateTrie for StateSmt { } impl StateSmt { - #[deprecated = "this should only be called from the frontend parsing"] pub(crate) fn new_unchecked( address2state: BTreeMap, hashed_out: BTreeMap, @@ -552,8 +551,17 @@ mod conv_hash { use ethereum_types::H256; use itertools::Itertools as _; - use plonky2::{field::goldilocks_field::GoldilocksField, hash::hash_types::HashOut}; + use plonky2::{ + field::{ + goldilocks_field::GoldilocksField, + types::{Field as _, PrimeField64}, + }, + hash::hash_types::HashOut, + }; + /// # Panics + /// - On certain inputs if `debug_assertions` are enabled. See + /// [`GoldilocksField::from_canonical_u64`] for more. pub fn eth2smt(H256(bytes): H256) -> smt_trie::smt::HashOut { let mut bytes = bytes.into_iter(); // (no unsafe, no unstable) @@ -561,9 +569,7 @@ mod conv_hash { elements: array::from_fn(|_ix| { let (a, b, c, d, e, f, g, h) = bytes.next_tuple().unwrap(); // REVIEW(0xaatif): what endianness? - // do we want the `canonical_u64` methods like - // the frontend uses? - GoldilocksField(u64::from_be_bytes([a, b, c, d, e, f, g, h])) + GoldilocksField::from_canonical_u64(u64::from_be_bytes([a, b, c, d, e, f, g, h])) }), }; assert_eq!(bytes.len(), 0); @@ -573,11 +579,9 @@ mod conv_hash { H256( build_array::ArrayBuilder::from_iter( elements - .into_iter() - // REVIEW(0xaatif): what endianness? - // do we want the `canonical_u64` methods - // like the frontend uses? - .flat_map(|GoldilocksField(u)| u.to_be_bytes()), + .iter() + .map(GoldilocksField::to_canonical_u64) + .flat_map(u64::to_be_bytes), ) .build_exact() .unwrap(), @@ -586,10 +590,12 @@ mod conv_hash { #[test] fn test() { + use plonky2::field::types::Field64 as _; + let mut max = std::iter::repeat(GoldilocksField::ORDER - 1).flat_map(u64::to_be_bytes); for h in [ H256::zero(), H256(array::from_fn(|ix| ix as u8)), - H256([u8::MAX; 32]), + H256(array::from_fn(|_| max.next().unwrap())), ] { assert_eq!(smt2eth(eth2smt(h)), h); } diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index fc5203067..a71761533 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -117,39 +117,36 @@ fn node2trie(node: Node) -> anyhow::Result { let mut hashes = BTreeMap::new(); let mut leaves = BTreeMap::new(); visit(&mut hashes, &mut leaves, Stack::new(), node)?; - Ok( - #[expect(deprecated, reason = "this is the frontend")] - StateSmt::new_unchecked( - leaves - .into_iter() - .map( - |( + Ok(StateSmt::new_unchecked( + leaves + .into_iter() + .map( + |( + addr, + CollatedLeaf { + balance, + nonce, + // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/707 + // we shouldn't ignore these fields + code: _, + code_length: _, + storage: _, + }, + )| { + ( addr, - CollatedLeaf { - balance, - nonce, - // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/707 - // we shouldn't ignore these fields - code: _, - code_length: _, - storage: _, + AccountRlp { + nonce: nonce.unwrap_or_default(), + balance: balance.unwrap_or_default(), + storage_root: H256::zero(), + code_hash: H256::zero(), }, - )| { - ( - addr, - AccountRlp { - nonce: nonce.unwrap_or_default(), - balance: balance.unwrap_or_default(), - storage_root: H256::zero(), - code_hash: H256::zero(), - }, - ) - }, - ) - .collect(), - hashes, - ), - ) + ) + }, + ) + .collect(), + hashes, + )) } fn visit( From 6b822b49f3d9efe55922814d5a0ecdf7b7f83c6f Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sat, 12 Oct 2024 01:05:58 +0100 Subject: [PATCH 13/13] doc: endianness --- trace_decoder/src/tries.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index f61dd9800..91add4d98 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -547,6 +547,11 @@ impl StateSmt { } mod conv_hash { + //! We [`u64::to_le_bytes`] because: + //! - Reference go code just puns the bytes: + //! - It's better to fix the endianness for correctness. + //! - Most (consumer) CPUs are little-endian. + use std::array; use ethereum_types::H256; @@ -568,8 +573,7 @@ mod conv_hash { let ret = HashOut { elements: array::from_fn(|_ix| { let (a, b, c, d, e, f, g, h) = bytes.next_tuple().unwrap(); - // REVIEW(0xaatif): what endianness? - GoldilocksField::from_canonical_u64(u64::from_be_bytes([a, b, c, d, e, f, g, h])) + GoldilocksField::from_canonical_u64(u64::from_le_bytes([a, b, c, d, e, f, g, h])) }), }; assert_eq!(bytes.len(), 0); @@ -581,7 +585,7 @@ mod conv_hash { elements .iter() .map(GoldilocksField::to_canonical_u64) - .flat_map(u64::to_be_bytes), + .flat_map(u64::to_le_bytes), ) .build_exact() .unwrap(), @@ -591,7 +595,7 @@ mod conv_hash { #[test] fn test() { use plonky2::field::types::Field64 as _; - let mut max = std::iter::repeat(GoldilocksField::ORDER - 1).flat_map(u64::to_be_bytes); + let mut max = std::iter::repeat(GoldilocksField::ORDER - 1).flat_map(u64::to_le_bytes); for h in [ H256::zero(), H256(array::from_fn(|ix| ix as u8)),