From 0384d31a52ed8c4e7a91c7479586939c5112512e Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 15 Jul 2024 13:25:23 +0100 Subject: [PATCH 1/6] mark: 0xaatif/typed-mpt From 6c9a20a021de015c8c91ce697aafa40def5c581c Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 15 Jul 2024 13:41:15 +0100 Subject: [PATCH 2/6] feat: typed_mpt --- Cargo.lock | 7 + evm_arithmetization/src/generation/mpt.rs | 2 +- trace_decoder/Cargo.toml | 1 + trace_decoder/src/lib.rs | 1 + trace_decoder/src/typed_mpt.rs | 247 ++++++++++++++++++++++ 5 files changed, 257 insertions(+), 1 deletion(-) create mode 100644 trace_decoder/src/typed_mpt.rs diff --git a/Cargo.lock b/Cargo.lock index 0302d8d8d..1f535b721 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1414,6 +1414,12 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9885fa71e26b8ab7855e2ec7cae6e9b380edff76cd052e07c683a0319d51b3a2" +[[package]] +name = "copyvec" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540a4ee3e6ea3547a6b492c5e76a4b9086e6b41178061e03ca82fc385912b1f" + [[package]] name = "core-foundation" version = "0.9.4" @@ -4951,6 +4957,7 @@ dependencies = [ "bytes", "ciborium", "ciborium-io", + "copyvec", "criterion", "either", "enum-as-inner", diff --git a/evm_arithmetization/src/generation/mpt.rs b/evm_arithmetization/src/generation/mpt.rs index cb91f06c9..e23bad805 100644 --- a/evm_arithmetization/src/generation/mpt.rs +++ b/evm_arithmetization/src/generation/mpt.rs @@ -15,7 +15,7 @@ use crate::util::h2u; use crate::witness::errors::{ProgramError, ProverInputError}; use crate::Node; -#[derive(RlpEncodable, RlpDecodable, Debug)] +#[derive(RlpEncodable, RlpDecodable, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct AccountRlp { pub nonce: U256, pub balance: U256, diff --git a/trace_decoder/Cargo.toml b/trace_decoder/Cargo.toml index 4d221d400..597880beb 100644 --- a/trace_decoder/Cargo.toml +++ b/trace_decoder/Cargo.toml @@ -16,6 +16,7 @@ bitvec = { workspace = true } bytes = { workspace = true } ciborium = { workspace = true } ciborium-io = { workspace = true } +copyvec = "0.2.0" either = { workspace = true } enum-as-inner = { workspace = true } ethereum-types = { workspace = true } diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index d5603012b..f62efbe00 100644 --- a/trace_decoder/src/lib.rs +++ b/trace_decoder/src/lib.rs @@ -91,6 +91,7 @@ mod type1; #[cfg(test)] #[allow(dead_code)] mod type2; +mod typed_mpt; mod wire; use std::collections::HashMap; diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs new file mode 100644 index 000000000..25e5595c2 --- /dev/null +++ b/trace_decoder/src/typed_mpt.rs @@ -0,0 +1,247 @@ +//! Principled MPT types used in this library. + +use std::collections::BTreeMap; + +use copyvec::CopyVec; +use either::Either; +use ethereum_types::{Address, H256}; +use evm_arithmetization::generation::mpt::AccountRlp; +use mpt_trie::partial_trie::PartialTrie as _; +use u4::{AsNibbles, U4}; + +/// Bounded sequence of [`U4`], +/// used as a key for [`TypedMpt`]. +/// +/// Semantically equivalent to [`mpt_trie::nibbles::Nibbles`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TriePath(CopyVec); + +impl TriePath { + pub fn new(components: impl IntoIterator) -> anyhow::Result { + Ok(TriePath(CopyVec::try_from_iter(components)?)) + } + fn into_hash_left_padded(mut self) -> H256 { + for _ in 0..self.0.spare_capacity_mut().len() { + self.0.insert(0, U4::Dec00) + } + let mut packed = [0u8; 32]; + AsNibbles(&mut packed).pack_from_slice(&self.0); + H256::from_slice(&packed) + } + fn from_address(address: Address) -> Self { + Self::from_hash(keccak_hash::keccak(address)) + } + fn from_hash(H256(bytes): H256) -> Self { + Self::new(AsNibbles(bytes)).expect("32 bytes is 64 nibbles, which fits") + } + fn from_txn_ix(txn_ix: usize) -> Self { + TriePath::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. + 9 < 32 + ", + ) + } + fn into_nibbles(self) -> mpt_trie::nibbles::Nibbles { + let mut theirs = mpt_trie::nibbles::Nibbles::default(); + for component in self.0 { + theirs.push_nibble_back(component as u8) + } + theirs + } +} + +/// Map where keys are [up to 64 nibbles](TriePath), and values are either an +/// out-of-band [hash](H256) or an inline [`rlp::Encodable`]/[`rlp::Decodable`] +/// value. +/// +/// [Merkle Patricia Trees](https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie) +/// are _maps_, where keys are typically _sequences_ of an _alphabet_. +/// +/// Map values are typically indirect (i.e a _hash_), +/// but in this structure may be stored _inline_. +/// +/// Semantically equivalent to a [`mpt_trie::partial_trie::HashedPartialTrie`]. +#[derive(Debug, Clone)] +struct TypedMpt { + map: BTreeMap>, +} + +impl Default for TypedMpt { + fn default() -> Self { + Self::new() + } +} + +impl TypedMpt { + pub const fn new() -> Self { + Self { + map: BTreeMap::new(), + } + } + pub fn remove(&mut self, path: TriePath) -> Option> { + self.map.remove(&path) + } + pub fn insert(&mut self, path: TriePath, value: T) -> Option> { + self.map.insert(path, Either::Right(value)) + } + pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option> { + self.map.insert(path, Either::Left(hash)) + } + pub fn get(&self, path: TriePath) -> Option> { + self.map.get(&path).map(|it| it.as_ref().map_left(|it| *it)) + } + pub fn root(&self) -> H256 + where + T: rlp::Encodable, + { + self.as_hashed_partial_trie().hash() + } + pub fn values(&self) -> impl Iterator { + self.map + .iter() + .filter_map(|(k, v)| Some((*k, v.as_ref().right()?))) + } + pub fn iter(&self) -> impl Iterator)> { + self.map + .iter() + .map(|(k, v)| (*k, v.as_ref().map_left(|h| *h))) + } + pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie + where + T: rlp::Encodable, + { + let mut theirs = mpt_trie::partial_trie::HashedPartialTrie::default(); + for (path, v) in &self.map { + let nibbles = path.into_nibbles(); + match v { + Either::Left(h) => theirs.insert(nibbles, *h), + Either::Right(v) => theirs.insert(nibbles, &*rlp::encode(v)), + } + .expect("internal error in legacy MPT library") + } + theirs + } +} + +impl<'a, T> IntoIterator for &'a TypedMpt { + type Item = (TriePath, Either); + + type IntoIter = Box)> + 'a>; + + fn into_iter(self) -> Self::IntoIter { + Box::new(self.iter()) + } +} + +/// Per-block, keyed by transaction index. +/// +/// See +#[derive(Debug, Clone, Default)] +pub struct TransactionTrie { + typed: TypedMpt>, +} + +impl TransactionTrie { + pub fn insert(&mut self, txn_ix: usize, val: Vec) -> Option>> { + self.typed.insert(TriePath::from_txn_ix(txn_ix), val) + } + pub fn root(&self) -> H256 { + self.typed.root() + } + pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + self.typed.as_hashed_partial_trie() + } +} + +/// Per-block, keyed by transaction index. +/// +/// See +#[derive(Debug, Clone, Default)] +pub struct ReceiptTrie { + typed: TypedMpt>, +} + +impl ReceiptTrie { + pub fn insert(&mut self, txn_ix: usize, val: Vec) { + self.typed.insert(TriePath::from_txn_ix(txn_ix), val); + } + pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + self.typed.as_hashed_partial_trie() + } +} + +/// Global, keyed by address +/// +/// See +#[derive(Debug, Clone, Default)] +pub struct StateTrie { + typed: TypedMpt, +} + +impl StateTrie { + pub fn insert_by_address( + &mut self, + address: Address, + account: AccountRlp, + ) -> Option> { + self.insert_by_path(TriePath::from_address(address), account) + } + fn insert_by_path( + &mut self, + path: TriePath, + account: AccountRlp, + ) -> Option> { + self.typed.insert(path, account) + } + pub fn insert_branch( + &mut self, + path: TriePath, + hash: H256, + ) -> Option> { + self.typed.insert_branch(path, hash) + } + pub fn get_by_path(&self, path: TriePath) -> Option> { + self.typed.map.get(&path).copied() + } + pub fn get_by_address(&self, address: Address) -> Option> { + self.get_by_path(TriePath::from_hash(keccak_hash::keccak(address))) + } + pub fn root(&self) -> H256 { + self.typed.root() + } + pub fn iter(&self) -> impl Iterator)> + '_ { + self.typed + .iter() + .map(|(path, eith)| (path, eith.map_right(|acct| *acct))) + } + pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + self.typed.as_hashed_partial_trie() + } +} + +/// Global, per-account. +/// +/// See +#[derive(Debug, Clone, Default)] +pub struct StorageTrie { + typed: TypedMpt>, +} +impl StorageTrie { + pub fn insert(&mut self, path: TriePath, value: Vec) -> Option>> { + self.typed.insert(path, value) + } + pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option>> { + self.typed.insert_branch(path, hash) + } + pub fn root(&self) -> H256 { + self.typed.root() + } + pub fn remove(&mut self, path: TriePath) -> Option>> { + self.typed.map.remove(&path) + } + pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + self.typed.as_hashed_partial_trie() + } +} From 0763bb9f58e5dc7c9a7264770bf829d636e63c10 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 15 Jul 2024 14:21:38 +0100 Subject: [PATCH 3/6] feat: type1 frontend uses typed_mpt --- Cargo.lock | 7 + trace_decoder/Cargo.toml | 1 + trace_decoder/src/decoding.rs | 4 +- trace_decoder/src/lib.rs | 14 +- trace_decoder/src/processed_block_trace.rs | 2 +- trace_decoder/src/type1.rs | 359 +++++++++------------ trace_decoder/src/typed_mpt.rs | 60 ++-- 7 files changed, 203 insertions(+), 244 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1f535b721..4ffd8fe41 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4527,6 +4527,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "stackstack" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8acdd7dbfcfb5dd6e46c63512508bf71c2043f70b8f143813ad75cb5e8a589f2" + [[package]] name = "starky" version = "0.4.0" @@ -4975,6 +4981,7 @@ dependencies = [ "serde", "serde_json", "smt_trie", + "stackstack", "thiserror", "u4", "winnow 0.6.13", diff --git a/trace_decoder/Cargo.toml b/trace_decoder/Cargo.toml index 597880beb..f52075a9e 100644 --- a/trace_decoder/Cargo.toml +++ b/trace_decoder/Cargo.toml @@ -32,6 +32,7 @@ plonky2 = { workspace = true } rlp = { workspace = true } serde = { workspace = true } smt_trie = { workspace = true } +stackstack = "0.3.0" thiserror = { workspace = true } u4 = { workspace = true } winnow = { workspace = true } diff --git a/trace_decoder/src/decoding.rs b/trace_decoder/src/decoding.rs index 40a17c887..f331373d3 100644 --- a/trace_decoder/src/decoding.rs +++ b/trace_decoder/src/decoding.rs @@ -66,7 +66,7 @@ impl std::fmt::Display for TraceParsingError { let h_slot = self.slot.map(|slot| { let mut buf = [0u8; 64]; slot.to_big_endian(&mut buf); - hash(&buf) + hash(buf) }); write!( f, @@ -379,7 +379,7 @@ impl ProcessedBlockTrace { for (slot, val) in storage_writes .iter() - .map(|(k, v)| (Nibbles::from_h256_be(hash(&k.bytes_be())), v)) + .map(|(k, v)| (Nibbles::from_h256_be(hash(k.bytes_be())), v)) { // If we are writing a zero, then we actually need to perform a delete. match val == &ZERO_STORAGE_SLOT_VAL_RLPED { diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index f62efbe00..5ef61e809 100644 --- a/trace_decoder/src/lib.rs +++ b/trace_decoder/src/lib.rs @@ -99,6 +99,7 @@ use std::collections::HashMap; use ethereum_types::{Address, U256}; use evm_arithmetization::proof::{BlockHashes, BlockMetadata}; use evm_arithmetization::GenerationInputs; +use keccak_hash::keccak as hash; use keccak_hash::H256; use mpt_trie::partial_trie::HashedPartialTrie; use serde::{Deserialize, Serialize}; @@ -338,8 +339,13 @@ pub fn entrypoint( } = type1::frontend(instructions)?; ProcessedBlockTracePreImages { tries: PartialTriePreImages { - state, - storage: storage.into_iter().collect(), + state: state.as_hashed_partial_trie(), + storage: storage + .into_iter() + .map(|(path, trie)| { + (path.into_hash_left_padded(), trie.as_hashed_partial_trie()) + }) + .collect(), }, extra_code_hash_mappings: match code.is_empty() { true => None, @@ -411,10 +417,6 @@ pub fn entrypoint( .into_txn_proof_gen_ir(other)?) } -fn hash(bytes: &[u8]) -> ethereum_types::H256 { - keccak_hash::keccak(bytes).0.into() -} - #[derive(Debug, Default)] struct PartialTriePreImages { pub state: HashedPartialTrie, diff --git a/trace_decoder/src/processed_block_trace.rs b/trace_decoder/src/processed_block_trace.rs index 7c9e89757..3463f662d 100644 --- a/trace_decoder/src/processed_block_trace.rs +++ b/trace_decoder/src/processed_block_trace.rs @@ -93,7 +93,7 @@ impl TxnInfo { nodes_used_by_txn.storage_accesses.push(( hashed_addr, storage_access_keys - .map(|k| Nibbles::from_h256_be(hash(&k.0))) + .map(|H256(bytes)| Nibbles::from_h256_be(hash(bytes))) .collect(), )); diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 5d93f6b4e..5e62f0bc6 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -1,27 +1,25 @@ //! Frontend for the witness format emitted by e.g the [`0xPolygonZero/erigon`](https://github.com/0xPolygonZero/erigon) //! Ethereum node (a.k.a "jerigon"). -use std::{ - array, - collections::{HashMap, HashSet}, - iter, -}; +use std::array; +use std::collections::{BTreeMap, BTreeSet}; use anyhow::{bail, ensure, Context as _}; use either::Either; -use mpt_trie::{ - partial_trie::{HashedPartialTrie, PartialTrie as _}, - trie_ops::ValOrHash, -}; +use evm_arithmetization::generation::mpt::AccountRlp; use nunny::NonEmpty; use u4::U4; +use crate::typed_mpt::{StateTrie, StorageTrie, TriePath}; use crate::wire::{Instruction, SmtLeaf}; +#[derive(Debug, Default, Clone)] pub struct Frontend { - pub state: HashedPartialTrie, - pub code: HashSet>>, - pub storage: HashMap, + pub state: StateTrie, + pub code: BTreeSet>>, + /// The key here matches the [`TriePath`] inside [`Self::state`] for + /// accounts which had inline storage + pub storage: BTreeMap, } pub fn frontend(instructions: impl IntoIterator) -> anyhow::Result { @@ -32,27 +30,141 @@ pub fn frontend(instructions: impl IntoIterator) -> anyhow:: ); let execution = executions.into_vec().remove(0); - let mut visitor = Visitor { - path: Vec::new(), - frontend: Frontend { - state: HashedPartialTrie::default(), - code: HashSet::new(), - storage: HashMap::new(), + let mut frontend = Frontend::default(); + visit( + &mut frontend, + &stackstack::Stack::new(), + match execution { + Execution::Leaf(it) => Node::Leaf(it), + Execution::Extension(it) => Node::Extension(it), + Execution::Branch(it) => Node::Branch(it), + Execution::Empty => Node::Empty, }, - }; - visitor.visit_node(match execution { - Execution::Leaf(it) => Node::Leaf(it), - Execution::Extension(it) => Node::Extension(it), - Execution::Branch(it) => Node::Branch(it), - Execution::Empty => Node::Empty, - })?; - let Visitor { path, frontend } = visitor; - - assert_eq!(Vec::::new(), path); + )?; Ok(frontend) } +fn visit( + frontend: &mut Frontend, + path: &stackstack::Stack<'_, U4>, + node: Node, +) -> anyhow::Result<()> { + match node { + Node::Hash(Hash { raw_hash }) => { + let clobbered = frontend + .state + .insert_branch(TriePath::new(path.iter().copied())?, raw_hash.into()); + ensure!(clobbered.is_none(), "duplicate hash") + } + Node::Leaf(Leaf { key, value }) => { + let path = TriePath::new(path.iter().copied().chain(key))?; + match value { + // TODO(0xaatif): what should this be interpreted as? + // (this branch isn't hit in our tests) + Either::Left(Value { .. }) => bail!("unsupported value node"), + Either::Right(Account { + nonce, + balance, + storage, + code, + }) => { + let account = AccountRlp { + nonce: nonce.into(), + balance, + storage_root: { + let storage = node2storagetrie(match storage { + Some(it) => *it, + None => Node::Empty, + })?; + let storage_root = storage.root(); + let clobbered = frontend.storage.insert(path, storage); + ensure!(clobbered.is_none(), "duplicate storage"); + storage_root + }, + code_hash: { + match code { + Some(Either::Left(Hash { raw_hash })) => raw_hash.into(), + Some(Either::Right(Code { code })) => { + let hash = crate::hash(&code); + frontend.code.insert(code); + hash + } + None => crate::hash([]), + } + }, + }; + let clobbered = frontend.state.insert_by_path(path, account); + ensure!(clobbered.is_none(), "duplicate account"); + } + } + } + Node::Extension(Extension { key, child }) => { + path.with_all(key, |path| visit(frontend, path, *child))? + } + Node::Branch(Branch { children }) => { + for (ix, node) in children.into_iter().enumerate() { + if let Some(node) = node { + path.with( + U4::new(ix.try_into().expect("ix is in range 0..16")) + .expect("ix is in range 0..16"), + |path| visit(frontend, path, *node), + )?; + } + } + } + Node::Code(Code { code }) => { + frontend.code.insert(code); + } + Node::Empty => {} + } + Ok(()) +} + +fn node2storagetrie(node: Node) -> anyhow::Result { + fn visit( + mpt: &mut StorageTrie, + path: &stackstack::Stack, + node: Node, + ) -> anyhow::Result<()> { + match node { + Node::Hash(Hash { raw_hash }) => { + mpt.insert_branch(TriePath::new(path.iter().copied())?, raw_hash.into()); + } + Node::Leaf(Leaf { key, value }) => { + match value { + Either::Left(Value { raw_value }) => mpt.insert( + TriePath::new(path.iter().copied().chain(key))?, + raw_value.into_vec(), + ), + Either::Right(_) => bail!("unexpected account node in storage trie"), + }; + } + Node::Extension(Extension { key, child }) => { + path.with_all(key, |path| visit(mpt, path, *child))? + } + Node::Branch(Branch { children }) => { + for (ix, node) in children.into_iter().enumerate() { + if let Some(node) = node { + path.with( + U4::new(ix.try_into().expect("ix is in range 0..16")) + .expect("ix is in range 0..16"), + |path| visit(mpt, path, *node), + )?; + } + } + } + Node::Code(_) => bail!("unexpected Code node in storage trie"), + Node::Empty => {} + } + Ok(()) + } + + let mut mpt = StorageTrie::default(); + visit(&mut mpt, &stackstack::Stack::new(), node)?; + Ok(mpt) +} + #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] struct Hash { raw_hash: [u8; 32], @@ -253,182 +365,9 @@ fn finish_stack(v: &mut Vec) -> anyhow::Result { } } -/// Visit a [`Node`], keeping track of the path and decorating the [`Frontend`] -/// as appropriate. -struct Visitor { - path: Vec, - frontend: Frontend, -} - -impl Visitor { - fn with_path( - &mut self, - path: impl IntoIterator, - f: impl FnOnce(&mut Self) -> T, - ) -> T { - let len = self.path.len(); - self.path.extend(path); - let ret = f(self); - self.path.truncate(len); - ret - } - fn visit_node(&mut self, it: Node) -> anyhow::Result<()> { - match it { - Node::Hash(Hash { raw_hash }) => { - self.frontend.state.insert( - nibbles2nibbles(self.path.clone()), - ValOrHash::Hash(raw_hash.into()), - )?; - Ok(()) - } - Node::Leaf(it) => self.visit_leaf(it), - Node::Extension(it) => self.visit_extension(it), - Node::Branch(it) => self.visit_branch(it), - Node::Code(Code { code }) => { - self.frontend.code.insert(code); - Ok(()) - } - Node::Empty => Ok(()), - } - } - fn visit_extension(&mut self, Extension { key, child }: Extension) -> anyhow::Result<()> { - self.with_path(key, |this| this.visit_node(*child)) - } - fn visit_branch(&mut self, Branch { children }: Branch) -> anyhow::Result<()> { - for (ix, node) in children.into_iter().enumerate() { - if let Some(node) = node { - self.with_path( - iter::once(U4::new(ix.try_into().unwrap()).unwrap()), - |this| this.visit_node(*node), - )? - } - } - Ok(()) - } - fn visit_leaf(&mut self, Leaf { key, value }: Leaf) -> anyhow::Result<()> { - let key = self.path.iter().copied().chain(key).collect::>(); - let value = match value { - Either::Left(Value { raw_value }) => rlp::encode(raw_value.as_vec()), - Either::Right(Account { - nonce, - balance, - storage, - code, - }) => rlp::encode(&evm_arithmetization::generation::mpt::AccountRlp { - nonce: nonce.into(), - balance, - code_hash: match code { - Some(Either::Left(Hash { raw_hash })) => raw_hash.into(), - Some(Either::Right(Code { code })) => { - let hash = crate::hash(&code); - self.frontend.code.insert(code); - hash - } - None => crate::hash(&[]), - }, - storage_root: { - let storage = node2trie(match storage { - Some(it) => *it, - None => Node::Empty, - }) - .context(format!( - "couldn't convert account storage to trie at path {:?}", - self.path - ))?; - let storage_root = storage.hash(); - self.frontend.storage.insert( - ethereum_types::H256::from_slice(&nibbles2nibbles(key.clone()).bytes_be()), - storage, - ); - storage_root - }, - }), - }; - // TODO(0xaatif): do consistency checks here. - self.frontend - .state - .insert(nibbles2nibbles(key), ValOrHash::Val(value.to_vec()))?; - Ok(()) - } -} - -/// # Panics -/// - internally in [`mpt_trie`]. -fn node2trie(node: Node) -> anyhow::Result { - let mut trie = HashedPartialTrie::default(); - for (k, v) in iter_leaves(node) { - trie.insert( - nibbles2nibbles(k), - match v { - IterLeaf::Hash(Hash { raw_hash }) => ValOrHash::Hash(raw_hash.into()), - IterLeaf::Value(Value { raw_value }) => { - ValOrHash::Val(rlp::encode(raw_value.as_vec()).to_vec()) - } - IterLeaf::Empty => continue, - IterLeaf::Account => bail!("unexpected Account node in storage trie"), - IterLeaf::Code => bail!("unexpected Code node in storage trie"), - }, - )?; - } - Ok(trie) -} - -/// # Panics -/// - If `ours` is too deep. -fn nibbles2nibbles(ours: Vec) -> mpt_trie::nibbles::Nibbles { - let mut theirs = mpt_trie::nibbles::Nibbles::default(); - for it in ours { - theirs.push_nibble_back(it as u8) - } - theirs -} - -/// Leaf in a [`Node`] tree, see [`iter_leaves`]. -enum IterLeaf { - Hash(Hash), - Value(Value), - Empty, - // we don't attach information to these variants because they're error cases - Account, - Code, -} - -/// Simple, inefficient visitor of all leaves of the [`Node`] tree. -#[allow(clippy::type_complexity)] -fn iter_leaves(node: Node) -> Box, IterLeaf)>> { - match node { - Node::Hash(it) => Box::new(iter::once((vec![], IterLeaf::Hash(it)))), - Node::Leaf(Leaf { key, value }) => match value { - Either::Left(it) => Box::new(iter::once((key.into(), IterLeaf::Value(it)))), - Either::Right(_) => Box::new(iter::once((key.into(), IterLeaf::Account))), - }, - Node::Extension(Extension { - key: parent_key, - child, - }) => Box::new(iter_leaves(*child).map(move |(child_key, v)| { - (parent_key.clone().into_iter().chain(child_key).collect(), v) - })), - Node::Branch(Branch { children }) => Box::new( - children - .into_iter() - .enumerate() - .flat_map(|(ix, child)| { - child.map(|it| (U4::new(ix.try_into().unwrap()).unwrap(), *it)) - }) - .flat_map(|(parent_key, child)| { - iter_leaves(child).map(move |(mut child_key, v)| { - child_key.insert(0, parent_key); - (child_key, v) - }) - }), - ), - Node::Code(_) => Box::new(iter::once((vec![], IterLeaf::Code))), - Node::Empty => Box::new(iter::once((vec![], IterLeaf::Empty))), - } -} - #[test] fn test() { + use mpt_trie::partial_trie::PartialTrie as _; for (ix, case) in serde_json::from_str::>(include_str!("test_cases/zero_jerigon.json")) .unwrap() @@ -438,17 +377,13 @@ fn test() { println!("case {}", ix); let instructions = crate::wire::parse(&case.bytes).unwrap(); let frontend = frontend(instructions).unwrap(); - assert_eq!(case.expected_state_root, frontend.state.hash()); + assert_eq!(case.expected_state_root, frontend.state.root()); - for (address, data) in frontend.state.items() { - if let ValOrHash::Val(bytes) = data { - let address = ethereum_types::H256::from_slice(&address.bytes_be()); - let account = - rlp::decode::(&bytes) - .unwrap(); - let storage_root = account.storage_root; - if storage_root != crate::processed_block_trace::EMPTY_TRIE_HASH { - assert!(frontend.storage.contains_key(&address)) + for (path, hash_or_acct) in frontend.state { + if let Either::Right(acct) = hash_or_acct { + if acct.storage_root != mpt_trie::partial_trie::HashedPartialTrie::default().hash() + { + assert!(frontend.storage.contains_key(&path)) } } } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 25e5595c2..7ad941c05 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -20,7 +20,7 @@ impl TriePath { pub fn new(components: impl IntoIterator) -> anyhow::Result { Ok(TriePath(CopyVec::try_from_iter(components)?)) } - fn into_hash_left_padded(mut self) -> H256 { + pub fn into_hash_left_padded(mut self) -> H256 { for _ in 0..self.0.spare_capacity_mut().len() { self.0.insert(0, U4::Dec00) } @@ -80,33 +80,26 @@ impl TypedMpt { map: BTreeMap::new(), } } - pub fn remove(&mut self, path: TriePath) -> Option> { - self.map.remove(&path) - } + // pub fn remove(&mut self, path: TriePath) -> Option> { + // self.map.remove(&path) + // } pub fn insert(&mut self, path: TriePath, value: T) -> Option> { self.map.insert(path, Either::Right(value)) } pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option> { self.map.insert(path, Either::Left(hash)) } - pub fn get(&self, path: TriePath) -> Option> { - self.map.get(&path).map(|it| it.as_ref().map_left(|it| *it)) - } + // pub fn get(&self, path: TriePath) -> Option> { + // self.map.get(&path).map(|it| it.as_ref().map_left(|it| *it)) + // } pub fn root(&self) -> H256 where T: rlp::Encodable, { self.as_hashed_partial_trie().hash() } - pub fn values(&self) -> impl Iterator { - self.map - .iter() - .filter_map(|(k, v)| Some((*k, v.as_ref().right()?))) - } - pub fn iter(&self) -> impl Iterator)> { - self.map - .iter() - .map(|(k, v)| (*k, v.as_ref().map_left(|h| *h))) + pub fn iter(&self) -> impl Iterator)> { + self.map.iter() } pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie where @@ -126,15 +119,28 @@ impl TypedMpt { } impl<'a, T> IntoIterator for &'a TypedMpt { - type Item = (TriePath, Either); + type Item = (&'a TriePath, &'a Either); - type IntoIter = Box)> + 'a>; + type IntoIter = Box)> + 'a>; fn into_iter(self) -> Self::IntoIter { Box::new(self.iter()) } } +impl IntoIterator for TypedMpt +where + T: 'static, +{ + type Item = (TriePath, Either); + + type IntoIter = Box)>>; + + fn into_iter(self) -> Self::IntoIter { + Box::new(self.map.into_iter()) + } +} + /// Per-block, keyed by transaction index. /// /// See @@ -188,7 +194,7 @@ impl StateTrie { ) -> Option> { self.insert_by_path(TriePath::from_address(address), account) } - fn insert_by_path( + pub fn insert_by_path( &mut self, path: TriePath, account: AccountRlp, @@ -211,16 +217,24 @@ impl StateTrie { pub fn root(&self) -> H256 { self.typed.root() } - pub fn iter(&self) -> impl Iterator)> + '_ { - self.typed - .iter() - .map(|(path, eith)| (path, eith.map_right(|acct| *acct))) + pub fn iter(&self) -> impl Iterator)> { + self.typed.iter() } pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { self.typed.as_hashed_partial_trie() } } +impl IntoIterator for StateTrie { + type Item = (TriePath, Either); + + type IntoIter = Box)>>; + + fn into_iter(self) -> Self::IntoIter { + self.typed.into_iter() + } +} + /// Global, per-account. /// /// See From 7807a7db38f2ba88d48fdd800a34dc715752d337 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Mon, 15 Jul 2024 16:57:33 +0100 Subject: [PATCH 4/6] fix: StorageTrie, and full frontend support --- trace_decoder/src/decoding.rs | 18 ++++++++--- trace_decoder/src/lib.rs | 56 +++++++++++++++++++++++++--------- trace_decoder/src/type1.rs | 2 +- trace_decoder/src/typed_mpt.rs | 37 +++++++++++++++++----- 4 files changed, 86 insertions(+), 27 deletions(-) diff --git a/trace_decoder/src/decoding.rs b/trace_decoder/src/decoding.rs index f331373d3..0ad829c9d 100644 --- a/trace_decoder/src/decoding.rs +++ b/trace_decoder/src/decoding.rs @@ -224,15 +224,25 @@ impl ProcessedBlockTrace { other_data: OtherBlockData, ) -> TraceParsingResult> { let mut curr_block_tries = PartialTrieState { - state: self.tries.state.clone(), - storage: self.tries.storage.clone(), + state: self.tries.state.as_hashed_partial_trie(), + storage: self + .tries + .storage + .iter() + .map(|(k, v)| (*k, v.as_hashed_partial_trie())) + .collect(), ..Default::default() }; // This is just a copy of `curr_block_tries`. let initial_tries_for_dummies = PartialTrieState { - state: self.tries.state, - storage: self.tries.storage, + state: self.tries.state.as_hashed_partial_trie(), + storage: self + .tries + .storage + .iter() + .map(|(k, v)| (*k, v.as_hashed_partial_trie())) + .collect(), ..Default::default() }; diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index 5ef61e809..f69dee624 100644 --- a/trace_decoder/src/lib.rs +++ b/trace_decoder/src/lib.rs @@ -103,6 +103,7 @@ use keccak_hash::keccak as hash; use keccak_hash::H256; use mpt_trie::partial_trie::HashedPartialTrie; use serde::{Deserialize, Serialize}; +use typed_mpt::{StateTrie, StorageTrie, TriePath}; /// Core payload needed to generate proof for a block. /// Additional data retrievable from the blockchain node (using standard ETH RPC @@ -297,7 +298,6 @@ pub fn entrypoint( resolve: impl Fn(H256) -> Vec, ) -> anyhow::Result> { use anyhow::Context as _; - use evm_arithmetization::generation::mpt::AccountRlp; use mpt_trie::partial_trie::PartialTrie as _; use crate::processed_block_trace::{ @@ -321,11 +321,42 @@ pub fn entrypoint( storage: SeparateStorageTriesPreImage::MultipleTries(storage), }) => ProcessedBlockTracePreImages { tries: PartialTriePreImages { - state, + state: state.items().try_fold( + StateTrie::default(), + |mut acc, (nibbles, hash_or_val)| { + let path = TriePath::from_nibbles(nibbles); + match hash_or_val { + mpt_trie::trie_ops::ValOrHash::Val(bytes) => acc.insert_by_path( + path, + rlp::decode(&bytes) + .context("invalid AccountRlp in direct state trie")?, + ), + mpt_trie::trie_ops::ValOrHash::Hash(h) => { + acc.insert_branch_by_path(path, h) + } + }; + anyhow::Ok(acc) + }, + )?, storage: storage .into_iter() - .map(|(k, SeparateTriePreImage::Direct(v))| (k, v)) - .collect(), + .map(|(k, SeparateTriePreImage::Direct(v))| { + v.items() + .try_fold(StorageTrie::default(), |mut acc, (nibbles, hash_or_val)| { + let path = TriePath::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_branch(path, h) + } + }; + anyhow::Ok(acc) + }) + .map(|v| (k, v)) + }) + .collect::>()?, }, extra_code_hash_mappings: None, }, @@ -339,12 +370,10 @@ pub fn entrypoint( } = type1::frontend(instructions)?; ProcessedBlockTracePreImages { tries: PartialTriePreImages { - state: state.as_hashed_partial_trie(), + state, storage: storage .into_iter() - .map(|(path, trie)| { - (path.into_hash_left_padded(), trie.as_hashed_partial_trie()) - }) + .map(|(path, trie)| (path.into_hash_left_padded(), trie)) .collect(), }, extra_code_hash_mappings: match code.is_empty() { @@ -362,11 +391,8 @@ pub fn entrypoint( let all_accounts_in_pre_images = pre_images .tries .state - .items() - .filter_map(|(addr, data)| { - data.as_val() - .map(|data| (addr.into(), rlp::decode::(data).unwrap())) - }) + .iter() + .filter_map(|(addr, data)| Some((addr.into_hash_left_padded(), data.right()?))) .collect::>(); let code_db = { @@ -419,8 +445,8 @@ pub fn entrypoint( #[derive(Debug, Default)] struct PartialTriePreImages { - pub state: HashedPartialTrie, - pub storage: HashMap, + pub state: StateTrie, + pub storage: HashMap, } /// Like `#[serde(with = "hex")`, but tolerates and emits leading `0x` prefixes diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 5e62f0bc6..f6046be7b 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -54,7 +54,7 @@ fn visit( Node::Hash(Hash { raw_hash }) => { let clobbered = frontend .state - .insert_branch(TriePath::new(path.iter().copied())?, raw_hash.into()); + .insert_branch_by_path(TriePath::new(path.iter().copied())?, raw_hash.into()); ensure!(clobbered.is_none(), "duplicate hash") } Node::Leaf(Leaf { key, value }) => { diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index 7ad941c05..f639401c7 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -50,6 +50,17 @@ impl TriePath { } theirs } + pub fn from_nibbles(mut theirs: mpt_trie::nibbles::Nibbles) -> Self { + let mut ours = CopyVec::new(); + while !theirs.is_empty() { + ours.try_push( + U4::new(theirs.pop_next_nibble_front()) + .expect("mpt_trie returned an invalid nibble"), + ) + .expect("mpt_trie should not have more than 64 nibbles") + } + Self(ours) + } } /// Map where keys are [up to 64 nibbles](TriePath), and values are either an @@ -201,7 +212,7 @@ impl StateTrie { ) -> Option> { self.typed.insert(path, account) } - pub fn insert_branch( + pub fn insert_branch_by_path( &mut self, path: TriePath, hash: H256, @@ -240,22 +251,34 @@ impl IntoIterator for StateTrie { /// See #[derive(Debug, Clone, Default)] pub struct StorageTrie { - typed: TypedMpt>, + /// This does NOT use [`TypedMpt`] - T could be anything! + map: BTreeMap>>, } impl StorageTrie { pub fn insert(&mut self, path: TriePath, value: Vec) -> Option>> { - self.typed.insert(path, value) + self.map.insert(path, Either::Right(value)) } pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option>> { - self.typed.insert_branch(path, hash) + self.map.insert(path, Either::Left(hash)) } pub fn root(&self) -> H256 { - self.typed.root() + self.as_hashed_partial_trie().hash() } pub fn remove(&mut self, path: TriePath) -> Option>> { - self.typed.map.remove(&path) + self.map.remove(&path) } pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { - self.typed.as_hashed_partial_trie() + let mut theirs = mpt_trie::partial_trie::HashedPartialTrie::default(); + for (k, v) in &self.map { + let nibbles = k.into_nibbles(); + match v { + Either::Left(h) => theirs.insert(nibbles, *h), + // TODO(0xaatif): why is RLP-encoding here the right thing to do? + // (Our tests fail without it). + Either::Right(v) => theirs.insert(nibbles, &*rlp::encode(v)), + } + .expect("error in mpt_trie") + } + theirs } } From 069a708320e1d4da27c57a26b396564cf01de640 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Wed, 17 Jul 2024 13:43:30 +0100 Subject: [PATCH 5/6] refactor: use BTreeMap> -> HashedPartialTrie --- trace_decoder/src/decoding.rs | 8 +- trace_decoder/src/lib.rs | 20 +-- trace_decoder/src/type1.rs | 20 +-- trace_decoder/src/typed_mpt.rs | 288 +++++++++++++++++---------------- 4 files changed, 170 insertions(+), 166 deletions(-) diff --git a/trace_decoder/src/decoding.rs b/trace_decoder/src/decoding.rs index 0ad829c9d..c6ac35106 100644 --- a/trace_decoder/src/decoding.rs +++ b/trace_decoder/src/decoding.rs @@ -224,24 +224,24 @@ impl ProcessedBlockTrace { other_data: OtherBlockData, ) -> TraceParsingResult> { let mut curr_block_tries = PartialTrieState { - state: self.tries.state.as_hashed_partial_trie(), + state: self.tries.state.as_hashed_partial_trie().clone(), storage: self .tries .storage .iter() - .map(|(k, v)| (*k, v.as_hashed_partial_trie())) + .map(|(k, v)| (*k, v.as_hashed_partial_trie().clone())) .collect(), ..Default::default() }; // This is just a copy of `curr_block_tries`. let initial_tries_for_dummies = PartialTrieState { - state: self.tries.state.as_hashed_partial_trie(), + state: self.tries.state.as_hashed_partial_trie().clone(), storage: self .tries .storage .iter() - .map(|(k, v)| (*k, v.as_hashed_partial_trie())) + .map(|(k, v)| (*k, v.as_hashed_partial_trie().clone())) .collect(), ..Default::default() }; diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index f69dee624..4c2ef6d69 100644 --- a/trace_decoder/src/lib.rs +++ b/trace_decoder/src/lib.rs @@ -326,13 +326,15 @@ pub fn entrypoint( |mut acc, (nibbles, hash_or_val)| { let path = TriePath::from_nibbles(nibbles); match hash_or_val { - mpt_trie::trie_ops::ValOrHash::Val(bytes) => acc.insert_by_path( - path, - rlp::decode(&bytes) - .context("invalid AccountRlp in direct state trie")?, - ), + mpt_trie::trie_ops::ValOrHash::Val(bytes) => { + acc.insert_by_path( + path, + rlp::decode(&bytes) + .context("invalid AccountRlp in direct state trie")?, + )?; + } mpt_trie::trie_ops::ValOrHash::Hash(h) => { - acc.insert_branch_by_path(path, h) + acc.insert_hash_by_path(path, h)?; } }; anyhow::Ok(acc) @@ -346,10 +348,10 @@ pub fn entrypoint( let path = TriePath::from_nibbles(nibbles); match hash_or_val { mpt_trie::trie_ops::ValOrHash::Val(value) => { - acc.insert(path, value) + acc.insert(path, value)?; } mpt_trie::trie_ops::ValOrHash::Hash(h) => { - acc.insert_branch(path, h) + acc.insert_hash(path, h)?; } }; anyhow::Ok(acc) @@ -392,7 +394,7 @@ pub fn entrypoint( .tries .state .iter() - .filter_map(|(addr, data)| Some((addr.into_hash_left_padded(), data.right()?))) + .map(|(addr, data)| (addr.into_hash_left_padded(), data)) .collect::>(); let code_db = { diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index f6046be7b..b50e2597c 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -52,10 +52,9 @@ fn visit( ) -> anyhow::Result<()> { match node { Node::Hash(Hash { raw_hash }) => { - let clobbered = frontend + frontend .state - .insert_branch_by_path(TriePath::new(path.iter().copied())?, raw_hash.into()); - ensure!(clobbered.is_none(), "duplicate hash") + .insert_hash_by_path(TriePath::new(path.iter().copied())?, raw_hash.into())?; } Node::Leaf(Leaf { key, value }) => { let path = TriePath::new(path.iter().copied().chain(key))?; @@ -94,7 +93,7 @@ fn visit( } }, }; - let clobbered = frontend.state.insert_by_path(path, account); + let clobbered = frontend.state.insert_by_path(path, account)?; ensure!(clobbered.is_none(), "duplicate account"); } } @@ -129,14 +128,14 @@ fn node2storagetrie(node: Node) -> anyhow::Result { ) -> anyhow::Result<()> { match node { Node::Hash(Hash { raw_hash }) => { - mpt.insert_branch(TriePath::new(path.iter().copied())?, raw_hash.into()); + mpt.insert_hash(TriePath::new(path.iter().copied())?, raw_hash.into())?; } Node::Leaf(Leaf { key, value }) => { match value { Either::Left(Value { raw_value }) => mpt.insert( TriePath::new(path.iter().copied().chain(key))?, raw_value.into_vec(), - ), + )?, Either::Right(_) => bail!("unexpected account node in storage trie"), }; } @@ -379,12 +378,9 @@ fn test() { let frontend = frontend(instructions).unwrap(); assert_eq!(case.expected_state_root, frontend.state.root()); - for (path, hash_or_acct) in frontend.state { - if let Either::Right(acct) = hash_or_acct { - if acct.storage_root != mpt_trie::partial_trie::HashedPartialTrie::default().hash() - { - assert!(frontend.storage.contains_key(&path)) - } + for (path, acct) in &frontend.state { + if acct.storage_root != mpt_trie::partial_trie::HashedPartialTrie::default().hash() { + assert!(frontend.storage.contains_key(&path)) } } } diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index f639401c7..2f68a0234 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -1,14 +1,125 @@ //! Principled MPT types used in this library. -use std::collections::BTreeMap; +use std::marker::PhantomData; use copyvec::CopyVec; -use either::Either; use ethereum_types::{Address, H256}; use evm_arithmetization::generation::mpt::AccountRlp; -use mpt_trie::partial_trie::PartialTrie as _; +use mpt_trie::{ + partial_trie::{HashedPartialTrie, Node, PartialTrie as _}, + trie_ops::TrieOpError, +}; use u4::{AsNibbles, U4}; +/// Map where keys are [up to 64 nibbles](TriePath), +/// and values are [`rlp::Encodable`]/[`rlp::Decodable`]. +/// +/// See . +/// +/// Portions of the trie may be deferred: see [`Self::insert_hash`]. +#[derive(Debug, Clone, PartialEq, Eq)] +struct TypedMpt { + inner: HashedPartialTrie, + _ty: PhantomData T>, +} + +impl TypedMpt { + fn new() -> Self { + Self { + inner: HashedPartialTrie::new(Node::Empty), + _ty: PhantomData, + } + } + /// Insert a node which represents an out-of-band sub-trie. + fn insert_hash(&mut self, path: TriePath, hash: H256) -> Result<(), Error> { + self.inner + .insert(path.into_nibbles(), hash) + .map_err(|source| Error { source }) + } + /// Returns an [`Error`] if the `path` crosses into a part of the trie that + /// isn't hydrated. + fn insert(&mut self, path: TriePath, value: T) -> Result, Error> + where + T: rlp::Encodable + rlp::Decodable, + { + let prev = self.get(path); + self.inner + .insert(path.into_nibbles(), rlp::encode(&value).to_vec()) + .map_err(|source| Error { source }) + .map(|_| prev) + } + /// Note that this returns [`None`] if `path` crosses into a part of the + /// trie that isn't hydrated. + /// + /// # Panics + /// - If [`rlp::decode`]-ing for `T` doesn't round-trip. + fn get(&self, path: TriePath) -> Option + where + T: rlp::Decodable, + { + let bytes = self.inner.get(path.into_nibbles())?; + Some(rlp::decode(bytes).expect( + "T encoding/decoding should round-trip,\ + and only encoded `T`s are ever inserted", + )) + } + /// # Panics + /// - If [`rlp::decode`]-ing for `T` doesn't round-trip. + fn remove(&mut self, path: TriePath) -> Result, Error> + where + T: rlp::Decodable, + { + match self.inner.delete(path.into_nibbles()) { + Ok(None) => Ok(None), + Ok(Some(bytes)) => Ok(Some(rlp::decode(&bytes).expect( + "T encoding/decoding should round-trip,\ + and only encoded `T`s are ever inserted", + ))), + // TODO(0xaatif): why is this fallible if `get` isn't? + Err(source) => Err(Error { source }), + } + } + fn as_hashed_partial_trie(&self) -> &HashedPartialTrie { + &self.inner + } + fn root(&self) -> H256 { + self.inner.hash() + } + /// Note that this returns owned paths and items. + fn iter(&self) -> impl Iterator + '_ + where + T: rlp::Decodable, + { + self.inner.keys().filter_map(|nib| { + let path = TriePath::from_nibbles(nib); + Some((path, self.get(path)?)) + }) + } +} + +impl Default for TypedMpt { + fn default() -> Self { + Self::new() + } +} + +impl<'a, T> IntoIterator for &'a TypedMpt +where + T: rlp::Decodable, +{ + type Item = (TriePath, T); + type IntoIter = Box + 'a>; + fn into_iter(self) -> Self::IntoIter { + Box::new(self.iter()) + } +} + +#[derive(thiserror::Error, Debug)] +#[error(transparent)] +pub struct Error { + source: TrieOpError, +} + /// Bounded sequence of [`U4`], /// used as a key for [`TypedMpt`]. /// @@ -39,8 +150,7 @@ impl TriePath { "\ rlp of an usize goes through a u64, which is 8 bytes, which will be 9 bytes RLP'ed. - 9 < 32 - ", + 9 < 32", ) } fn into_nibbles(self) -> mpt_trie::nibbles::Nibbles { @@ -63,96 +173,7 @@ impl TriePath { } } -/// Map where keys are [up to 64 nibbles](TriePath), and values are either an -/// out-of-band [hash](H256) or an inline [`rlp::Encodable`]/[`rlp::Decodable`] -/// value. -/// -/// [Merkle Patricia Trees](https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie) -/// are _maps_, where keys are typically _sequences_ of an _alphabet_. -/// -/// Map values are typically indirect (i.e a _hash_), -/// but in this structure may be stored _inline_. -/// -/// Semantically equivalent to a [`mpt_trie::partial_trie::HashedPartialTrie`]. -#[derive(Debug, Clone)] -struct TypedMpt { - map: BTreeMap>, -} - -impl Default for TypedMpt { - fn default() -> Self { - Self::new() - } -} - -impl TypedMpt { - pub const fn new() -> Self { - Self { - map: BTreeMap::new(), - } - } - // pub fn remove(&mut self, path: TriePath) -> Option> { - // self.map.remove(&path) - // } - pub fn insert(&mut self, path: TriePath, value: T) -> Option> { - self.map.insert(path, Either::Right(value)) - } - pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option> { - self.map.insert(path, Either::Left(hash)) - } - // pub fn get(&self, path: TriePath) -> Option> { - // self.map.get(&path).map(|it| it.as_ref().map_left(|it| *it)) - // } - pub fn root(&self) -> H256 - where - T: rlp::Encodable, - { - self.as_hashed_partial_trie().hash() - } - pub fn iter(&self) -> impl Iterator)> { - self.map.iter() - } - pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie - where - T: rlp::Encodable, - { - let mut theirs = mpt_trie::partial_trie::HashedPartialTrie::default(); - for (path, v) in &self.map { - let nibbles = path.into_nibbles(); - match v { - Either::Left(h) => theirs.insert(nibbles, *h), - Either::Right(v) => theirs.insert(nibbles, &*rlp::encode(v)), - } - .expect("internal error in legacy MPT library") - } - theirs - } -} - -impl<'a, T> IntoIterator for &'a TypedMpt { - type Item = (&'a TriePath, &'a Either); - - type IntoIter = Box)> + 'a>; - - fn into_iter(self) -> Self::IntoIter { - Box::new(self.iter()) - } -} - -impl IntoIterator for TypedMpt -where - T: 'static, -{ - type Item = (TriePath, Either); - - type IntoIter = Box)>>; - - fn into_iter(self) -> Self::IntoIter { - Box::new(self.map.into_iter()) - } -} - -/// Per-block, keyed by transaction index. +/// Per-block, `txn_ix -> [u8]`. /// /// See #[derive(Debug, Clone, Default)] @@ -161,18 +182,18 @@ pub struct TransactionTrie { } impl TransactionTrie { - pub fn insert(&mut self, txn_ix: usize, val: Vec) -> Option>> { + pub fn insert(&mut self, txn_ix: usize, val: Vec) -> Result>, Error> { self.typed.insert(TriePath::from_txn_ix(txn_ix), val) } pub fn root(&self) -> H256 { self.typed.root() } - pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { self.typed.as_hashed_partial_trie() } } -/// Per-block, keyed by transaction index. +/// Per-block, `txn_ix -> [u8]`. /// /// See #[derive(Debug, Clone, Default)] @@ -181,15 +202,15 @@ pub struct ReceiptTrie { } impl ReceiptTrie { - pub fn insert(&mut self, txn_ix: usize, val: Vec) { - self.typed.insert(TriePath::from_txn_ix(txn_ix), val); + pub fn insert(&mut self, txn_ix: usize, val: Vec) -> Result>, Error> { + self.typed.insert(TriePath::from_txn_ix(txn_ix), val) } - pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { self.typed.as_hashed_partial_trie() } } -/// Global, keyed by address +/// Global, [`Address`] `->` [`AccountRlp`]. /// /// See #[derive(Debug, Clone, Default)] @@ -202,44 +223,40 @@ impl StateTrie { &mut self, address: Address, account: AccountRlp, - ) -> Option> { + ) -> Result, Error> { self.insert_by_path(TriePath::from_address(address), account) } pub fn insert_by_path( &mut self, path: TriePath, account: AccountRlp, - ) -> Option> { + ) -> Result, Error> { self.typed.insert(path, account) } - pub fn insert_branch_by_path( - &mut self, - path: TriePath, - hash: H256, - ) -> Option> { - self.typed.insert_branch(path, hash) + pub fn insert_hash_by_path(&mut self, path: TriePath, hash: H256) -> Result<(), Error> { + self.typed.insert_hash(path, hash) } - pub fn get_by_path(&self, path: TriePath) -> Option> { - self.typed.map.get(&path).copied() + pub fn get_by_path(&self, path: TriePath) -> Option { + self.typed.get(path) } - pub fn get_by_address(&self, address: Address) -> Option> { + pub fn get_by_address(&self, address: Address) -> Option { self.get_by_path(TriePath::from_hash(keccak_hash::keccak(address))) } pub fn root(&self) -> H256 { self.typed.root() } - pub fn iter(&self) -> impl Iterator)> { + pub fn iter(&self) -> impl Iterator + '_ { self.typed.iter() } - pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { + pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { self.typed.as_hashed_partial_trie() } } -impl IntoIterator for StateTrie { - type Item = (TriePath, Either); +impl<'a> IntoIterator for &'a StateTrie { + type Item = (TriePath, AccountRlp); - type IntoIter = Box)>>; + type IntoIter = Box + 'a>; fn into_iter(self) -> Self::IntoIter { self.typed.into_iter() @@ -252,33 +269,22 @@ impl IntoIterator for StateTrie { #[derive(Debug, Clone, Default)] pub struct StorageTrie { /// This does NOT use [`TypedMpt`] - T could be anything! - map: BTreeMap>>, + typed: TypedMpt>, } impl StorageTrie { - pub fn insert(&mut self, path: TriePath, value: Vec) -> Option>> { - self.map.insert(path, Either::Right(value)) + pub fn insert(&mut self, path: TriePath, value: Vec) -> Result>, Error> { + self.typed.insert(path, value) } - pub fn insert_branch(&mut self, path: TriePath, hash: H256) -> Option>> { - self.map.insert(path, Either::Left(hash)) + pub fn insert_hash(&mut self, path: TriePath, hash: H256) -> Result<(), Error> { + self.typed.insert_hash(path, hash) } pub fn root(&self) -> H256 { - self.as_hashed_partial_trie().hash() - } - pub fn remove(&mut self, path: TriePath) -> Option>> { - self.map.remove(&path) - } - pub fn as_hashed_partial_trie(&self) -> mpt_trie::partial_trie::HashedPartialTrie { - let mut theirs = mpt_trie::partial_trie::HashedPartialTrie::default(); - for (k, v) in &self.map { - let nibbles = k.into_nibbles(); - match v { - Either::Left(h) => theirs.insert(nibbles, *h), - // TODO(0xaatif): why is RLP-encoding here the right thing to do? - // (Our tests fail without it). - Either::Right(v) => theirs.insert(nibbles, &*rlp::encode(v)), - } - .expect("error in mpt_trie") - } - theirs + self.typed.root() + } + pub fn remove(&mut self, path: TriePath) -> Result>, Error> { + self.typed.remove(path) + } + pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { + self.typed.as_hashed_partial_trie() } } From 992f223978373b3af6eb2c9f4215cd02a2ecad00 Mon Sep 17 00:00:00 2001 From: 0xaatif Date: Thu, 18 Jul 2024 18:24:14 +0100 Subject: [PATCH 6/6] review: tweaks --- trace_decoder/src/type1.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index 96e153fca..9f0cc171b 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -18,7 +18,7 @@ pub struct Frontend { pub state: StateTrie, pub code: BTreeSet>>, /// The key here matches the [`TriePath`] inside [`Self::state`] for - /// accounts which had inline storage + /// accounts which had inline storage. pub storage: BTreeMap, } @@ -59,9 +59,7 @@ fn visit( Node::Leaf(Leaf { key, value }) => { let path = TriePath::new(path.iter().copied().chain(key))?; match value { - // TODO(0xaatif): what should this be interpreted as? - // (this branch isn't hit in our tests) - Either::Left(Value { .. }) => bail!("unsupported value node"), + Either::Left(Value { .. }) => bail!("unsupported value node at top level"), Either::Right(Account { nonce, balance,