From da7fea8f4f57b14496123a6fd7b44a2130c0f6fe Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Wed, 2 Oct 2024 17:17:02 +0100 Subject: [PATCH 01/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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 ab6ba910bb7af8274e398d26d03b6baa5bd08510 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 7 Oct 2024 17:47:03 +0100 Subject: [PATCH 13/25] mark: 0xaatif/storage-trait From 437be6b7f95fcc765bea294cdba7e0aa9c51a724 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Tue, 8 Oct 2024 03:24:39 +0100 Subject: [PATCH 14/25] chore: naming --- trace_decoder/src/core.rs | 68 +++++++++++++++--------------- trace_decoder/src/tries.rs | 84 +++++++++++++++++++++----------------- trace_decoder/src/type1.rs | 10 ++--- trace_decoder/src/type2.rs | 8 ++-- 4 files changed, 90 insertions(+), 80 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 46495030f..8d1704a19 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}, - tries::StateSmt, + tries::Type2World, }; use crate::{ - tries::{MptKey, ReceiptTrie, StateMpt, StateTrie, StorageTrie, TransactionTrie}, + tries::{MptKey, ReceiptTrie, StorageTrie, TransactionTrie, Type1World, World}, BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage, OtherBlockData, SeparateStorageTriesPreImage, SeparateTriePreImage, SeparateTriePreImages, TxnInfo, TxnMeta, TxnTrace, @@ -46,7 +46,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); @@ -179,7 +179,7 @@ fn start( pre_images: BlockTraceTriePreImages, wire_disposition: WireDisposition, ) -> anyhow::Result<( - Either, + Either, BTreeMap, Hash2Code, )> { @@ -191,7 +191,7 @@ fn start( storage: SeparateStorageTriesPreImage::MultipleTries(storage), }) => { let state = state.items().try_fold( - StateMpt::default(), + Type1World::default(), |mut acc, (nibbles, hash_or_val)| { let path = MptKey::from_nibbles(nibbles); match hash_or_val { @@ -391,7 +391,7 @@ pub struct FatalMissingCode(pub bool); /// Does the main work mentioned in the [module documentation](super). #[allow(clippy::too_many_arguments)] -fn middle( +fn middle + Clone>( // state at the beginning of the block mut state_trie: StateTrieT, // storage at the beginning of the block @@ -409,7 +409,7 @@ fn middle( observer: &mut impl Observer, ) -> anyhow::Result>> where - StateTrieT::Key: Ord + From
, + StateTrieT::StateKey: Ord + From
, { // Initialise the storage tries. for (haddr, acct) in state_trie.iter() { @@ -451,7 +451,7 @@ where // 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( @@ -511,7 +511,7 @@ where .context(format!("couldn't decode receipt in txn {tx_hash:x}"))?; let (mut acct, born) = state_trie - .get_by_address(addr) + .get_account_info(addr) .map(|acct| (acct, false)) .unwrap_or((AccountRlp::default(), true)); @@ -524,7 +524,7 @@ where if born || just_access { state_trie .clone() - .insert_by_address(addr, acct) + .insert_account_info(addr, acct) .context(format!( "couldn't reach state of {} address {addr:x} in txn {tx_hash:x}", match born { @@ -606,11 +606,11 @@ where acct.storage_root = storage.root(); } - state_trie.insert_by_address(addr, acct)?; - state_mask.insert(::from(addr)); + state_trie.insert_account_info(addr, acct)?; + state_mask.insert(::from(addr)); } else { // Simple state access - state_mask.insert(::from(addr)); + state_mask.insert(::from(addr)); } if self_destructed { @@ -633,13 +633,13 @@ where withdrawals: match loop_ix == loop_len { true => { for (addr, amt) in &withdrawals { - state_mask.insert(::from(*addr)); + state_mask.insert(::from(*addr)); let mut acct = state_trie - .get_by_address(*addr) + .get_account_info(*addr) .context(format!("missing address {addr:x} for withdrawal"))?; acct.balance += *amt; state_trie - .insert_by_address(*addr, acct) + .insert_account_info(*addr, acct) // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 // Add an entry API .expect("insert must succeed with the same key as a successful `get`"); @@ -687,16 +687,16 @@ where } /// Performs all the pre-txn execution rules of the targeted network. -fn do_pre_execution( +fn do_pre_execution + Clone>( block: &BlockMetadata, 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<()> where - StateTrieT::Key: From
+ Ord, + StateTrieT::StateKey: From
+ Ord, { // Ethereum mainnet: EIP-4788 if cfg!(feature = "eth_mainnet") { @@ -729,16 +729,16 @@ where /// /// This is Polygon-CDK-specific, and runs at the start of the block, /// before any transactions (as per the Etrog specification). -fn do_scalable_hook( +fn do_scalable_hook + Clone>( block: &BlockMetadata, 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<()> where - StateTrieT::Key: From
+ Ord, + StateTrieT::StateKey: From
+ Ord, { use evm_arithmetization::testing_utils::{ ADDRESS_SCALABLE_L2, ADDRESS_SCALABLE_L2_ADDRESS_HASHED, GLOBAL_EXIT_ROOT_ADDRESS, @@ -786,13 +786,13 @@ where scalable_storage.insert(slot, alloy::rlp::encode(prev_block_root_hash.compat()))?; scalable_trim.insert(slot); - trim_state.insert(::from(ADDRESS_SCALABLE_L2)); + trim_state.insert(::from(ADDRESS_SCALABLE_L2)); let mut scalable_acct = state_trie - .get_by_address(ADDRESS_SCALABLE_L2) + .get_account_info(ADDRESS_SCALABLE_L2) .context("missing scalable contract address")?; scalable_acct.storage_root = scalable_storage.root(); state_trie - .insert_by_address(ADDRESS_SCALABLE_L2, scalable_acct) + .insert_account_info(ADDRESS_SCALABLE_L2, scalable_acct) // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 // Add an entry API .expect("insert must succeed with the same key as a successful `get`"); @@ -812,13 +812,13 @@ where ger_storage.insert(slot, alloy::rlp::encode(l1blockhash.compat()))?; ger_trim.insert(slot); - trim_state.insert(::from(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) + .get_account_info(GLOBAL_EXIT_ROOT_ADDRESS) .context("missing GER contract address")?; ger_acct.storage_root = ger_storage.root(); state_trie - .insert_by_address(GLOBAL_EXIT_ROOT_ADDRESS, ger_acct) + .insert_account_info(GLOBAL_EXIT_ROOT_ADDRESS, ger_acct) // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 // Add an entry API .expect("insert must succeed with the same key as a successful `get`"); @@ -832,16 +832,16 @@ where /// /// This is Cancun-specific, and runs at the start of the block, /// before any transactions (as per the EIP). -fn do_beacon_hook( +fn do_beacon_hook + Clone>( block_timestamp: U256, 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<()> where - StateTrieT::Key: From
+ Ord, + StateTrieT::StateKey: From
+ Ord, { use evm_arithmetization::testing_utils::{ BEACON_ROOTS_CONTRACT_ADDRESS, BEACON_ROOTS_CONTRACT_ADDRESS_HASHED, HISTORY_BUFFER_LENGTH, @@ -874,13 +874,13 @@ where } } } - trim_state.insert(::from(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) + .get_account_info(BEACON_ROOTS_CONTRACT_ADDRESS) .context("missing beacon contract address")?; beacon_acct.storage_root = beacon_storage.root(); state_trie - .insert_by_address(BEACON_ROOTS_CONTRACT_ADDRESS, beacon_acct) + .insert_account_info(BEACON_ROOTS_CONTRACT_ADDRESS, beacon_acct) // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 // Add an entry API .expect("insert must succeed with the same key as a successful `get`"); diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index f61dd9800..bdb4bb58a 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -366,15 +366,23 @@ impl From for HashedPartialTrie { } } -/// 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<()>; - 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`. - fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()>; - fn iter(&self) -> impl Iterator + '_; +/// The state and storage of all accounts. +/// +/// Some parts of the tries may be _hashed out_. +pub trait World { + type StateKey; + type AccountInfo; + fn insert_account_info( + &mut self, + address: Address, + account: Self::AccountInfo, + ) -> anyhow::Result<()>; + fn get_account_info(&self, address: Address) -> Option; + /// Hacky method to workaround MPT shenanigans. + fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; + /// _Hash out_ parts of the (state) trie that aren't in `addresses`. + fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()>; + fn iter(&self) -> impl Iterator + '_; fn root(&self) -> H256; } @@ -382,14 +390,14 @@ pub trait StateTrie { /// /// See #[derive(Debug, Clone, Default)] -pub struct StateMpt { - typed: TypedMpt, +pub struct Type1World { + state: TypedMpt, } -impl StateMpt { +impl Type1World { pub fn new(strategy: OnOrphanedHashNode) -> Self { Self { - typed: TypedMpt { + state: TypedMpt { inner: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), _ty: PhantomData, }, @@ -397,7 +405,7 @@ impl StateMpt { } /// Insert a _hashed out_ part of the trie pub fn insert_hash_by_key(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { - self.typed.insert_hash(key, hash) + self.state.insert_hash(key, hash) } #[deprecated = "prefer operations on `Address` where possible, as SMT support requires this"] pub fn insert_by_hashed_address( @@ -405,61 +413,62 @@ impl StateMpt { key: H256, account: AccountRlp, ) -> anyhow::Result<()> { - self.typed.insert(MptKey::from_hash(key), account) + self.state.insert(MptKey::from_hash(key), account) } pub fn iter(&self) -> impl Iterator + '_ { - self.typed + self.state .iter() .map(|(key, rlp)| (key.into_hash().expect("key is always H256"), rlp)) } pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { - self.typed.as_hashed_partial_trie() + self.state.as_hashed_partial_trie() } } -impl StateTrie for StateMpt { - type Key = MptKey; - fn insert_by_address(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { +impl World for Type1World { + type StateKey = MptKey; + type AccountInfo = AccountRlp; + fn insert_account_info(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { #[expect(deprecated)] self.insert_by_hashed_address(keccak_hash::keccak(address), account) } - fn get_by_address(&self, address: Address) -> Option { - self.typed + fn get_account_info(&self, address: Address) -> Option { + self.state .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> { delete_node_and_report_remaining_key_if_branch_collapsed( - self.typed.as_mut_hashed_partial_trie_unchecked(), + self.state.as_mut_hashed_partial_trie_unchecked(), MptKey::from_address(address), ) } fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()> { let inner = mpt_trie::trie_subsets::create_trie_subset( - self.typed.as_hashed_partial_trie(), + self.state.as_hashed_partial_trie(), addresses.into_iter().map(MptKey::into_nibbles), )?; - self.typed = TypedMpt { + self.state = TypedMpt { inner, _ty: PhantomData, }; Ok(()) } fn iter(&self) -> impl Iterator + '_ { - self.typed + self.state .iter() .map(|(key, rlp)| (key.into_hash().expect("key is always H256"), rlp)) } fn root(&self) -> H256 { - self.typed.root() + self.state.root() } } -impl From for HashedPartialTrie { - fn from(value: StateMpt) -> Self { - let StateMpt { - typed: TypedMpt { inner, _ty }, +impl From for HashedPartialTrie { + fn from(value: Type1World) -> Self { + let Type1World { + state: TypedMpt { inner, _ty }, } = value; inner } @@ -470,18 +479,19 @@ impl From for HashedPartialTrie { // - 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 { +pub struct Type2World { address2state: BTreeMap, hashed_out: BTreeMap, } -impl StateTrie for StateSmt { - type Key = SmtKey; - fn insert_by_address(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { +impl World for Type2World { + type StateKey = SmtKey; + type AccountInfo = AccountRlp; + fn insert_account_info(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { self.address2state.insert(address, account); Ok(()) } - fn get_by_address(&self, address: Address) -> Option { + fn get_account_info(&self, address: Address) -> Option { self.address2state.get(&address).copied() } fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { @@ -502,7 +512,7 @@ impl StateTrie for StateSmt { } } -impl StateSmt { +impl Type2World { pub(crate) fn new_unchecked( address2state: BTreeMap, hashed_out: BTreeMap, diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index c44beaec7..1c028a799 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -12,12 +12,12 @@ use mpt_trie::partial_trie::OnOrphanedHashNode; use nunny::NonEmpty; use u4::U4; -use crate::tries::{MptKey, StateMpt, StorageTrie}; +use crate::tries::{MptKey, StorageTrie, Type1World}; use crate::wire::{Instruction, SmtLeaf}; #[derive(Debug, Clone)] pub struct Frontend { - pub state: StateMpt, + pub state: Type1World, pub code: BTreeSet>>, pub storage: BTreeMap, } @@ -27,7 +27,7 @@ impl Default for Frontend { // which covers branch-to-extension collapse edge cases. fn default() -> Self { Self { - state: StateMpt::new(OnOrphanedHashNode::CollapseToExtension), + state: Type1World::new(OnOrphanedHashNode::CollapseToExtension), code: BTreeSet::new(), storage: BTreeMap::new(), } @@ -379,7 +379,7 @@ fn finish_stack(v: &mut Vec) -> anyhow::Result { #[test] fn test_tries() { - use crate::tries::StateTrie as _; + use crate::tries::World as _; for (ix, case) in serde_json::from_str::>(include_str!("cases/zero_jerigon.json")) @@ -393,7 +393,7 @@ fn test_tries() { assert_eq!(case.expected_state_root, frontend.state.root()); for (haddr, acct) in frontend.state.iter() { - if acct.storage_root != StateMpt::default().root() { + if acct.storage_root != Type1World::default().root() { assert!(frontend.storage.contains_key(&haddr)) } } diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index a71761533..4ad31fdeb 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::{ - tries::{SmtKey, StateSmt}, + tries::{SmtKey, Type2World}, wire::{Instruction, SmtLeaf, SmtLeafType}, }; @@ -27,7 +27,7 @@ pub struct CollatedLeaf { } pub struct Frontend { - pub trie: StateSmt, + pub trie: Type2World, pub code: HashSet>>, } @@ -113,11 +113,11 @@ fn fold1(instructions: impl IntoIterator) -> anyhow::Result< } } -fn node2trie(node: Node) -> anyhow::Result { +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(StateSmt::new_unchecked( + Ok(Type2World::new_unchecked( leaves .into_iter() .map( From 42c4e4cb9c1455fbdc2d338ce01a13aac920fd04 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Wed, 9 Oct 2024 02:49:05 +0100 Subject: [PATCH 15/25] refactor: store_int, store_hash --- trace_decoder/src/core.rs | 19 +++++++++---------- trace_decoder/src/tries.rs | 39 +++++++++++++++++++++++++++++--------- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 8d1704a19..e866c5edf 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -412,7 +412,7 @@ where StateTrieT::StateKey: Ord + From
, { // Initialise the storage tries. - for (haddr, acct) in state_trie.iter() { + for (haddr, acct) in state_trie.iter_account_info() { let storage = storage_tries.entry(haddr).or_insert({ let mut it = StorageTrie::default(); it.insert_hash(MptKey::default(), acct.storage_root) @@ -599,7 +599,7 @@ where // this is actually a delete true => storage_mask.extend(storage.reporting_remove(slot)?), false => { - storage.insert(slot, rlp::encode(&v).to_vec())?; + storage.store_int(slot, v)?; } } } @@ -757,9 +757,8 @@ where let timestamp_slot_key = MptKey::from_slot_position(U256::from(TIMESTAMP_STORAGE_POS.1)); let timestamp = scalable_storage - .get(×tamp_slot_key) - .map(rlp::decode::) - .unwrap_or(Ok(0.into()))?; + .load_int(timestamp_slot_key) + .unwrap_or_default(); let timestamp = core::cmp::max(timestamp, block.block_timestamp); // Store block number and largest timestamp @@ -770,8 +769,8 @@ where ] { let slot = MptKey::from_slot_position(ix); - // These values are never 0. - scalable_storage.insert(slot, alloy::rlp::encode(u.compat()))?; + ensure!(u != U256::zero()); + scalable_storage.store_int(slot, u)?; scalable_trim.insert(slot); } @@ -783,7 +782,7 @@ where U256::from(STATE_ROOT_STORAGE_POS.1).to_big_endian(&mut arr[32..64]); let slot = MptKey::from_hash(keccak_hash::keccak(arr)); - scalable_storage.insert(slot, alloy::rlp::encode(prev_block_root_hash.compat()))?; + scalable_storage.store_hash(slot, prev_block_root_hash)?; scalable_trim.insert(slot); trim_state.insert(::from(ADDRESS_SCALABLE_L2)); @@ -809,7 +808,7 @@ where U256::from(GLOBAL_EXIT_ROOT_STORAGE_POS.1).to_big_endian(&mut arr[32..64]); let slot = MptKey::from_hash(keccak_hash::keccak(arr)); - ger_storage.insert(slot, alloy::rlp::encode(l1blockhash.compat()))?; + ger_storage.store_hash(slot, l1blockhash)?; ger_trim.insert(slot); trim_state.insert(::from(GLOBAL_EXIT_ROOT_ADDRESS)); @@ -869,7 +868,7 @@ where match u.is_zero() { true => beacon_trim.extend(beacon_storage.reporting_remove(slot)?), false => { - beacon_storage.insert(slot, alloy::rlp::encode(u.compat()))?; + beacon_storage.store_int(slot, u)?; beacon_trim.insert(slot); } } diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index bdb4bb58a..897243944 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -3,7 +3,8 @@ use core::fmt; use std::{cmp, collections::BTreeMap, marker::PhantomData}; -use anyhow::ensure; +use alloy_compat::Compat as _; +use anyhow::{ensure, Context}; use bitvec::{array::BitArray, slice::BitSlice}; use copyvec::CopyVec; use ethereum_types::{Address, BigEndianHash as _, H256, U256}; @@ -382,8 +383,9 @@ pub trait World { fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; /// _Hash out_ parts of the (state) trie that aren't in `addresses`. fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()>; - fn iter(&self) -> impl Iterator + '_; + fn iter_account_info(&self) -> impl Iterator + '_; fn root(&self) -> H256; + fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()>; } /// Global, [`Address`] `->` [`AccountRlp`]. @@ -392,6 +394,7 @@ pub trait World { #[derive(Debug, Clone, Default)] pub struct Type1World { state: TypedMpt, + storage: BTreeMap, } impl Type1World { @@ -401,6 +404,7 @@ impl Type1World { inner: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), _ty: PhantomData, }, + storage: BTreeMap::new(), } } /// Insert a _hashed out_ part of the trie @@ -455,7 +459,7 @@ impl World for Type1World { }; Ok(()) } - fn iter(&self) -> impl Iterator + '_ { + fn iter_account_info(&self) -> impl Iterator + '_ { self.state .iter() .map(|(key, rlp)| (key.into_hash().expect("key is always H256"), rlp)) @@ -463,12 +467,16 @@ impl World for Type1World { fn root(&self) -> H256 { self.state.root() } + fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()> { + todo!() + } } impl From for HashedPartialTrie { fn from(value: Type1World) -> Self { let Type1World { state: TypedMpt { inner, _ty }, + storage: _, } = value; inner } @@ -502,7 +510,7 @@ impl World for Type2World { let _ = address; Ok(()) } - fn iter(&self) -> impl Iterator + '_ { + fn iter_account_info(&self) -> impl Iterator + '_ { self.address2state .iter() .map(|(addr, acct)| (keccak_hash::keccak(addr), *acct)) @@ -510,6 +518,9 @@ impl World for Type2World { fn root(&self) -> H256 { conv_hash::smt2eth(self.as_smt().root) } + fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()> { + todo!() + } } impl Type2World { @@ -625,13 +636,23 @@ impl StorageTrie { untyped: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), } } - pub fn get(&mut self, key: &MptKey) -> Option<&[u8]> { - self.untyped.get(key.into_nibbles()) + pub fn load_int(&mut self, key: MptKey) -> anyhow::Result { + let bytes = self.untyped.get(key.into_nibbles()).context("no item")?; + Ok(rlp::decode(bytes)?) + } + pub fn store_int(&mut self, key: MptKey, value: U256) -> anyhow::Result<()> { + self.untyped + .insert(key.into_nibbles(), alloy::rlp::encode(value.compat()))?; + Ok(()) + } + pub fn store_hash(&mut self, key: MptKey, value: H256) -> anyhow::Result<()> { + self.untyped + .insert(key.into_nibbles(), alloy::rlp::encode(value.compat()))?; + Ok(()) } - pub fn insert(&mut self, key: MptKey, value: Vec) -> anyhow::Result>> { - let prev = self.get(&key).map(Vec::from); + pub fn insert(&mut self, key: MptKey, value: Vec) -> anyhow::Result<()> { self.untyped.insert(key.into_nibbles(), value)?; - Ok(prev) + Ok(()) } pub fn insert_hash(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { self.untyped.insert(key.into_nibbles(), hash)?; From d24c9f1abf2d40b371079fc4d2e37be7dea093e6 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Wed, 9 Oct 2024 03:54:53 +0100 Subject: [PATCH 16/25] refactor: store_int -> store_int_at_slot --- trace_decoder/src/core.rs | 21 ++++++++++----------- trace_decoder/src/interface.rs | 2 +- trace_decoder/src/tries.rs | 9 ++++++--- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index e866c5edf..840db1ba3 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -4,10 +4,9 @@ use std::{ mem, }; -use alloy_compat::Compat as _; use anyhow::{anyhow, bail, ensure, Context as _}; use either::Either; -use ethereum_types::{Address, U256}; +use ethereum_types::{Address, BigEndianHash as _, U256}; use evm_arithmetization::{ generation::{mpt::AccountRlp, TrieInputs}, proof::{BlockMetadata, TrieRoots}, @@ -594,12 +593,15 @@ where }; for (k, v) in storage_written { - 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)?), + true => { + storage_mask.extend(storage.reporting_remove( + MptKey::from_hash(keccak_hash::keccak(k)), + )?) + } false => { - storage.store_int(slot, v)?; + storage.store_int_at_slot(k.into_uint(), v)?; } } } @@ -767,11 +769,9 @@ where (U256::from(LAST_BLOCK_STORAGE_POS.1), block.block_number), (U256::from(TIMESTAMP_STORAGE_POS.1), timestamp), ] { - let slot = MptKey::from_slot_position(ix); - ensure!(u != U256::zero()); - scalable_storage.store_int(slot, u)?; - scalable_trim.insert(slot); + scalable_storage.store_int_at_slot(ix, u)?; + scalable_trim.insert(MptKey::from_slot_position(ix)); } // Store previous block root hash @@ -863,12 +863,11 @@ where ), ] { let slot = MptKey::from_slot_position(ix); - beacon_trim.insert(slot); match u.is_zero() { true => beacon_trim.extend(beacon_storage.reporting_remove(slot)?), false => { - beacon_storage.store_int(slot, u)?; + beacon_storage.store_int_at_slot(ix, u)?; beacon_trim.insert(slot); } } diff --git a/trace_decoder/src/interface.rs b/trace_decoder/src/interface.rs index abe3b0af0..84d57c5b9 100644 --- a/trace_decoder/src/interface.rs +++ b/trace_decoder/src/interface.rs @@ -135,9 +135,9 @@ pub struct TxnTrace { /// hash([Address]) of storages written by the transaction, /// with their new value. + // TODO(0xaatif): this should be a map from U256 -> U256 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] pub storage_written: BTreeMap, - /// Contract code that this account has accessed or created #[serde(skip_serializing_if = "Option::is_none")] pub code_usage: Option, diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 897243944..aadc7d016 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -468,6 +468,7 @@ impl World for Type1World { self.state.root() } fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()> { + // self.storage.entry(keccak_hash::keccak(address)).or_default() todo!() } } @@ -640,9 +641,11 @@ impl StorageTrie { let bytes = self.untyped.get(key.into_nibbles()).context("no item")?; Ok(rlp::decode(bytes)?) } - pub fn store_int(&mut self, key: MptKey, value: U256) -> anyhow::Result<()> { - self.untyped - .insert(key.into_nibbles(), alloy::rlp::encode(value.compat()))?; + pub fn store_int_at_slot(&mut self, slot: U256, value: U256) -> anyhow::Result<()> { + self.untyped.insert( + MptKey::from_slot_position(slot).into_nibbles(), + alloy::rlp::encode(value.compat()), + )?; Ok(()) } pub fn store_hash(&mut self, key: MptKey, value: H256) -> anyhow::Result<()> { From df960feb932b38d598fe5c62ad32f307f543279b Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Wed, 9 Oct 2024 10:50:51 +0100 Subject: [PATCH 17/25] wip: factor out methods on StorageTrie --- trace_decoder/src/core.rs | 67 +++++++++++++------------------------- trace_decoder/src/tries.rs | 49 ++++++++++++++++++++++++---- 2 files changed, 65 insertions(+), 51 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 840db1ba3..6500c1b7a 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -743,84 +743,63 @@ where StateTrieT::StateKey: 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, - STATE_ROOT_STORAGE_POS, TIMESTAMP_STORAGE_POS, + ADDRESS_SCALABLE_L2, GLOBAL_EXIT_ROOT_ADDRESS, GLOBAL_EXIT_ROOT_STORAGE_POS, + LAST_BLOCK_STORAGE_POS, STATE_ROOT_STORAGE_POS, TIMESTAMP_STORAGE_POS, }; - if block.block_number.is_zero() { - return Err(anyhow!("Attempted to prove the Genesis block!")); - } - let scalable_storage = storage - .get_mut(&ADDRESS_SCALABLE_L2_ADDRESS_HASHED) - .context("missing scalable contract storage trie")?; - let scalable_trim = trim_storage.entry(ADDRESS_SCALABLE_L2).or_default(); + ensure!( + block.block_number != U256::zero(), + "attempted to prove the genesis block" + ); - let timestamp_slot_key = MptKey::from_slot_position(U256::from(TIMESTAMP_STORAGE_POS.1)); + let scalable_trim = trim_storage.entry(ADDRESS_SCALABLE_L2).or_default(); - let timestamp = scalable_storage - .load_int(timestamp_slot_key) + let timestamp = state_trie + .load_int(ADDRESS_SCALABLE_L2, U256::from(TIMESTAMP_STORAGE_POS.1)) .unwrap_or_default(); - let timestamp = core::cmp::max(timestamp, block.block_timestamp); - - // Store block number and largest timestamp + let timestamp = cmp::max(timestamp, block.block_timestamp); for (ix, u) in [ (U256::from(LAST_BLOCK_STORAGE_POS.1), block.block_number), (U256::from(TIMESTAMP_STORAGE_POS.1), timestamp), ] { ensure!(u != U256::zero()); - scalable_storage.store_int_at_slot(ix, u)?; + state_trie.store_int(ADDRESS_SCALABLE_L2, ix, u)?; scalable_trim.insert(MptKey::from_slot_position(ix)); } // Store previous block root hash - let prev_block_root_hash = state_trie.root(); 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 = MptKey::from_hash(keccak_hash::keccak(arr)); - scalable_storage.store_hash(slot, prev_block_root_hash)?; + let prev_block_root_hash = state_trie.root(); + state_trie.store_hash( + ADDRESS_SCALABLE_L2, + keccak_hash::keccak(arr), + prev_block_root_hash, + )?; scalable_trim.insert(slot); trim_state.insert(::from(ADDRESS_SCALABLE_L2)); - let mut scalable_acct = state_trie - .get_account_info(ADDRESS_SCALABLE_L2) - .context("missing scalable contract address")?; - scalable_acct.storage_root = scalable_storage.root(); - state_trie - .insert_account_info(ADDRESS_SCALABLE_L2, scalable_acct) - // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 - // Add an entry API - .expect("insert must succeed with the same key as a successful `get`"); // Update GER contract's storage if necessary if let Some((root, l1blockhash)) = ger_data { - let ger_storage = storage - .get_mut(&GLOBAL_EXIT_ROOT_ADDRESS_HASHED) - .context("missing GER contract storage trie")?; let ger_trim = trim_storage.entry(GLOBAL_EXIT_ROOT_ADDRESS).or_default(); 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 = MptKey::from_hash(keccak_hash::keccak(arr)); - - ger_storage.store_hash(slot, l1blockhash)?; + state_trie.store_hash( + GLOBAL_EXIT_ROOT_ADDRESS, + keccak_hash::keccak(arr), + l1blockhash, + )?; ger_trim.insert(slot); - trim_state.insert(::from(GLOBAL_EXIT_ROOT_ADDRESS)); - let mut ger_acct = state_trie - .get_account_info(GLOBAL_EXIT_ROOT_ADDRESS) - .context("missing GER contract address")?; - ger_acct.storage_root = ger_storage.root(); - state_trie - .insert_account_info(GLOBAL_EXIT_ROOT_ADDRESS, ger_acct) - // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 - // Add an entry API - .expect("insert must succeed with the same key as a successful `get`"); } Ok(()) @@ -867,7 +846,7 @@ where match u.is_zero() { true => beacon_trim.extend(beacon_storage.reporting_remove(slot)?), false => { - beacon_storage.store_int_at_slot(ix, u)?; + state_trie.store_int(BEACON_ROOTS_CONTRACT_ADDRESS, ix, u)?; beacon_trim.insert(slot); } } diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index aadc7d016..f99bafc69 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -386,6 +386,8 @@ pub trait World { fn iter_account_info(&self) -> impl Iterator + '_; fn root(&self) -> H256; fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()>; + fn store_hash(&mut self, address: Address, position: H256, value: H256) -> anyhow::Result<()>; + fn load_int(&self, address: Address, slot: U256) -> anyhow::Result; } /// Global, [`Address`] `->` [`AccountRlp`]. @@ -468,9 +470,36 @@ impl World for Type1World { self.state.root() } fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()> { - // self.storage.entry(keccak_hash::keccak(address)).or_default() - todo!() + on_storage_trie(self, address, |storage| { + storage.store_int_at_slot(slot, value) + }) + } + fn load_int(&self, address: Address, slot: U256) -> anyhow::Result { + self.storage + .get(&keccak_hash::keccak(address)) + .context("storage for address")? + .load_int(MptKey::from_slot_position(slot)) } + fn store_hash(&mut self, address: Address, position: H256, value: H256) -> anyhow::Result<()> { + on_storage_trie(self, address, |storage| { + storage.store_hash(MptKey::from_hash(position), value) + }) + } +} + +fn on_storage_trie( + world: &mut Type1World, + address: Address, + f: impl FnOnce(&mut StorageTrie) -> anyhow::Result<()>, +) -> anyhow::Result<()> { + let mut account = world.get_account_info(address).context("no such account")?; + let storage = world + .storage + .entry(keccak_hash::keccak(address)) + .or_default(); + f(storage)?; + account.storage_root = storage.root(); + world.insert_account_info(address, account) } impl From for HashedPartialTrie { @@ -520,6 +549,15 @@ impl World for Type2World { conv_hash::smt2eth(self.as_smt().root) } fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()> { + let _ = (address, slot, value); + todo!() + } + fn load_int(&self, address: Address, slot: U256) -> anyhow::Result { + let _ = (address, slot); + todo!() + } + fn store_hash(&mut self, address: Address, position: H256, value: H256) -> anyhow::Result<()> { + let _ = (address, position, value); todo!() } } @@ -637,7 +675,7 @@ impl StorageTrie { untyped: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), } } - pub fn load_int(&mut self, key: MptKey) -> anyhow::Result { + pub fn load_int(&self, key: MptKey) -> anyhow::Result { let bytes = self.untyped.get(key.into_nibbles()).context("no item")?; Ok(rlp::decode(bytes)?) } @@ -648,7 +686,7 @@ impl StorageTrie { )?; Ok(()) } - pub fn store_hash(&mut self, key: MptKey, value: H256) -> anyhow::Result<()> { + fn store_hash(&mut self, key: MptKey, value: H256) -> anyhow::Result<()> { self.untyped .insert(key.into_nibbles(), alloy::rlp::encode(value.compat()))?; Ok(()) @@ -670,9 +708,6 @@ impl StorageTrie { 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<()> { self.untyped = mpt_trie::trie_subsets::create_trie_subset( From ff4c901abf56df20670f30e5e159d41b4f944b6e Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Thu, 10 Oct 2024 15:15:45 +0100 Subject: [PATCH 18/25] wip --- trace_decoder/src/core.rs | 64 +++++++------------------------------- trace_decoder/src/tries.rs | 58 ++++++++++++++++++++++++++++------ 2 files changed, 60 insertions(+), 62 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 6500c1b7a..7c6b52cd3 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -212,23 +212,8 @@ fn start( )?; let storage = storage .into_iter() - .map(|(k, SeparateTriePreImage::Direct(v))| { - v.items() - .try_fold(StorageTrie::default(), |mut acc, (nibbles, hash_or_val)| { - let path = MptKey::from_nibbles(nibbles); - match hash_or_val { - mpt_trie::trie_ops::ValOrHash::Val(value) => { - acc.insert(path, value)?; - } - mpt_trie::trie_ops::ValOrHash::Hash(h) => { - acc.insert_hash(path, h)?; - } - }; - anyhow::Ok(acc) - }) - .map(|v| (k, v)) - }) - .collect::>()?; + .map(|(k, SeparateTriePreImage::Direct(v))| (k, StorageTrie::from(v))) + .collect(); (Either::Left(state), storage, Hash2Code::new()) } BlockTraceTriePreImages::Combined(CombinedPreImages { compact }) => { @@ -456,7 +441,6 @@ where do_pre_execution( block, ger_data, - &mut storage_tries, &mut storage_masks, &mut state_mask, &mut state_trie, @@ -595,14 +579,10 @@ where for (k, v) in storage_written { match v.is_zero() { // this is actually a delete - true => { - storage_mask.extend(storage.reporting_remove( - MptKey::from_hash(keccak_hash::keccak(k)), - )?) - } - false => { - storage.store_int_at_slot(k.into_uint(), v)?; - } + true => storage_mask.extend( + state_trie.reporting_remove_storage(addr, k.into_uint())?, + ), + false => state_trie.store_int(addr, k.into_uint(), v)?, } } acct.storage_root = storage.root(); @@ -617,7 +597,7 @@ where if self_destructed { storage_tries.remove(&keccak_hash::keccak(addr)); - state_mask.extend(state_trie.reporting_remove(addr)?) + state_mask.extend(state_trie.reporting_remove_account_info(addr)?) } } @@ -692,7 +672,6 @@ where fn do_pre_execution + Clone>( block: &BlockMetadata, ger_data: Option<(H256, H256)>, - storage: &mut BTreeMap, trim_storage: &mut BTreeMap>, trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, @@ -704,7 +683,6 @@ where if cfg!(feature = "eth_mainnet") { return do_beacon_hook( block.block_timestamp, - storage, trim_storage, block.parent_beacon_block_root, trim_state, @@ -713,14 +691,7 @@ where } if cfg!(feature = "cdk_erigon") { - return do_scalable_hook( - block, - ger_data, - storage, - trim_storage, - trim_state, - state_trie, - ); + return do_scalable_hook(block, ger_data, trim_storage, trim_state, state_trie); } Ok(()) @@ -734,7 +705,6 @@ where fn do_scalable_hook + Clone>( block: &BlockMetadata, ger_data: Option<(H256, H256)>, - storage: &mut BTreeMap, trim_storage: &mut BTreeMap>, trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, @@ -812,7 +782,6 @@ where /// before any transactions (as per the EIP). fn do_beacon_hook + Clone>( block_timestamp: U256, - storage: &mut BTreeMap, trim_storage: &mut BTreeMap>, parent_beacon_block_root: H256, trim_state: &mut BTreeSet, @@ -822,14 +791,11 @@ where StateTrieT::StateKey: From
+ Ord, { use evm_arithmetization::testing_utils::{ - BEACON_ROOTS_CONTRACT_ADDRESS, BEACON_ROOTS_CONTRACT_ADDRESS_HASHED, HISTORY_BUFFER_LENGTH, + BEACON_ROOTS_CONTRACT_ADDRESS, HISTORY_BUFFER_LENGTH, }; let timestamp_idx = block_timestamp % HISTORY_BUFFER_LENGTH.value; let root_idx = timestamp_idx + HISTORY_BUFFER_LENGTH.value; - let beacon_storage = storage - .get_mut(&BEACON_ROOTS_CONTRACT_ADDRESS_HASHED) - .context("missing beacon contract storage trie")?; let beacon_trim = trim_storage .entry(BEACON_ROOTS_CONTRACT_ADDRESS) .or_default(); @@ -844,7 +810,8 @@ where let slot = MptKey::from_slot_position(ix); match u.is_zero() { - true => beacon_trim.extend(beacon_storage.reporting_remove(slot)?), + true => beacon_trim + .extend(state_trie.reporting_remove_storage(BEACON_ROOTS_CONTRACT_ADDRESS, ix)?), false => { state_trie.store_int(BEACON_ROOTS_CONTRACT_ADDRESS, ix, u)?; beacon_trim.insert(slot); @@ -852,15 +819,6 @@ where } } trim_state.insert(::from(BEACON_ROOTS_CONTRACT_ADDRESS)); - let mut beacon_acct = state_trie - .get_account_info(BEACON_ROOTS_CONTRACT_ADDRESS) - .context("missing beacon contract address")?; - beacon_acct.storage_root = beacon_storage.root(); - state_trie - .insert_account_info(BEACON_ROOTS_CONTRACT_ADDRESS, beacon_acct) - // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 - // Add an entry API - .expect("insert must succeed with the same key as a successful `get`"); Ok(()) } diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index f99bafc69..8a2e0bdcd 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -380,7 +380,10 @@ pub trait World { ) -> anyhow::Result<()>; fn get_account_info(&self, address: Address) -> Option; /// Hacky method to workaround MPT shenanigans. - fn reporting_remove(&mut self, address: Address) -> anyhow::Result>; + fn reporting_remove_account_info( + &mut self, + address: Address, + ) -> anyhow::Result>; /// _Hash out_ parts of the (state) trie that aren't in `addresses`. fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()>; fn iter_account_info(&self) -> impl Iterator + '_; @@ -388,6 +391,11 @@ pub trait World { fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()>; fn store_hash(&mut self, address: Address, position: H256, value: H256) -> anyhow::Result<()>; fn load_int(&self, address: Address, slot: U256) -> anyhow::Result; + fn reporting_remove_storage( + &mut self, + address: Address, + slot: U256, + ) -> anyhow::Result>; } /// Global, [`Address`] `->` [`AccountRlp`]. @@ -444,7 +452,10 @@ impl World for Type1World { } /// Delete the account at `address`, returning any remaining branch on /// collapse - fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { + fn reporting_remove_account_info( + &mut self, + address: Address, + ) -> anyhow::Result> { delete_node_and_report_remaining_key_if_branch_collapsed( self.state.as_mut_hashed_partial_trie_unchecked(), MptKey::from_address(address), @@ -485,6 +496,21 @@ impl World for Type1World { storage.store_hash(MptKey::from_hash(position), value) }) } + fn reporting_remove_storage( + &mut self, + address: Address, + slot: U256, + ) -> anyhow::Result> { + let mut account = self.get_account_info(address).context("no such account")?; + let storage = self + .storage + .get_mut(&keccak_hash::keccak(address)) + .context("no such account")?; + let report = storage.reporting_remove(MptKey::from_slot_position(slot))?; + account.storage_root = storage.root(); + self.insert_account_info(address, account)?; + Ok(report) + } } fn on_storage_trie( @@ -532,7 +558,10 @@ impl World for Type2World { fn get_account_info(&self, address: Address) -> Option { self.address2state.get(&address).copied() } - fn reporting_remove(&mut self, address: Address) -> anyhow::Result> { + fn reporting_remove_account_info( + &mut self, + address: Address, + ) -> anyhow::Result> { self.address2state.remove(&address); Ok(None) } @@ -560,6 +589,14 @@ impl World for Type2World { let _ = (address, position, value); todo!() } + fn reporting_remove_storage( + &mut self, + address: Address, + slot: U256, + ) -> anyhow::Result> { + let _ = (address, slot); + todo!() + } } impl Type2World { @@ -675,11 +712,11 @@ impl StorageTrie { untyped: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), } } - pub fn load_int(&self, key: MptKey) -> anyhow::Result { + fn load_int(&self, key: MptKey) -> anyhow::Result { let bytes = self.untyped.get(key.into_nibbles()).context("no item")?; Ok(rlp::decode(bytes)?) } - pub fn store_int_at_slot(&mut self, slot: U256, value: U256) -> anyhow::Result<()> { + fn store_int_at_slot(&mut self, slot: U256, value: U256) -> anyhow::Result<()> { self.untyped.insert( MptKey::from_slot_position(slot).into_nibbles(), alloy::rlp::encode(value.compat()), @@ -702,10 +739,7 @@ impl StorageTrie { pub fn root(&self) -> H256 { self.untyped.hash() } - pub const fn as_hashed_partial_trie(&self) -> &HashedPartialTrie { - &self.untyped - } - pub fn reporting_remove(&mut self, key: MptKey) -> anyhow::Result> { + fn reporting_remove(&mut self, key: MptKey) -> anyhow::Result> { delete_node_and_report_remaining_key_if_branch_collapsed(&mut self.untyped, key) } /// _Hash out_ the parts of the trie that aren't in `paths`. @@ -718,6 +752,12 @@ impl StorageTrie { } } +impl From for StorageTrie { + fn from(untyped: HashedPartialTrie) -> Self { + Self { untyped } + } +} + impl From for HashedPartialTrie { fn from(value: StorageTrie) -> Self { value.untyped From 17df67705e609d694b208cc153e9f774259ef9e8 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 13 Oct 2024 14:40:36 +0100 Subject: [PATCH 19/25] wip --- trace_decoder/src/core.rs | 142 ++++++++++++++----------------- trace_decoder/src/observer.rs | 9 +- trace_decoder/src/tries.rs | 152 ++++++++++++++++++++++++++++------ 3 files changed, 193 insertions(+), 110 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 7c6b52cd3..7cbd5682f 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -20,7 +20,7 @@ use zk_evm_common::gwei_to_wei; use crate::{ observer::{DummyObserver, Observer}, - tries::Type2World, + tries::{Key, Type2World}, }; use crate::{ tries::{MptKey, ReceiptTrie, StorageTrie, TransactionTrie, Type1World, World}, @@ -128,41 +128,46 @@ pub fn entrypoint( before: IntraBlockTries { state, - storage, transaction, receipt, }, after, withdrawals, - }| GenerationInputs { - txn_number_before: first_txn_ix.into(), - gas_used_before: running_gas_used.into(), - gas_used_after: { - running_gas_used += gas_used; - running_gas_used.into() - }, - signed_txns: byte_code.into_iter().map(Into::into).collect(), - withdrawals, - ger_data, - tries: TrieInputs { - state_trie: match state { - Either::Left(mpt) => mpt.into(), - Either::Right(_) => todo!("evm_arithmetization accepts an SMT"), + }| { + let (state_trie, storage_tries) = match state { + Either::Left(Type1World { state, storage }) => (state, storage), + Either::Right(_) => todo!("evm_arithmetization accepts an SMT"), + }; + GenerationInputs { + txn_number_before: first_txn_ix.into(), + gas_used_before: running_gas_used.into(), + gas_used_after: { + running_gas_used += gas_used; + running_gas_used.into() }, - transactions_trie: transaction.into(), - receipts_trie: receipt.into(), - storage_tries: storage.into_iter().map(|(k, v)| (k, v.into())).collect(), - }, - trie_roots_after: after, - checkpoint_state_trie_root, - checkpoint_consolidated_hash, - contract_code: contract_code - .into_iter() - .map(|it| (keccak_hash::keccak(&it), it)) - .collect(), - block_metadata: b_meta.clone(), - block_hashes: b_hashes.clone(), - burn_addr, + signed_txns: byte_code.into_iter().map(Into::into).collect(), + withdrawals, + ger_data, + tries: TrieInputs { + state_trie: state_trie.into(), + transactions_trie: transaction.into(), + receipts_trie: receipt.into(), + storage_tries: storage_tries + .into_iter() + .map(|(k, v)| (k, v.into())) + .collect(), + }, + trie_roots_after: after, + checkpoint_state_trie_root, + checkpoint_consolidated_hash, + contract_code: contract_code + .into_iter() + .map(|it| (keccak_hash::keccak(&it), it)) + .collect(), + block_metadata: b_meta.clone(), + block_hashes: b_hashes.clone(), + burn_addr, + } }, ) .collect()) @@ -344,7 +349,6 @@ impl Batch { #[derive(Debug)] pub struct IntraBlockTries { pub state: StateTrieT, - pub storage: BTreeMap, pub transaction: TransactionTrie, pub receipt: ReceiptTrie, } @@ -353,13 +357,11 @@ 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, } @@ -379,7 +381,7 @@ fn middle + Clone>( // state at the beginning of the block mut state_trie: StateTrieT, // storage at the beginning of the block - mut storage_tries: BTreeMap, + mut __storage_tries: BTreeMap, // None represents a dummy transaction that should not increment the transaction index // all batches SHOULD not be empty batches: Vec>>, @@ -393,11 +395,11 @@ fn middle + Clone>( observer: &mut impl Observer, ) -> anyhow::Result>> where - StateTrieT::StateKey: Ord + From
, + StateTrieT::StateKey: Ord + Key, { // Initialise the storage tries. for (haddr, acct) in state_trie.iter_account_info() { - let storage = storage_tries.entry(haddr).or_insert({ + let storage = __storage_tries.entry(haddr).or_insert({ let mut it = StorageTrie::default(); it.insert_hash(MptKey::default(), acct.storage_root) .expect("empty trie insert cannot fail"); @@ -428,13 +430,12 @@ where state: state_trie.clone(), transaction: transaction_trie.clone(), receipt: receipt_trie.clone(), - storage: storage_tries.clone(), }; // 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 { @@ -499,9 +500,7 @@ where .unwrap_or((AccountRlp::default(), true)); if born { - // Empty accounts cannot have non-empty storage, - // so we can safely insert a default trie. - storage_tries.insert(keccak_hash::keccak(addr), StorageTrie::default()); + state_trie.create_storage(addr)?; } if born || just_access { @@ -530,7 +529,7 @@ where storage_written .keys() .chain(&storage_read) - .map(|it| MptKey::from_hash(keccak_hash::keccak(it))), + .map(|it| StateTrieT::StateKey::from_hash(keccak_hash::keccak(it))), ); if do_writes { @@ -567,15 +566,6 @@ where .unwrap_or(acct.code_hash); if !storage_written.is_empty() { - let storage = match born { - true => storage_tries.entry(keccak_hash::keccak(addr)).or_default(), - false => storage_tries.get_mut(&keccak_hash::keccak(addr)).context( - format!( - "missing storage trie for address {addr:x} in txn {tx_hash:x}" - ), - )?, - }; - for (k, v) in storage_written { match v.is_zero() { // this is actually a delete @@ -585,18 +575,17 @@ where false => state_trie.store_int(addr, k.into_uint(), v)?, } } - acct.storage_root = storage.root(); } state_trie.insert_account_info(addr, acct)?; - state_mask.insert(::from(addr)); + state_mask.insert(StateTrieT::StateKey::from_address(addr)); } else { // Simple state access - state_mask.insert(::from(addr)); + state_mask.insert(StateTrieT::StateKey::from_address(addr)); } if self_destructed { - storage_tries.remove(&keccak_hash::keccak(addr)); + state_trie.destroy_storage(addr)?; state_mask.extend(state_trie.reporting_remove_account_info(addr)?) } } @@ -615,7 +604,7 @@ where withdrawals: match loop_ix == loop_len { true => { for (addr, amt) in &withdrawals { - state_mask.insert(::from(*addr)); + state_mask.insert(StateTrieT::StateKey::from_address(*addr)); let mut acct = state_trie .get_account_info(*addr) .context(format!("missing address {addr:x} for withdrawal"))?; @@ -631,20 +620,14 @@ where false => vec![], }, before: { - before.state.mask(state_mask)?; + before.state.mask_accounts(state_mask)?; before.receipt.mask(batch_first_txn_ix..txn_ix)?; before.transaction.mask(batch_first_txn_ix..txn_ix)?; - let keep = storage_masks - .keys() - .map(keccak_hash::keccak) - .collect::>(); - before.storage.retain(|haddr, _| keep.contains(haddr)); + state_trie.retain_storage(storage_masks.keys().copied())?; for (addr, mask) in storage_masks { - if let Some(it) = before.storage.get_mut(&keccak_hash::keccak(addr)) { - it.mask(mask)? - } // else must have self-destructed + state_trie.mask_storage(addr, mask)? } before }, @@ -659,7 +642,6 @@ where block.block_number, batch_index, &state_trie, - &storage_tries, &transaction_trie, &receipt_trie, ) @@ -672,12 +654,12 @@ where fn do_pre_execution + Clone>( block: &BlockMetadata, ger_data: Option<(H256, H256)>, - trim_storage: &mut BTreeMap>, + trim_storage: &mut BTreeMap>, trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, ) -> anyhow::Result<()> where - StateTrieT::StateKey: From
+ Ord, + StateTrieT::StateKey: Ord + Key, { // Ethereum mainnet: EIP-4788 if cfg!(feature = "eth_mainnet") { @@ -705,12 +687,12 @@ where fn do_scalable_hook + Clone>( block: &BlockMetadata, ger_data: Option<(H256, H256)>, - trim_storage: &mut BTreeMap>, + trim_storage: &mut BTreeMap>, trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, ) -> anyhow::Result<()> where - StateTrieT::StateKey: From
+ Ord, + StateTrieT::StateKey: Ord + Key, { use evm_arithmetization::testing_utils::{ ADDRESS_SCALABLE_L2, GLOBAL_EXIT_ROOT_ADDRESS, GLOBAL_EXIT_ROOT_STORAGE_POS, @@ -735,7 +717,7 @@ where ] { ensure!(u != U256::zero()); state_trie.store_int(ADDRESS_SCALABLE_L2, ix, u)?; - scalable_trim.insert(MptKey::from_slot_position(ix)); + scalable_trim.insert(StateTrieT::StateKey::from_slot_position(ix)); } // Store previous block root hash @@ -743,7 +725,7 @@ where 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 = MptKey::from_hash(keccak_hash::keccak(arr)); + let slot = StateTrieT::StateKey::from_hash(keccak_hash::keccak(arr)); let prev_block_root_hash = state_trie.root(); state_trie.store_hash( @@ -753,7 +735,7 @@ where )?; scalable_trim.insert(slot); - trim_state.insert(::from(ADDRESS_SCALABLE_L2)); + trim_state.insert(StateTrieT::StateKey::from_address(ADDRESS_SCALABLE_L2)); // Update GER contract's storage if necessary if let Some((root, l1blockhash)) = ger_data { @@ -762,14 +744,14 @@ where 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 = MptKey::from_hash(keccak_hash::keccak(arr)); + let slot = StateTrieT::StateKey::from_hash(keccak_hash::keccak(arr)); state_trie.store_hash( GLOBAL_EXIT_ROOT_ADDRESS, keccak_hash::keccak(arr), l1blockhash, )?; ger_trim.insert(slot); - trim_state.insert(::from(GLOBAL_EXIT_ROOT_ADDRESS)); + trim_state.insert(StateTrieT::StateKey::from_address(GLOBAL_EXIT_ROOT_ADDRESS)); } Ok(()) @@ -782,13 +764,13 @@ where /// before any transactions (as per the EIP). fn do_beacon_hook + Clone>( block_timestamp: U256, - trim_storage: &mut BTreeMap>, + trim_storage: &mut BTreeMap>, parent_beacon_block_root: H256, trim_state: &mut BTreeSet, state_trie: &mut StateTrieT, ) -> anyhow::Result<()> where - StateTrieT::StateKey: From
+ Ord, + StateTrieT::StateKey: Ord + Key, { use evm_arithmetization::testing_utils::{ BEACON_ROOTS_CONTRACT_ADDRESS, HISTORY_BUFFER_LENGTH, @@ -807,7 +789,7 @@ where U256::from_big_endian(parent_beacon_block_root.as_bytes()), ), ] { - let slot = MptKey::from_slot_position(ix); + let slot = ::from_slot_position(ix); match u.is_zero() { true => beacon_trim @@ -818,7 +800,9 @@ where } } } - trim_state.insert(::from(BEACON_ROOTS_CONTRACT_ADDRESS)); + trim_state.insert(StateTrieT::StateKey::from_address( + BEACON_ROOTS_CONTRACT_ADDRESS, + )); Ok(()) } diff --git a/trace_decoder/src/observer.rs b/trace_decoder/src/observer.rs index f9811e87c..bba477ce5 100644 --- a/trace_decoder/src/observer.rs +++ b/trace_decoder/src/observer.rs @@ -1,10 +1,9 @@ -use std::collections::BTreeMap; use std::marker::PhantomData; -use ethereum_types::{H256, U256}; +use ethereum_types::U256; use crate::core::IntraBlockTries; -use crate::tries::{ReceiptTrie, StorageTrie, TransactionTrie}; +use crate::tries::{ReceiptTrie, TransactionTrie}; /// Observer API for the trace decoder. /// Observer is used to collect various debugging and metadata info @@ -20,7 +19,6 @@ pub trait Observer { block: U256, batch: usize, state_trie: &StateTrieT, - storage: &BTreeMap, transaction_trie: &TransactionTrie, receipt_trie: &ReceiptTrie, ); @@ -61,7 +59,6 @@ impl Observer for TriesObserver { block: U256, batch: usize, state_trie: &StateTrieT, - storage: &BTreeMap, transaction_trie: &TransactionTrie, receipt_trie: &ReceiptTrie, ) { @@ -70,7 +67,6 @@ impl Observer for TriesObserver { batch, tries: IntraBlockTries { state: state_trie.clone(), - storage: storage.clone(), transaction: transaction_trie.clone(), receipt: receipt_trie.clone(), }, @@ -105,7 +101,6 @@ impl Observer for DummyObserver { _block: U256, _batch: usize, _state_trie: &StateTrieT, - _storage: &BTreeMap, _transaction_trie: &TransactionTrie, _receipt_trie: &ReceiptTrie, ) { diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 8a2e0bdcd..76e63f62a 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -1,11 +1,15 @@ //! Principled trie types and abstractions used in this library. use core::fmt; -use std::{cmp, collections::BTreeMap, marker::PhantomData}; +use std::{ + cmp, + collections::{BTreeMap, HashSet}, + marker::PhantomData, +}; use alloy_compat::Compat as _; -use anyhow::{ensure, Context}; -use bitvec::{array::BitArray, slice::BitSlice}; +use anyhow::{ensure, Context as _}; +use bitvec::slice::BitSlice; use copyvec::CopyVec; use ethereum_types::{Address, BigEndianHash as _, H256, U256}; use evm_arithmetization::generation::mpt::AccountRlp; @@ -16,7 +20,7 @@ use u4::{AsNibbles, U4}; /// /// Portions of the trie may be _hashed out_: see [`Self::insert_hash`]. #[derive(Debug, Clone, PartialEq, Eq)] -struct TypedMpt { +pub struct TypedMpt { inner: HashedPartialTrie, _ty: PhantomData T>, } @@ -97,6 +101,12 @@ where } } +impl From> for HashedPartialTrie { + fn from(value: TypedMpt) -> Self { + value.inner + } +} + /// Bounded sequence of [`U4`], /// used as a key for [`TypedMpt`]. /// @@ -172,12 +182,6 @@ impl MptKey { } } -impl From
for MptKey { - fn from(value: Address) -> Self { - Self::from_hash(keccak_hash::keccak(value)) - } -} - #[test] fn mpt_key_into_hash() { assert_eq!(MptKey::new([]).unwrap().into_hash(), None); @@ -253,13 +257,6 @@ impl SmtKey { } } -impl From
for SmtKey { - fn from(addr: Address) -> Self { - let H256(bytes) = keccak_hash::keccak(addr); - Self::new(BitArray::<_>::new(bytes)).expect("SmtKey has room for 256 bits") - } -} - impl Ord for SmtKey { fn cmp(&self, other: &Self) -> cmp::Ordering { self.as_bitslice().cmp(other.as_bitslice()) @@ -367,6 +364,42 @@ impl From for HashedPartialTrie { } } +pub trait Key { + fn from_address(address: Address) -> Self; + fn from_hash(hash: H256) -> Self; + fn from_slot_position(ix: U256) -> Self; +} + +impl Key for SmtKey { + fn from_hash(hash: H256) -> Self { + let _ = hash; + todo!() + } + fn from_slot_position(ix: U256) -> Self { + let _ = ix; + todo!() + } + + fn from_address(address: Address) -> Self { + let _ = address; + todo!() + } +} + +impl Key for MptKey { + fn from_hash(hash: H256) -> Self { + Self::from_hash(hash) + } + + fn from_slot_position(ix: U256) -> Self { + Self::from_slot_position(ix) + } + + fn from_address(address: Address) -> Self { + Self::from_address(address) + } +} + /// The state and storage of all accounts. /// /// Some parts of the tries may be _hashed out_. @@ -385,17 +418,31 @@ pub trait World { address: Address, ) -> anyhow::Result>; /// _Hash out_ parts of the (state) trie that aren't in `addresses`. - fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()>; + fn mask_accounts( + &mut self, + addresses: impl IntoIterator, + ) -> anyhow::Result<()>; fn iter_account_info(&self) -> impl Iterator + '_; fn root(&self) -> H256; fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()>; fn store_hash(&mut self, address: Address, position: H256, value: H256) -> anyhow::Result<()>; fn load_int(&self, address: Address, slot: U256) -> anyhow::Result; + fn create_storage(&mut self, address: Address) -> anyhow::Result<()>; fn reporting_remove_storage( &mut self, address: Address, slot: U256, - ) -> anyhow::Result>; + ) -> anyhow::Result>; + fn destroy_storage(&mut self, address: Address) -> anyhow::Result<()>; + fn retain_storage( + &mut self, + addresses: impl IntoIterator, + ) -> anyhow::Result<()>; + fn mask_storage( + &mut self, + address: Address, + keys: impl IntoIterator, + ) -> anyhow::Result<()>; } /// Global, [`Address`] `->` [`AccountRlp`]. @@ -403,8 +450,8 @@ pub trait World { /// See #[derive(Debug, Clone, Default)] pub struct Type1World { - state: TypedMpt, - storage: BTreeMap, + pub state: TypedMpt, + pub storage: BTreeMap, } impl Type1World { @@ -461,7 +508,7 @@ impl World for Type1World { MptKey::from_address(address), ) } - fn mask(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()> { + fn mask_accounts(&mut self, addresses: impl IntoIterator) -> anyhow::Result<()> { let inner = mpt_trie::trie_subsets::create_trie_subset( self.state.as_hashed_partial_trie(), addresses.into_iter().map(MptKey::into_nibbles), @@ -496,6 +543,13 @@ impl World for Type1World { storage.store_hash(MptKey::from_hash(position), value) }) } + fn create_storage(&mut self, address: Address) -> anyhow::Result<()> { + let clobbered = self + .storage + .insert(keccak_hash::keccak(address), StorageTrie::default()); + ensure!(clobbered.is_none()); + Ok(()) + } fn reporting_remove_storage( &mut self, address: Address, @@ -511,6 +565,32 @@ impl World for Type1World { self.insert_account_info(address, account)?; Ok(report) } + fn destroy_storage(&mut self, address: Address) -> anyhow::Result<()> { + let removed = self.storage.remove(&keccak_hash::keccak(address)); + ensure!(removed.is_some()); + Ok(()) + } + fn retain_storage( + &mut self, + addresses: impl IntoIterator, + ) -> anyhow::Result<()> { + let haddrs = addresses + .into_iter() + .map(keccak_hash::keccak) + .collect::>(); + self.storage.retain(|haddr, _| haddrs.contains(haddr)); + Ok(()) + } + fn mask_storage( + &mut self, + address: Address, + keys: impl IntoIterator, + ) -> anyhow::Result<()> { + if let Some(storage) = self.storage.get_mut(&keccak_hash::keccak(address)) { + storage.mask(keys)?; + } + Ok(()) + } } fn on_storage_trie( @@ -565,7 +645,7 @@ impl World for Type2World { self.address2state.remove(&address); Ok(None) } - fn mask(&mut self, address: impl IntoIterator) -> anyhow::Result<()> { + fn mask_accounts(&mut self, address: impl IntoIterator) -> anyhow::Result<()> { let _ = address; Ok(()) } @@ -589,14 +669,38 @@ impl World for Type2World { let _ = (address, position, value); todo!() } + fn create_storage(&mut self, address: Address) -> anyhow::Result<()> { + let _ = address; + todo!() + } fn reporting_remove_storage( &mut self, address: Address, slot: U256, - ) -> anyhow::Result> { + ) -> anyhow::Result> { let _ = (address, slot); todo!() } + fn destroy_storage(&mut self, address: Address) -> anyhow::Result<()> { + let _ = address; + todo!() + } + fn retain_storage( + &mut self, + addresses: impl IntoIterator, + ) -> anyhow::Result<()> { + let _ = addresses; + todo!() + } + + fn mask_storage( + &mut self, + address: Address, + keys: impl IntoIterator, + ) -> anyhow::Result<()> { + let _ = (address, keys); + todo!() + } } impl Type2World { From 0b72a680e0e798eacbff526a72c7169258535cd4 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 13 Oct 2024 16:54:19 +0100 Subject: [PATCH 20/25] wip --- trace_decoder/src/core.rs | 72 +++++++++----------------------------- trace_decoder/src/tries.rs | 72 ++++++++++++++++++++++++-------------- trace_decoder/src/type1.rs | 29 +++++---------- 3 files changed, 71 insertions(+), 102 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index 7cbd5682f..b1f6ac100 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -14,7 +14,6 @@ use evm_arithmetization::{ }; use itertools::Itertools as _; use keccak_hash::H256; -use mpt_trie::partial_trie::PartialTrie as _; use nunny::NonEmpty; use zk_evm_common::gwei_to_wei; @@ -23,7 +22,7 @@ use crate::{ tries::{Key, Type2World}, }; use crate::{ - tries::{MptKey, ReceiptTrie, StorageTrie, TransactionTrie, Type1World, World}, + tries::{ReceiptTrie, TransactionTrie, Type1World, World}, BlockLevelData, BlockTrace, BlockTraceTriePreImages, CombinedPreImages, ContractCodeUsage, OtherBlockData, SeparateStorageTriesPreImage, SeparateTriePreImage, SeparateTriePreImages, TxnInfo, TxnMeta, TxnTrace, @@ -60,7 +59,7 @@ pub fn entrypoint( BlockTraceTriePreImages::Separate(_) => FatalMissingCode(true), BlockTraceTriePreImages::Combined(_) => FatalMissingCode(false), }; - let (state, storage, mut code) = start(trie_pre_images, wire_disposition)?; + let (state, mut code) = start(trie_pre_images, wire_disposition)?; code.extend(code_db); @@ -85,7 +84,6 @@ pub fn entrypoint( Either::Left(mpt) => Either::Left( middle( mpt, - storage, batch(txn_info, batch_size_hint), &mut code, &b_meta, @@ -101,7 +99,6 @@ pub fn entrypoint( Either::Right( middle( smt, - storage, batch(txn_info, batch_size_hint), &mut code, &b_meta, @@ -182,11 +179,7 @@ pub fn entrypoint( fn start( pre_images: BlockTraceTriePreImages, wire_disposition: WireDisposition, -) -> anyhow::Result<( - Either, - BTreeMap, - Hash2Code, -)> { +) -> anyhow::Result<(Either, Hash2Code)> { Ok(match pre_images { // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/401 // refactor our convoluted input types @@ -194,46 +187,30 @@ fn start( state: SeparateTriePreImage::Direct(state), storage: SeparateStorageTriesPreImage::MultipleTries(storage), }) => { - let state = state.items().try_fold( - Type1World::default(), - |mut acc, (nibbles, hash_or_val)| { - let path = MptKey::from_nibbles(nibbles); - match hash_or_val { - mpt_trie::trie_ops::ValOrHash::Val(bytes) => { - #[expect(deprecated)] // this is MPT specific - acc.insert_by_hashed_address( - path.into_hash() - .context("invalid path length in direct state trie")?, - rlp::decode(&bytes) - .context("invalid AccountRlp in direct state trie")?, - )?; - } - mpt_trie::trie_ops::ValOrHash::Hash(h) => { - acc.insert_hash_by_key(path, h)?; - } - }; - anyhow::Ok(acc) - }, + let world = Type1World::new( + state, + storage + .into_iter() + .map(|(k, SeparateTriePreImage::Direct(v))| (k, v)), )?; - let storage = storage - .into_iter() - .map(|(k, SeparateTriePreImage::Direct(v))| (k, StorageTrie::from(v))) - .collect(); - (Either::Left(state), storage, Hash2Code::new()) + (Either::Left(world), Hash2Code::new()) } BlockTraceTriePreImages::Combined(CombinedPreImages { compact }) => { let instructions = crate::wire::parse(&compact) .context("couldn't parse instructions from binary format")?; - let (state, storage, code) = match wire_disposition { + let (state, code) = match wire_disposition { WireDisposition::Type1 => { let crate::type1::Frontend { state, storage, code, } = crate::type1::frontend(instructions)?; + ( - Either::Left(state), - storage, + Either::Left(Type1World::new( + state.into(), + storage.into_iter().map(|(k, v)| (k, v.into())), + )?), Hash2Code::from_iter(code.into_iter().map(NonEmpty::into_vec)), ) } @@ -242,12 +219,11 @@ fn start( crate::type2::frontend(instructions)?; ( Either::Right(trie), - BTreeMap::new(), Hash2Code::from_iter(code.into_iter().map(NonEmpty::into_vec)), ) } }; - (state, storage, code) + (state, code) } }) } @@ -380,8 +356,6 @@ pub struct FatalMissingCode(pub bool); fn middle + Clone>( // state at the beginning of the block mut state_trie: StateTrieT, - // storage at the beginning of the block - mut __storage_tries: BTreeMap, // None represents a dummy transaction that should not increment the transaction index // all batches SHOULD not be empty batches: Vec>>, @@ -397,20 +371,6 @@ fn middle + Clone>( where StateTrieT::StateKey: Ord + Key, { - // Initialise the storage tries. - for (haddr, acct) in state_trie.iter_account_info() { - let storage = __storage_tries.entry(haddr).or_insert({ - let mut it = StorageTrie::default(); - it.insert_hash(MptKey::default(), acct.storage_root) - .expect("empty trie insert cannot fail"); - it - }); - ensure!( - storage.root() == acct.storage_root, - "inconsistent initial storage for hashed address {haddr:x}" - ) - } - // These are the per-block tries. let mut transaction_trie = TransactionTrie::new(); let mut receipt_trie = ReceiptTrie::new(); diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 76e63f62a..1c6f4005c 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -13,7 +13,10 @@ use bitvec::slice::BitSlice; use copyvec::CopyVec; 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 mpt_trie::{ + partial_trie::{HashedPartialTrie, Node, OnOrphanedHashNode, PartialTrie as _}, + trie_ops::ValOrHash, +}; use u4::{AsNibbles, U4}; /// See . @@ -30,20 +33,23 @@ impl TypedMpt { and only encoded `T`s are ever inserted"; fn new() -> Self { Self { - inner: HashedPartialTrie::new(Node::Empty), + inner: HashedPartialTrie::new_with_strategy( + Node::Empty, + OnOrphanedHashNode::CollapseToExtension, + ), _ty: PhantomData, } } /// Insert a node which represents an out-of-band sub-trie. /// /// See [module documentation](super) for more. - fn insert_hash(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { + pub fn insert_hash_by_key(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { self.inner.insert(key.into_nibbles(), hash)?; Ok(()) } /// 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<()> + pub fn insert_value_by_key(&mut self, key: MptKey, value: T) -> anyhow::Result<()> where T: rlp::Encodable + rlp::Decodable, { @@ -69,11 +75,11 @@ impl TypedMpt { fn as_mut_hashed_partial_trie_unchecked(&mut self) -> &mut HashedPartialTrie { &mut self.inner } - fn root(&self) -> H256 { + pub fn root(&self) -> H256 { self.inner.hash() } /// Note that this returns owned paths and items. - fn iter(&self) -> impl Iterator + '_ + pub fn iter(&self) -> impl Iterator + '_ where T: rlp::Decodable, { @@ -422,7 +428,6 @@ pub trait World { &mut self, addresses: impl IntoIterator, ) -> anyhow::Result<()>; - fn iter_account_info(&self) -> impl Iterator + '_; fn root(&self) -> H256; fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()>; fn store_hash(&mut self, address: Address, position: H256, value: H256) -> anyhow::Result<()>; @@ -455,18 +460,42 @@ pub struct Type1World { } impl Type1World { - pub fn new(strategy: OnOrphanedHashNode) -> Self { - Self { - state: TypedMpt { - inner: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), - _ty: PhantomData, - }, - storage: BTreeMap::new(), + pub fn new( + state: HashedPartialTrie, + storage: impl IntoIterator, + ) -> anyhow::Result { + let mut storage = storage + .into_iter() + .map(|(k, v)| (k, StorageTrie::from(v))) + .collect::>(); + let mut typed = TypedMpt::default(); + for (key, vorh) in state.items() { + let key = MptKey::from_nibbles(key); + let haddr = key.into_hash().context("invalid key length")?; + match vorh { + ValOrHash::Val(vec) => { + let acct = rlp::decode::(&vec)?; + let storage = storage.entry(haddr).or_insert_with(|| { + HashedPartialTrie::new_with_strategy( + Node::Hash(acct.storage_root), + OnOrphanedHashNode::CollapseToExtension, + ) + .into() + }); + ensure!(storage.root() == acct.storage_root); + typed.insert_value_by_key(key, acct)? + } + ValOrHash::Hash(h256) => typed.insert_hash_by_key(key, h256)?, + } } + Ok(Self { + state: typed, + storage, + }) } /// Insert a _hashed out_ part of the trie pub fn insert_hash_by_key(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { - self.state.insert_hash(key, hash) + self.state.insert_hash_by_key(key, hash) } #[deprecated = "prefer operations on `Address` where possible, as SMT support requires this"] pub fn insert_by_hashed_address( @@ -474,7 +503,8 @@ impl Type1World { key: H256, account: AccountRlp, ) -> anyhow::Result<()> { - self.state.insert(MptKey::from_hash(key), account) + self.state + .insert_value_by_key(MptKey::from_hash(key), account) } pub fn iter(&self) -> impl Iterator + '_ { self.state @@ -519,11 +549,6 @@ impl World for Type1World { }; Ok(()) } - fn iter_account_info(&self) -> impl Iterator + '_ { - self.state - .iter() - .map(|(key, rlp)| (key.into_hash().expect("key is always H256"), rlp)) - } fn root(&self) -> H256 { self.state.root() } @@ -649,11 +674,6 @@ impl World for Type2World { let _ = address; Ok(()) } - fn iter_account_info(&self) -> impl Iterator + '_ { - self.address2state - .iter() - .map(|(addr, acct)| (keccak_hash::keccak(addr), *acct)) - } fn root(&self) -> H256 { conv_hash::smt2eth(self.as_smt().root) } diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 1c028a799..99340c409 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -12,28 +12,16 @@ use mpt_trie::partial_trie::OnOrphanedHashNode; use nunny::NonEmpty; use u4::U4; -use crate::tries::{MptKey, StorageTrie, Type1World}; +use crate::tries::{MptKey, StorageTrie, TypedMpt}; use crate::wire::{Instruction, SmtLeaf}; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct Frontend { - pub state: Type1World, + pub state: TypedMpt, pub code: BTreeSet>>, pub storage: BTreeMap, } -impl Default for Frontend { - // This frontend is intended to be used with our custom `zeroTracer`, - // which covers branch-to-extension collapse edge cases. - fn default() -> Self { - Self { - state: Type1World::new(OnOrphanedHashNode::CollapseToExtension), - code: BTreeSet::new(), - storage: BTreeMap::new(), - } - } -} - pub fn frontend(instructions: impl IntoIterator) -> anyhow::Result { let executions = execute(instructions)?; ensure!( @@ -105,8 +93,9 @@ fn visit( } }, }; - #[expect(deprecated)] // this is MPT-specific code - frontend.state.insert_by_hashed_address(path, account)?; + frontend + .state + .insert_value_by_key(MptKey::from_hash(path), account)?; } } } @@ -379,7 +368,7 @@ fn finish_stack(v: &mut Vec) -> anyhow::Result { #[test] fn test_tries() { - use crate::tries::World as _; + use crate::tries::{Type1World, World as _}; for (ix, case) in serde_json::from_str::>(include_str!("cases/zero_jerigon.json")) @@ -392,9 +381,9 @@ fn test_tries() { let frontend = frontend(instructions).unwrap(); assert_eq!(case.expected_state_root, frontend.state.root()); - for (haddr, acct) in frontend.state.iter() { + for (key, acct) in frontend.state.iter() { if acct.storage_root != Type1World::default().root() { - assert!(frontend.storage.contains_key(&haddr)) + assert!(frontend.storage.contains_key(&key.into_hash().unwrap())) } } } From 26de3290970d85ec734bd63d178d49078df2a238 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 13 Oct 2024 17:43:13 +0100 Subject: [PATCH 21/25] cleanup --- trace_decoder/src/core.rs | 164 ++++++++++++++++------------------ trace_decoder/src/observer.rs | 36 ++++---- trace_decoder/src/tries.rs | 124 +++++++++++-------------- trace_decoder/src/type1.rs | 7 +- zero/src/bin/trie_diff.rs | 5 +- 5 files changed, 153 insertions(+), 183 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index b1f6ac100..a152e5a24 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -59,7 +59,7 @@ pub fn entrypoint( BlockTraceTriePreImages::Separate(_) => FatalMissingCode(true), BlockTraceTriePreImages::Combined(_) => FatalMissingCode(false), }; - let (state, mut code) = start(trie_pre_images, wire_disposition)?; + let (world, mut code) = start(trie_pre_images, wire_disposition)?; code.extend(code_db); @@ -80,10 +80,10 @@ pub fn entrypoint( *amt = gwei_to_wei(*amt) } - let batches = match state { - Either::Left(mpt) => Either::Left( + let batches = match world { + Either::Left(type1world) => Either::Left( middle( - mpt, + type1world, batch(txn_info, batch_size_hint), &mut code, &b_meta, @@ -95,17 +95,17 @@ pub fn entrypoint( .into_iter() .map(|it| it.map(Either::Left)), ), - Either::Right(smt) => { + Either::Right(type2world) => { Either::Right( middle( - smt, + type2world, batch(txn_info, batch_size_hint), &mut code, &b_meta, ger_data, withdrawals, fatal_missing_code, - &mut DummyObserver::new(), // TODO(0xaatif) + &mut DummyObserver::new(), // TODO(zero): https://github.com/0xPolygonZero/zk_evm/issues/720 )? .into_iter() .map(|it| it.map(Either::Right)), @@ -124,14 +124,14 @@ pub fn entrypoint( byte_code, before: IntraBlockTries { - state, + world, transaction, receipt, }, after, withdrawals, }| { - let (state_trie, storage_tries) = match state { + let (state, storage) = match world { Either::Left(Type1World { state, storage }) => (state, storage), Either::Right(_) => todo!("evm_arithmetization accepts an SMT"), }; @@ -146,13 +146,10 @@ pub fn entrypoint( withdrawals, ger_data, tries: TrieInputs { - state_trie: state_trie.into(), + state_trie: state.into(), transactions_trie: transaction.into(), receipts_trie: receipt.into(), - storage_tries: storage_tries - .into_iter() - .map(|(k, v)| (k, v.into())) - .collect(), + storage_tries: storage.into_iter().map(|(k, v)| (k, v.into())).collect(), }, trie_roots_after: after, checkpoint_state_trie_root, @@ -198,7 +195,7 @@ fn start( BlockTraceTriePreImages::Combined(CombinedPreImages { compact }) => { let instructions = crate::wire::parse(&compact) .context("couldn't parse instructions from binary format")?; - let (state, code) = match wire_disposition { + let (world, code) = match wire_disposition { WireDisposition::Type1 => { let crate::type1::Frontend { state, @@ -223,7 +220,7 @@ fn start( ) } }; - (state, code) + (world, code) } }) } @@ -282,7 +279,7 @@ fn test_batch() { } #[derive(Debug)] -struct Batch { +struct Batch { pub first_txn_ix: usize, pub gas_used: u64, /// See [`GenerationInputs::contract_code`]. @@ -290,7 +287,7 @@ struct Batch { /// For each transaction in batch, in order. pub byte_code: Vec>>, - pub before: IntraBlockTries, + pub before: IntraBlockTries, pub after: TrieRoots, /// Empty for all but the final batch @@ -323,8 +320,8 @@ impl Batch { /// [`evm_arithmetization::generation::TrieInputs`], /// generic over state trie representation. #[derive(Debug)] -pub struct IntraBlockTries { - pub state: StateTrieT, +pub struct IntraBlockTries { + pub world: WorldT, pub transaction: TransactionTrie, pub receipt: ReceiptTrie, } @@ -332,12 +329,12 @@ pub struct IntraBlockTries { impl IntraBlockTries { fn map(self, mut f: impl FnMut(T) -> U) -> IntraBlockTries { let Self { - state, + world, transaction, receipt, } = self; IntraBlockTries { - state: f(state), + world: f(world), transaction, receipt, } @@ -353,9 +350,9 @@ pub struct FatalMissingCode(pub bool); /// Does the main work mentioned in the [module documentation](super). #[allow(clippy::too_many_arguments)] -fn middle + Clone>( +fn middle + Clone>( // state at the beginning of the block - mut state_trie: StateTrieT, + mut world: WorldT, // None represents a dummy transaction that should not increment the transaction index // all batches SHOULD not be empty batches: Vec>>, @@ -366,10 +363,10 @@ fn middle + Clone>( mut withdrawals: Vec<(Address, U256)>, fatal_missing_code: FatalMissingCode, // called with the untrimmed tries after each batch - observer: &mut impl Observer, -) -> anyhow::Result>> + observer: &mut impl Observer, +) -> anyhow::Result>> where - StateTrieT::StateKey: Ord + Key, + WorldT::Key: Ord + Key, { // These are the per-block tries. let mut transaction_trie = TransactionTrie::new(); @@ -387,7 +384,7 @@ where let mut batch_contract_code = BTreeSet::from([vec![]]); // always include empty code let mut before = IntraBlockTries { - state: state_trie.clone(), + world: world.clone(), transaction: transaction_trie.clone(), receipt: receipt_trie.clone(), }; @@ -395,8 +392,8 @@ where // 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 state_mask = BTreeSet::::new(); + let mut storage_masks = BTreeMap::<_, BTreeSet>::new(); + let mut state_mask = BTreeSet::::new(); if txn_ix == 0 { do_pre_execution( @@ -404,7 +401,7 @@ where ger_data, &mut storage_masks, &mut state_mask, - &mut state_trie, + &mut world, )?; } @@ -424,7 +421,7 @@ where if let Ok(nonempty) = nunny::Vec::new(byte_code) { batch_byte_code.push(nonempty.clone()); - transaction_trie.insert(txn_ix, nonempty.into())?; + transaction_trie.insert_value_by_index(txn_ix, nonempty.into())?; receipt_trie.insert( txn_ix, map_receipt_bytes(new_receipt_trie_node_byte.clone())?, @@ -454,17 +451,17 @@ where .map_err(|e| anyhow!("{e:?}")) .context(format!("couldn't decode receipt in txn {tx_hash:x}"))?; - let (mut acct, born) = state_trie + let (mut acct, born) = world .get_account_info(addr) .map(|acct| (acct, false)) .unwrap_or((AccountRlp::default(), true)); if born { - state_trie.create_storage(addr)?; + world.create_storage(addr)?; } if born || just_access { - state_trie + world .clone() .insert_account_info(addr, acct) .context(format!( @@ -489,7 +486,7 @@ where storage_written .keys() .chain(&storage_read) - .map(|it| StateTrieT::StateKey::from_hash(keccak_hash::keccak(it))), + .map(|it| WorldT::Key::from_hash(keccak_hash::keccak(it))), ); if do_writes { @@ -529,24 +526,23 @@ where for (k, v) in storage_written { match v.is_zero() { // this is actually a delete - true => storage_mask.extend( - state_trie.reporting_remove_storage(addr, k.into_uint())?, - ), - false => state_trie.store_int(addr, k.into_uint(), v)?, + true => storage_mask + .extend(world.reporting_remove_storage(addr, k.into_uint())?), + false => world.store_int(addr, k.into_uint(), v)?, } } } - state_trie.insert_account_info(addr, acct)?; - state_mask.insert(StateTrieT::StateKey::from_address(addr)); + world.insert_account_info(addr, acct)?; + state_mask.insert(WorldT::Key::from_address(addr)); } else { // Simple state access - state_mask.insert(StateTrieT::StateKey::from_address(addr)); + state_mask.insert(WorldT::Key::from_address(addr)); } if self_destructed { - state_trie.destroy_storage(addr)?; - state_mask.extend(state_trie.reporting_remove_account_info(addr)?) + world.destroy_storage(addr)?; + state_mask.extend(world.reporting_remove_account_info(addr)?) } } @@ -564,35 +560,35 @@ where withdrawals: match loop_ix == loop_len { true => { for (addr, amt) in &withdrawals { - state_mask.insert(StateTrieT::StateKey::from_address(*addr)); - let mut acct = state_trie + state_mask.insert(WorldT::Key::from_address(*addr)); + let mut acct = world .get_account_info(*addr) .context(format!("missing address {addr:x} for withdrawal"))?; acct.balance += *amt; - state_trie + world .insert_account_info(*addr, acct) // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 // Add an entry API - .expect("insert must succeed with the same key as a successful `get`"); + .expect("insert must succeed with the same key as a successful `get_account_info`"); } mem::take(&mut withdrawals) } false => vec![], }, before: { - before.state.mask_accounts(state_mask)?; + before.world.mask_accounts(state_mask)?; before.receipt.mask(batch_first_txn_ix..txn_ix)?; before.transaction.mask(batch_first_txn_ix..txn_ix)?; - state_trie.retain_storage(storage_masks.keys().copied())?; + world.retain_storage(storage_masks.keys().copied())?; for (addr, mask) in storage_masks { - state_trie.mask_storage(addr, mask)? + world.mask_storage(addr, mask)? } before }, after: TrieRoots { - state_root: state_trie.root(), + state_root: world.root(), transactions_root: transaction_trie.root(), receipts_root: receipt_trie.root(), }, @@ -601,7 +597,7 @@ where observer.collect_tries( block.block_number, batch_index, - &state_trie, + &world, &transaction_trie, &receipt_trie, ) @@ -611,15 +607,15 @@ where } /// Performs all the pre-txn execution rules of the targeted network. -fn do_pre_execution + Clone>( +fn do_pre_execution + Clone>( block: &BlockMetadata, ger_data: Option<(H256, H256)>, - trim_storage: &mut BTreeMap>, - trim_state: &mut BTreeSet, - state_trie: &mut StateTrieT, + trim_storage: &mut BTreeMap>, + trim_state: &mut BTreeSet, + state_trie: &mut WorldT, ) -> anyhow::Result<()> where - StateTrieT::StateKey: Ord + Key, + WorldT::Key: Ord + Key, { // Ethereum mainnet: EIP-4788 if cfg!(feature = "eth_mainnet") { @@ -644,15 +640,15 @@ where /// /// This is Polygon-CDK-specific, and runs at the start of the block, /// before any transactions (as per the Etrog specification). -fn do_scalable_hook + Clone>( +fn do_scalable_hook + Clone>( block: &BlockMetadata, ger_data: Option<(H256, H256)>, - trim_storage: &mut BTreeMap>, - trim_state: &mut BTreeSet, - state_trie: &mut StateTrieT, + trim_storage: &mut BTreeMap>, + trim_state: &mut BTreeSet, + world: &mut WorldT, ) -> anyhow::Result<()> where - StateTrieT::StateKey: Ord + Key, + WorldT::Key: Ord + Key, { use evm_arithmetization::testing_utils::{ ADDRESS_SCALABLE_L2, GLOBAL_EXIT_ROOT_ADDRESS, GLOBAL_EXIT_ROOT_STORAGE_POS, @@ -666,7 +662,7 @@ where let scalable_trim = trim_storage.entry(ADDRESS_SCALABLE_L2).or_default(); - let timestamp = state_trie + let timestamp = world .load_int(ADDRESS_SCALABLE_L2, U256::from(TIMESTAMP_STORAGE_POS.1)) .unwrap_or_default(); let timestamp = cmp::max(timestamp, block.block_timestamp); @@ -676,8 +672,8 @@ where (U256::from(TIMESTAMP_STORAGE_POS.1), timestamp), ] { ensure!(u != U256::zero()); - state_trie.store_int(ADDRESS_SCALABLE_L2, ix, u)?; - scalable_trim.insert(StateTrieT::StateKey::from_slot_position(ix)); + world.store_int(ADDRESS_SCALABLE_L2, ix, u)?; + scalable_trim.insert(WorldT::Key::from_slot_position(ix)); } // Store previous block root hash @@ -685,17 +681,17 @@ where 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 = StateTrieT::StateKey::from_hash(keccak_hash::keccak(arr)); + let slot = WorldT::Key::from_hash(keccak_hash::keccak(arr)); - let prev_block_root_hash = state_trie.root(); - state_trie.store_hash( + let prev_block_root_hash = world.root(); + world.store_hash( ADDRESS_SCALABLE_L2, keccak_hash::keccak(arr), prev_block_root_hash, )?; scalable_trim.insert(slot); - trim_state.insert(StateTrieT::StateKey::from_address(ADDRESS_SCALABLE_L2)); + trim_state.insert(WorldT::Key::from_address(ADDRESS_SCALABLE_L2)); // Update GER contract's storage if necessary if let Some((root, l1blockhash)) = ger_data { @@ -704,14 +700,14 @@ where 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 = StateTrieT::StateKey::from_hash(keccak_hash::keccak(arr)); - state_trie.store_hash( + let slot = WorldT::Key::from_hash(keccak_hash::keccak(arr)); + world.store_hash( GLOBAL_EXIT_ROOT_ADDRESS, keccak_hash::keccak(arr), l1blockhash, )?; ger_trim.insert(slot); - trim_state.insert(StateTrieT::StateKey::from_address(GLOBAL_EXIT_ROOT_ADDRESS)); + trim_state.insert(WorldT::Key::from_address(GLOBAL_EXIT_ROOT_ADDRESS)); } Ok(()) @@ -722,15 +718,15 @@ where /// /// This is Cancun-specific, and runs at the start of the block, /// before any transactions (as per the EIP). -fn do_beacon_hook + Clone>( +fn do_beacon_hook + Clone>( block_timestamp: U256, - trim_storage: &mut BTreeMap>, + trim_storage: &mut BTreeMap>, parent_beacon_block_root: H256, - trim_state: &mut BTreeSet, - state_trie: &mut StateTrieT, + trim_state: &mut BTreeSet, + world: &mut WorldT, ) -> anyhow::Result<()> where - StateTrieT::StateKey: Ord + Key, + WorldT::Key: Ord + Key, { use evm_arithmetization::testing_utils::{ BEACON_ROOTS_CONTRACT_ADDRESS, HISTORY_BUFFER_LENGTH, @@ -749,20 +745,18 @@ where U256::from_big_endian(parent_beacon_block_root.as_bytes()), ), ] { - let slot = ::from_slot_position(ix); + let slot = ::from_slot_position(ix); match u.is_zero() { true => beacon_trim - .extend(state_trie.reporting_remove_storage(BEACON_ROOTS_CONTRACT_ADDRESS, ix)?), + .extend(world.reporting_remove_storage(BEACON_ROOTS_CONTRACT_ADDRESS, ix)?), false => { - state_trie.store_int(BEACON_ROOTS_CONTRACT_ADDRESS, ix, u)?; + world.store_int(BEACON_ROOTS_CONTRACT_ADDRESS, ix, u)?; beacon_trim.insert(slot); } } } - trim_state.insert(StateTrieT::StateKey::from_address( - BEACON_ROOTS_CONTRACT_ADDRESS, - )); + trim_state.insert(WorldT::Key::from_address(BEACON_ROOTS_CONTRACT_ADDRESS)); Ok(()) } diff --git a/trace_decoder/src/observer.rs b/trace_decoder/src/observer.rs index bba477ce5..ba06f16fb 100644 --- a/trace_decoder/src/observer.rs +++ b/trace_decoder/src/observer.rs @@ -8,7 +8,7 @@ use crate::tries::{ReceiptTrie, TransactionTrie}; /// Observer API for the trace decoder. /// Observer is used to collect various debugging and metadata info /// from the trace decoder run. -pub trait Observer { +pub trait Observer { /// Collect tries after the transaction/batch execution. /// /// Passing the arguments one by one through reference, because @@ -18,7 +18,7 @@ pub trait Observer { &mut self, block: U256, batch: usize, - state_trie: &StateTrieT, + world: &WorldT, transaction_trie: &TransactionTrie, receipt_trie: &ReceiptTrie, ); @@ -28,37 +28,37 @@ pub trait Observer { /// Tries observer collected data element - contains /// the data collected during the trace decoder processing of the batches in a /// block, one element is retrieved after every batch. -pub struct TriesObserverElement { +pub struct TriesObserverElement { /// Block where the tries are collected. pub block: U256, /// Tries were collected after trace decoder processes batch number `batch`. pub batch: usize, /// State, transaction, and receipt tries after the batch /// execution (how the trace decoder sees them). - pub tries: IntraBlockTries, + pub tries: IntraBlockTries, } /// Observer for collection of post-execution tries from the /// trace decoder run. #[derive(Debug)] -pub struct TriesObserver { +pub struct TriesObserver { /// Collected data in the observer pass - pub data: Vec>, + pub data: Vec>, } -impl TriesObserver { +impl TriesObserver { /// Create new tries collecting observer. pub fn new() -> Self { - TriesObserver:: { data: Vec::new() } + TriesObserver:: { data: Vec::new() } } } -impl Observer for TriesObserver { +impl Observer for TriesObserver { fn collect_tries( &mut self, block: U256, batch: usize, - state_trie: &StateTrieT, + world: &WorldT, transaction_trie: &TransactionTrie, receipt_trie: &ReceiptTrie, ) { @@ -66,7 +66,7 @@ impl Observer for TriesObserver { block, batch, tries: IntraBlockTries { - state: state_trie.clone(), + world: world.clone(), transaction: transaction_trie.clone(), receipt: receipt_trie.clone(), }, @@ -74,7 +74,7 @@ impl Observer for TriesObserver { } } -impl Default for TriesObserver { +impl Default for TriesObserver { fn default() -> Self { Self::new() } @@ -82,25 +82,25 @@ impl Default for TriesObserver { /// Dummy observer which does not collect any data. #[derive(Default, Debug)] -pub struct DummyObserver { - phantom: PhantomData, +pub struct DummyObserver { + phantom: PhantomData, } -impl DummyObserver { +impl DummyObserver { /// Create a new dummy observer. pub fn new() -> Self { - DummyObserver:: { + DummyObserver:: { phantom: Default::default(), } } } -impl Observer for DummyObserver { +impl Observer for DummyObserver { fn collect_tries( &mut self, _block: U256, _batch: usize, - _state_trie: &StateTrieT, + _world: &WorldT, _transaction_trie: &TransactionTrie, _receipt_trie: &ReceiptTrie, ) { diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 1c6f4005c..9cc3ba3e1 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -1,9 +1,9 @@ //! Principled trie types and abstractions used in this library. -use core::fmt; use std::{ cmp, collections::{BTreeMap, HashSet}, + fmt, marker::PhantomData, }; @@ -19,9 +19,11 @@ use mpt_trie::{ }; use u4::{AsNibbles, U4}; +/// A Merkle PATRICIA Trie, where the values are always [RLP](rlp)-encoded `T`s. +/// /// See . /// -/// Portions of the trie may be _hashed out_: see [`Self::insert_hash`]. +/// Portions of the trie may be _hashed out_: see [`Self::insert_hash_by_key`]. #[derive(Debug, Clone, PartialEq, Eq)] pub struct TypedMpt { inner: HashedPartialTrie, @@ -35,6 +37,7 @@ impl TypedMpt { Self { inner: HashedPartialTrie::new_with_strategy( Node::Empty, + // These are used in the Type1World, which requires this behaviour. OnOrphanedHashNode::CollapseToExtension, ), _ty: PhantomData, @@ -96,17 +99,6 @@ impl Default for TypedMpt { } } -impl<'a, T> IntoIterator for &'a TypedMpt -where - T: rlp::Decodable, -{ - type Item = (MptKey, T); - type IntoIter = Box + 'a>; - fn into_iter(self) -> Self::IntoIter { - Box::new(self.iter()) - } -} - impl From> for HashedPartialTrie { fn from(value: TypedMpt) -> Self { value.inner @@ -114,14 +106,14 @@ impl From> for HashedPartialTrie { } /// Bounded sequence of [`U4`], -/// used as a key for [`TypedMpt`]. +/// used as a key for [`TypedMpt`] and [`Type1World`]. /// /// Semantically equivalent to [`mpt_trie::nibbles::Nibbles`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] pub struct MptKey(CopyVec); impl fmt::Display for MptKey { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { for u in self.0 { f.write_fmt(format_args!("{:x}", u))? } @@ -141,18 +133,6 @@ impl MptKey { AsNibbles(&mut packed).pack_from_slice(&self.0); H256::from_slice(&packed) } - pub fn from_address(address: Address) -> Self { - Self::from_hash(keccak_hash::keccak(address)) - } - pub fn from_slot_position(pos: U256) -> Self { - let mut bytes = [0; 32]; - pos.to_big_endian(&mut bytes); - Self::from_hash(keccak_hash::keccak(H256::from_slice(&bytes))) - } - pub fn from_hash(H256(bytes): H256) -> Self { - Self::new(AsNibbles(bytes)).expect("32 bytes is 64 nibbles, which fits") - } - pub fn from_txn_ix(txn_ix: usize) -> Self { MptKey::new(AsNibbles(rlp::encode(&txn_ix))).expect( "\ @@ -200,7 +180,7 @@ fn mpt_key_into_hash() { } /// Bounded sequence of bits, -/// used as a key for [`StateSmt`]. +/// used as a key for [`Type2World`]. /// /// Semantically equivalent to [`smt_trie::bits::Bits`]. #[derive(Clone, Copy)] @@ -292,7 +272,11 @@ impl TransactionTrie { pub fn new() -> Self { Self::default() } - pub fn insert(&mut self, txn_ix: usize, val: Vec) -> anyhow::Result>> { + pub fn insert_value_by_index( + &mut self, + txn_ix: usize, + val: Vec, + ) -> anyhow::Result>> { let prev = self .untyped .get(MptKey::from_txn_ix(txn_ix).into_nibbles()) @@ -370,6 +354,7 @@ impl From for HashedPartialTrie { } } +/// [`World::Key`]-agnostic operations. pub trait Key { fn from_address(address: Address) -> Self; fn from_hash(hash: H256) -> Self; @@ -394,15 +379,18 @@ impl Key for SmtKey { impl Key for MptKey { fn from_hash(hash: H256) -> Self { - Self::from_hash(hash) + let H256(bytes) = hash; + Self::new(AsNibbles(bytes)).expect("32 bytes is 64 nibbles, which fits") } fn from_slot_position(ix: U256) -> Self { - Self::from_slot_position(ix) + let mut bytes = [0; 32]; + ix.to_big_endian(&mut bytes); + Self::from_hash(keccak_hash::keccak(H256::from_slice(&bytes))) } fn from_address(address: Address) -> Self { - Self::from_address(address) + Self::from_hash(keccak_hash::keccak(address)) } } @@ -410,7 +398,7 @@ impl Key for MptKey { /// /// Some parts of the tries may be _hashed out_. pub trait World { - type StateKey; + type Key; type AccountInfo; fn insert_account_info( &mut self, @@ -418,15 +406,15 @@ pub trait World { account: Self::AccountInfo, ) -> anyhow::Result<()>; fn get_account_info(&self, address: Address) -> Option; - /// Hacky method to workaround MPT shenanigans. + /// Workaround MPT quirks. fn reporting_remove_account_info( &mut self, address: Address, - ) -> anyhow::Result>; + ) -> anyhow::Result>; /// _Hash out_ parts of the (state) trie that aren't in `addresses`. fn mask_accounts( &mut self, - addresses: impl IntoIterator, + addresses: impl IntoIterator, ) -> anyhow::Result<()>; fn root(&self) -> H256; fn store_int(&mut self, address: Address, slot: U256, value: U256) -> anyhow::Result<()>; @@ -437,7 +425,7 @@ pub trait World { &mut self, address: Address, slot: U256, - ) -> anyhow::Result>; + ) -> anyhow::Result>; fn destroy_storage(&mut self, address: Address) -> anyhow::Result<()>; fn retain_storage( &mut self, @@ -446,13 +434,11 @@ pub trait World { fn mask_storage( &mut self, address: Address, - keys: impl IntoIterator, + keys: impl IntoIterator, ) -> anyhow::Result<()>; } -/// Global, [`Address`] `->` [`AccountRlp`]. -/// -/// See +/// State and storage based on distinct [MPTs](HashedPartialTrie). #[derive(Debug, Clone, Default)] pub struct Type1World { pub state: TypedMpt, @@ -493,35 +479,15 @@ impl Type1World { storage, }) } - /// Insert a _hashed out_ part of the trie - pub fn insert_hash_by_key(&mut self, key: MptKey, hash: H256) -> anyhow::Result<()> { - self.state.insert_hash_by_key(key, hash) - } - #[deprecated = "prefer operations on `Address` where possible, as SMT support requires this"] - pub fn insert_by_hashed_address( - &mut self, - key: H256, - account: AccountRlp, - ) -> anyhow::Result<()> { - self.state - .insert_value_by_key(MptKey::from_hash(key), account) - } - pub fn iter(&self) -> impl Iterator + '_ { - self.state - .iter() - .map(|(key, rlp)| (key.into_hash().expect("key is always H256"), rlp)) - } - pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { - self.state.as_hashed_partial_trie() - } } impl World for Type1World { - type StateKey = MptKey; + type Key = MptKey; type AccountInfo = AccountRlp; fn insert_account_info(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { - #[expect(deprecated)] - self.insert_by_hashed_address(keccak_hash::keccak(address), account) + let key = keccak_hash::keccak(address); + self.state + .insert_value_by_key(MptKey::from_hash(key), account) } fn get_account_info(&self, address: Address) -> Option { self.state @@ -579,7 +545,7 @@ impl World for Type1World { &mut self, address: Address, slot: U256, - ) -> anyhow::Result> { + ) -> anyhow::Result> { let mut account = self.get_account_info(address).context("no such account")?; let storage = self .storage @@ -609,7 +575,7 @@ impl World for Type1World { fn mask_storage( &mut self, address: Address, - keys: impl IntoIterator, + keys: impl IntoIterator, ) -> anyhow::Result<()> { if let Some(storage) = self.storage.get_mut(&keccak_hash::keccak(address)) { storage.mask(keys)?; @@ -618,6 +584,7 @@ impl World for Type1World { } } +/// Update the [`AccountRlp::storage_root`] after running `f`. fn on_storage_trie( world: &mut Type1World, address: Address, @@ -654,7 +621,7 @@ pub struct Type2World { } impl World for Type2World { - type StateKey = SmtKey; + type Key = SmtKey; type AccountInfo = AccountRlp; fn insert_account_info(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { self.address2state.insert(address, account); @@ -716,7 +683,7 @@ impl World for Type2World { fn mask_storage( &mut self, address: Address, - keys: impl IntoIterator, + keys: impl IntoIterator, ) -> anyhow::Result<()> { let _ = (address, keys); todo!() @@ -826,14 +793,17 @@ mod conv_hash { /// Global, per-account. /// /// See -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub struct StorageTrie { untyped: HashedPartialTrie, } impl StorageTrie { - pub fn new(strategy: OnOrphanedHashNode) -> Self { + pub fn new() -> Self { Self { - untyped: HashedPartialTrie::new_with_strategy(Node::Empty, strategy), + untyped: HashedPartialTrie::new_with_strategy( + Node::Empty, + OnOrphanedHashNode::CollapseToExtension, + ), } } fn load_int(&self, key: MptKey) -> anyhow::Result { @@ -852,7 +822,7 @@ impl StorageTrie { .insert(key.into_nibbles(), alloy::rlp::encode(value.compat()))?; Ok(()) } - pub fn insert(&mut self, key: MptKey, value: Vec) -> anyhow::Result<()> { + pub fn insert_value(&mut self, key: MptKey, value: Vec) -> anyhow::Result<()> { self.untyped.insert(key.into_nibbles(), value)?; Ok(()) } @@ -867,7 +837,7 @@ impl StorageTrie { delete_node_and_report_remaining_key_if_branch_collapsed(&mut self.untyped, key) } /// _Hash out_ the parts of the trie that aren't in `paths`. - pub fn mask(&mut self, paths: impl IntoIterator) -> anyhow::Result<()> { + fn mask(&mut self, paths: impl IntoIterator) -> anyhow::Result<()> { self.untyped = mpt_trie::trie_subsets::create_trie_subset( &self.untyped, paths.into_iter().map(MptKey::into_nibbles), @@ -876,6 +846,12 @@ impl StorageTrie { } } +impl Default for StorageTrie { + fn default() -> Self { + Self::new() + } +} + impl From for StorageTrie { fn from(untyped: HashedPartialTrie) -> Self { Self { untyped } diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 99340c409..bdb0fa413 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -8,11 +8,10 @@ use anyhow::{bail, ensure, Context as _}; use either::Either; use evm_arithmetization::generation::mpt::AccountRlp; use keccak_hash::H256; -use mpt_trie::partial_trie::OnOrphanedHashNode; use nunny::NonEmpty; use u4::U4; -use crate::tries::{MptKey, StorageTrie, TypedMpt}; +use crate::tries::{Key as _, MptKey, StorageTrie, TypedMpt}; use crate::wire::{Instruction, SmtLeaf}; #[derive(Debug, Clone, Default)] @@ -133,7 +132,7 @@ fn node2storagetrie(node: Node) -> anyhow::Result { } Node::Leaf(Leaf { key, value }) => { match value { - Either::Left(Value { raw_value }) => mpt.insert( + Either::Left(Value { raw_value }) => mpt.insert_value( MptKey::new(path.iter().copied().chain(key))?, rlp::encode(&raw_value.as_slice()).to_vec(), )?, @@ -160,7 +159,7 @@ fn node2storagetrie(node: Node) -> anyhow::Result { Ok(()) } - let mut mpt = StorageTrie::new(OnOrphanedHashNode::CollapseToExtension); + let mut mpt = StorageTrie::new(); visit(&mut mpt, &stackstack::Stack::new(), node)?; Ok(mpt) } diff --git a/zero/src/bin/trie_diff.rs b/zero/src/bin/trie_diff.rs index c211cc528..8986f96d0 100644 --- a/zero/src/bin/trie_diff.rs +++ b/zero/src/bin/trie_diff.rs @@ -147,9 +147,10 @@ async fn main() -> Result<()> { &DebugOutputTries { state_trie: observer.data[prover_tries.batch_index] .tries + .world .state - .as_hashed_partial_trie() - .clone(), + .clone() + .into(), transaction_trie: observer.data[prover_tries.batch_index] .tries .transaction From e6adadbaa413ff5d1bb0fd03ce764c4ab4581b27 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sat, 12 Oct 2024 01:05:58 +0100 Subject: [PATCH 22/25] 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 9cc3ba3e1..ceea03364 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -735,6 +735,11 @@ impl Type2World { } 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; @@ -756,8 +761,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); @@ -769,7 +773,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(), @@ -779,7 +783,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)), From 7567b4968da1dbb478e31ed4efc2d3173b00b10b Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 13 Oct 2024 20:15:03 +0100 Subject: [PATCH 23/25] done? --- trace_decoder/src/core.rs | 52 +++++++++--------- trace_decoder/src/tries.rs | 107 +++++++++++++++++++++++++++++-------- trace_decoder/src/type2.rs | 24 ++++----- 3 files changed, 125 insertions(+), 58 deletions(-) diff --git a/trace_decoder/src/core.rs b/trace_decoder/src/core.rs index a152e5a24..a106caf56 100644 --- a/trace_decoder/src/core.rs +++ b/trace_decoder/src/core.rs @@ -8,7 +8,7 @@ use anyhow::{anyhow, bail, ensure, Context as _}; use either::Either; use ethereum_types::{Address, BigEndianHash as _, U256}; use evm_arithmetization::{ - generation::{mpt::AccountRlp, TrieInputs}, + generation::TrieInputs, proof::{BlockMetadata, TrieRoots}, GenerationInputs, }; @@ -19,7 +19,7 @@ use zk_evm_common::gwei_to_wei; use crate::{ observer::{DummyObserver, Observer}, - tries::{Key, Type2World}, + tries::{Account, Key, Type2World}, }; use crate::{ tries::{ReceiptTrie, TransactionTrie, Type1World, World}, @@ -212,7 +212,7 @@ fn start( ) } WireDisposition::Type2 => { - let crate::type2::Frontend { trie, code } = + let crate::type2::Frontend { world: trie, code } = crate::type2::frontend(instructions)?; ( Either::Right(trie), @@ -350,7 +350,7 @@ pub struct FatalMissingCode(pub bool); /// Does the main work mentioned in the [module documentation](super). #[allow(clippy::too_many_arguments)] -fn middle + Clone>( +fn middle( // state at the beginning of the block mut world: WorldT, // None represents a dummy transaction that should not increment the transaction index @@ -367,6 +367,7 @@ fn middle + Clone>( ) -> anyhow::Result>> where WorldT::Key: Ord + Key, + WorldT::Account: Account + Default + Clone, { // These are the per-block tries. let mut transaction_trie = TransactionTrie::new(); @@ -454,7 +455,7 @@ where let (mut acct, born) = world .get_account_info(addr) .map(|acct| (acct, false)) - .unwrap_or((AccountRlp::default(), true)); + .unwrap_or((WorldT::Account::default(), true)); if born { world.create_storage(addr)?; @@ -463,7 +464,7 @@ where if born || just_access { world .clone() - .insert_account_info(addr, acct) + .insert_account_info(addr, acct.clone()) .context(format!( "couldn't reach state of {} address {addr:x} in txn {tx_hash:x}", match born { @@ -490,37 +491,40 @@ where ); if do_writes { - acct.balance = balance.unwrap_or(acct.balance); - acct.nonce = nonce.unwrap_or(acct.nonce); - acct.code_hash = code_usage - .map(|it| match it { + *acct.balance_mut() = balance.unwrap_or(acct.balance()); + *acct.nonce_mut() = nonce.unwrap_or(acct.nonce()); + if let Some(usage) = code_usage { + let (hash, len) = match usage { + // TODO(Nashtare): https://github.com/0xPolygonZero/zk_evm/issues/700 + // This is a bug in the zero tracer, which shouldn't be giving us + // this read at all. Workaround for now. ContractCodeUsage::Read(hash) => { - // TODO(Nashtare): https://github.com/0xPolygonZero/zk_evm/issues/700 - // This is a bug in the zero tracer, which shouldn't be giving us - // this read at all. Workaround for now. match (fatal_missing_code, code.get(hash)) { (FatalMissingCode(true), None) => { bail!("no code for hash {hash:x}") } (_, Some(byte_code)) => { + let len = U256::from(byte_code.len()); batch_contract_code.insert(byte_code); + (hash, Some(len)) } (_, None) => { - log::warn!("no code for {hash:x}") + log::warn!("no code for {hash:x}"); + (hash, None) } } - - anyhow::Ok(hash) } ContractCodeUsage::Write(bytes) => { code.insert(bytes.clone()); let hash = keccak_hash::keccak(&bytes); + let len = U256::from(bytes.len()); batch_contract_code.insert(bytes); - Ok(hash) + (hash, Some(len)) } - }) - .transpose()? - .unwrap_or(acct.code_hash); + }; + acct.set_code(hash); + acct.set_code_length(len); + } if !storage_written.is_empty() { for (k, v) in storage_written { @@ -564,7 +568,7 @@ where let mut acct = world .get_account_info(*addr) .context(format!("missing address {addr:x} for withdrawal"))?; - acct.balance += *amt; + *acct.balance_mut() += *amt; world .insert_account_info(*addr, acct) // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 @@ -607,7 +611,7 @@ where } /// Performs all the pre-txn execution rules of the targeted network. -fn do_pre_execution + Clone>( +fn do_pre_execution( block: &BlockMetadata, ger_data: Option<(H256, H256)>, trim_storage: &mut BTreeMap>, @@ -640,7 +644,7 @@ where /// /// This is Polygon-CDK-specific, and runs at the start of the block, /// before any transactions (as per the Etrog specification). -fn do_scalable_hook + Clone>( +fn do_scalable_hook( block: &BlockMetadata, ger_data: Option<(H256, H256)>, trim_storage: &mut BTreeMap>, @@ -718,7 +722,7 @@ where /// /// This is Cancun-specific, and runs at the start of the block, /// before any transactions (as per the EIP). -fn do_beacon_hook + Clone>( +fn do_beacon_hook( block_timestamp: U256, trim_storage: &mut BTreeMap>, parent_beacon_block_root: H256, diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index ceea03364..96aecd115 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -19,6 +19,68 @@ use mpt_trie::{ }; use u4::{AsNibbles, U4}; +pub type Type1Account = AccountRlp; + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] +pub struct Type2Account { + pub balance: U256, + pub nonce: U256, + pub code: U256, + pub code_length: U256, + pub storage: BTreeMap, +} + +pub trait Account { + fn balance(&self) -> U256; + fn balance_mut(&mut self) -> &mut U256; + fn nonce(&self) -> U256; + fn nonce_mut(&mut self) -> &mut U256; + fn set_code(&mut self, hash: H256); + fn set_code_length(&mut self, length: Option); +} + +impl Account for Type1Account { + fn balance(&self) -> U256 { + self.balance + } + fn balance_mut(&mut self) -> &mut U256 { + &mut self.balance + } + fn nonce(&self) -> U256 { + self.nonce + } + fn nonce_mut(&mut self) -> &mut U256 { + &mut self.nonce + } + fn set_code(&mut self, hash: H256) { + self.code_hash = hash + } + fn set_code_length(&mut self, length: Option) { + let _ = length; + } +} +impl Account for Type2Account { + fn balance(&self) -> U256 { + self.balance + } + fn balance_mut(&mut self) -> &mut U256 { + &mut self.balance + } + fn nonce(&self) -> U256 { + self.nonce + } + fn nonce_mut(&mut self) -> &mut U256 { + &mut self.nonce + } + fn set_code(&mut self, hash: H256) { + self.code = hash.into_uint() + } + fn set_code_length(&mut self, length: Option) { + if let Some(length) = length { + self.code_length = length + } + } +} /// A Merkle PATRICIA Trie, where the values are always [RLP](rlp)-encoded `T`s. /// /// See . @@ -399,13 +461,13 @@ impl Key for MptKey { /// Some parts of the tries may be _hashed out_. pub trait World { type Key; - type AccountInfo; + type Account; fn insert_account_info( &mut self, address: Address, - account: Self::AccountInfo, + account: Self::Account, ) -> anyhow::Result<()>; - fn get_account_info(&self, address: Address) -> Option; + fn get_account_info(&self, address: Address) -> Option; /// Workaround MPT quirks. fn reporting_remove_account_info( &mut self, @@ -483,7 +545,7 @@ impl Type1World { impl World for Type1World { type Key = MptKey; - type AccountInfo = AccountRlp; + type Account = AccountRlp; fn insert_account_info(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { let key = keccak_hash::keccak(address); self.state @@ -616,19 +678,23 @@ impl From for HashedPartialTrie { // - it documents a requirement that `set_hash` is called before `set`. #[derive(Clone, Debug)] pub struct Type2World { - address2state: BTreeMap, + address2state: BTreeMap, hashed_out: BTreeMap, } impl World for Type2World { type Key = SmtKey; - type AccountInfo = AccountRlp; - fn insert_account_info(&mut self, address: Address, account: AccountRlp) -> anyhow::Result<()> { + type Account = Type2Account; + fn insert_account_info( + &mut self, + address: Address, + account: Type2Account, + ) -> anyhow::Result<()> { self.address2state.insert(address, account); Ok(()) } - fn get_account_info(&self, address: Address) -> Option { - self.address2state.get(&address).copied() + fn get_account_info(&self, address: Address) -> Option { + self.address2state.get(&address).cloned() } fn reporting_remove_account_info( &mut self, @@ -692,7 +758,7 @@ impl World for Type2World { impl Type2World { pub(crate) fn new_unchecked( - address2state: BTreeMap, + address2state: BTreeMap, hashed_out: BTreeMap, ) -> Self { Self { @@ -712,23 +778,22 @@ impl Type2World { } for ( addr, - AccountRlp { - nonce, + Type2Account { balance, - storage_root, - code_hash, + nonce, + code, + code_length, + storage, }, ) 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( - // 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(), - ); + smt.set(smt_trie::keys::key_code(*addr), *code); + smt.set(smt_trie::keys::key_code_length(*addr), *code_length); + for (slot, val) in storage { + smt.set(smt_trie::keys::key_storage(*addr, *slot), *val); + } } smt } diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index 4ad31fdeb..1ff0f5a23 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -5,14 +5,13 @@ use std::collections::{BTreeMap, HashSet}; use anyhow::{bail, ensure, Context as _}; use ethereum_types::{Address, U256}; -use evm_arithmetization::generation::mpt::AccountRlp; use itertools::EitherOrBoth; use keccak_hash::H256; use nunny::NonEmpty; use stackstack::Stack; use crate::{ - tries::{SmtKey, Type2World}, + tries::{SmtKey, Type2Account, Type2World}, wire::{Instruction, SmtLeaf, SmtLeafType}, }; @@ -27,7 +26,7 @@ pub struct CollatedLeaf { } pub struct Frontend { - pub trie: Type2World, + pub world: Type2World, pub code: HashSet>>, } @@ -37,7 +36,7 @@ pub struct Frontend { pub fn frontend(instructions: impl IntoIterator) -> anyhow::Result { let (node, code) = fold(instructions).context("couldn't fold smt from instructions")?; let trie = node2trie(node).context("couldn't construct trie and collation from folded node")?; - Ok(Frontend { trie, code }) + Ok(Frontend { world: trie, code }) } /// Node in a binary (SMT) tree. @@ -126,20 +125,19 @@ fn node2trie(node: Node) -> anyhow::Result { CollatedLeaf { balance, nonce, - // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/707 - // we shouldn't ignore these fields - code: _, - code_length: _, - storage: _, + code, + code_length, + storage, }, )| { ( addr, - AccountRlp { - nonce: nonce.unwrap_or_default(), + Type2Account { balance: balance.unwrap_or_default(), - storage_root: H256::zero(), - code_hash: H256::zero(), + nonce: nonce.unwrap_or_default(), + code: code.unwrap_or_default(), + code_length: code_length.unwrap_or_default(), + storage, }, ) }, From 32b61e602e414cc6bdff23dde5396fa87661dad1 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 13 Oct 2024 20:30:29 +0100 Subject: [PATCH 24/25] fix: bug --- trace_decoder/src/tries.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 96aecd115..1a393c383 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -144,6 +144,7 @@ impl TypedMpt { self.inner.hash() } /// Note that this returns owned paths and items. + #[cfg(test)] pub fn iter(&self) -> impl Iterator + '_ where T: rlp::Decodable, @@ -519,9 +520,9 @@ impl Type1World { let mut typed = TypedMpt::default(); for (key, vorh) in state.items() { let key = MptKey::from_nibbles(key); - let haddr = key.into_hash().context("invalid key length")?; match vorh { ValOrHash::Val(vec) => { + let haddr = key.into_hash().context("invalid key length")?; let acct = rlp::decode::(&vec)?; let storage = storage.entry(haddr).or_insert_with(|| { HashedPartialTrie::new_with_strategy( From 39c0f6684a3a190e8f6b99c552e95f42feb04944 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Sun, 13 Oct 2024 20:43:06 +0100 Subject: [PATCH 25/25] fix: accept clobbered storage --- trace_decoder/src/tries.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/trace_decoder/src/tries.rs b/trace_decoder/src/tries.rs index 1a393c383..d6a20e73b 100644 --- a/trace_decoder/src/tries.rs +++ b/trace_decoder/src/tries.rs @@ -601,7 +601,9 @@ impl World for Type1World { let clobbered = self .storage .insert(keccak_hash::keccak(address), StorageTrie::default()); - ensure!(clobbered.is_none()); + if clobbered.is_some() { + log::warn!("clobbered storage for {address:x}") + } Ok(()) } fn reporting_remove_storage(