From fb1aa3d3d8856762019874788e780626b8641ebd Mon Sep 17 00:00:00 2001 From: 0xaatif <169152398+0xaatif@users.noreply.github.com> Date: Thu, 15 Aug 2024 14:47:51 +0100 Subject: [PATCH] refactor: use typed_mpt in the backend (#494) * mark: 0xaatif/refactor-trace-decoder-decoding * refactor: remove inappropriate static methods * refactor: remove inappropriate method * refactor: remove TraceParsingErrorReason * refactor: remove LocatedError * refactor: remove EMPTY_ACCOUNT_BYTES_RLPED * refactor: remove update_val_if_some * refactor: impl Display for TrieType * wibble * refactor: inline TraceParsingResult * wibble * wibble: order * refactor: remove unused variable * wibble: inline TrieRoots etc * wibble * review: WithHash -> CustomFmt * mark: 0xaatif/typed-backend3 * refactor: trace_decoder::typed_mpt::TriePath -> TrieKey * refactor: path: TriePath -> key: TrieKey * refactor: use typed_mpt in the backend * tweaks * refactor: remove allow(unused) * refactor: remove dead methods * bump * fix: do not force precompile address access in case of txn reversion (#488) * Fix precompile insertion in state trie * Reword * fix indexing in error message (#490) * refactor: more path -> key --------- Co-authored-by: Robin Salen <30937548+Nashtare@users.noreply.github.com> --- trace_decoder/src/decoding.rs | 249 ++++++++++----------- trace_decoder/src/lib.rs | 16 +- trace_decoder/src/processed_block_trace.rs | 38 ++-- trace_decoder/src/type1.rs | 17 +- trace_decoder/src/type2.rs | 3 +- trace_decoder/src/typed_mpt.rs | 147 ++++++------ 6 files changed, 240 insertions(+), 230 deletions(-) diff --git a/trace_decoder/src/decoding.rs b/trace_decoder/src/decoding.rs index 39f0c745f..60b6be301 100644 --- a/trace_decoder/src/decoding.rs +++ b/trace_decoder/src/decoding.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, fmt}; +use std::collections::HashMap; use anyhow::{anyhow, Context as _}; use ethereum_types::{Address, BigEndianHash, H256, U256, U512}; @@ -12,7 +12,7 @@ use evm_arithmetization::{ }; use mpt_trie::{ nibbles::Nibbles, - partial_trie::{HashedPartialTrie, Node, PartialTrie}, + partial_trie::{HashedPartialTrie, PartialTrie as _}, special_query::path_for_query, utils::{IntoTrieKey as _, TriePath}, }; @@ -22,6 +22,7 @@ use crate::{ processed_block_trace::{ NodesUsedByTxn, ProcessedBlockTrace, ProcessedTxnInfo, StateTrieWrites, TxnMetaState, }, + typed_mpt::{ReceiptTrie, StateTrie, StorageTrie, TransactionTrie, TrieKey}, OtherBlockData, PartialTriePreImages, }; @@ -29,10 +30,10 @@ use crate::{ /// after every txn we process in the trace. #[derive(Clone, Debug, Default)] struct PartialTrieState { - state: HashedPartialTrie, - storage: HashMap, - txn: HashedPartialTrie, - receipt: HashedPartialTrie, + state: StateTrie, + storage: HashMap, + txn: TransactionTrie, + receipt: ReceiptTrie, } /// Additional information discovered during delta application. @@ -40,8 +41,8 @@ struct PartialTrieState { struct TrieDeltaApplicationOutput { // During delta application, if a delete occurs, we may have to make sure additional nodes // that are not accessed by the txn remain unhashed. - additional_state_trie_paths_to_not_hash: Vec, - additional_storage_trie_paths_to_not_hash: HashMap>, + additional_state_trie_paths_to_not_hash: Vec, + additional_storage_trie_paths_to_not_hash: HashMap>, } pub fn into_txn_proof_gen_ir( @@ -53,11 +54,8 @@ pub fn into_txn_proof_gen_ir( other_data: OtherBlockData, ) -> anyhow::Result> { let mut curr_block_tries = PartialTrieState { - state: state.as_hashed_partial_trie().clone(), - storage: storage - .iter() - .map(|(k, v)| (*k, v.as_hashed_partial_trie().clone())) - .collect(), + state: state.clone(), + storage: storage.iter().map(|(k, v)| (*k, v.clone())).collect(), ..Default::default() }; @@ -134,23 +132,21 @@ fn update_beacon_block_root_contract_storage( let mut slots_nibbles = vec![]; - for (slot, val) in [(timestamp_idx, timestamp), (root_idx, calldata)] - .iter() - .map(|(k, v)| { - ( - Nibbles::from_h256_be(hash(Nibbles::from_h256_be(H256::from_uint(k)).bytes_be())), - v, - ) - }) - { + for (ix, val) in [(timestamp_idx, timestamp), (root_idx, calldata)] { + // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 + // document this + let slot = TrieKey::from_nibbles(Nibbles::from_h256_be(hash( + Nibbles::from_h256_be(H256::from_uint(&ix)).bytes_be(), + ))); + slots_nibbles.push(slot); // If we are writing a zero, then we actually need to perform a delete. - match val == &ZERO_STORAGE_SLOT_VAL_RLPED { + match val == ZERO_STORAGE_SLOT_VAL_RLPED { false => { storage_trie.insert(slot, val.clone()).context(format!( - "at slot {} with value {}", - CustomFmt(U512::from_big_endian(slot.bytes_be().as_slice())), + "at slot {:?} with value {}", + slot, U512::from_big_endian(val.as_slice()) ))?; @@ -162,7 +158,10 @@ fn update_beacon_block_root_contract_storage( } true => { if let Ok(Some(remaining_slot_key)) = - delete_node_and_report_remaining_key_if_branch_collapsed(storage_trie, &slot) + delete_node_and_report_remaining_key_if_branch_collapsed( + storage_trie.as_mut_hashed_partial_trie_unchecked(), + &slot, + ) { delta_out .additional_storage_trie_paths_to_not_hash @@ -176,26 +175,23 @@ fn update_beacon_block_root_contract_storage( nodes_used.storage_accesses.push((ADDRESS, slots_nibbles)); - let addr_nibbles = Nibbles::from_h256_be(ADDRESS); + let addr_nibbles = TrieKey::from_hash(ADDRESS); delta_out .additional_state_trie_paths_to_not_hash .push(addr_nibbles); - let addr_bytes = trie_state + let mut account = trie_state .state - .get(addr_nibbles) + .get_by_key(addr_nibbles) .context(format!("missing account storage trie {:x}", ADDRESS))?; - let mut account = account_from_rlped_bytes(addr_bytes)?; - account.storage_root = storage_trie.hash(); + account.storage_root = storage_trie.root(); - let updated_account_bytes = rlp::encode(&account); trie_state .state - .insert(addr_nibbles, updated_account_bytes.to_vec()) - .context(format!( - "at slot {}", - CustomFmt(U512::from_big_endian(addr_nibbles.bytes_be().as_slice())) - ))?; + .insert_by_key(addr_nibbles, account) + // 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(()) } @@ -210,19 +206,18 @@ fn update_txn_and_receipt_tries( return Ok(()); } - let txn_k = Nibbles::from_bytes_be(&rlp::encode(&txn_idx)).unwrap(); - trie_state.txn.insert(txn_k, meta.txn_bytes())?; - - Ok(trie_state + trie_state.txn.insert(txn_idx, meta.txn_bytes())?; + trie_state .receipt - .insert(txn_k, meta.receipt_node_bytes.as_ref())?) + .insert(txn_idx, meta.receipt_node_bytes.clone())?; + Ok(()) } /// If the account does not have a storage trie or does but is not /// accessed by any txns, then we still need to manually create an entry for /// them. fn init_any_needed_empty_storage_tries<'a>( - storage_tries: &mut HashMap, + storage_tries: &mut HashMap, accounts_with_storage: impl Iterator, state_accounts_with_no_accesses_but_storage_tries: &'a HashMap, ) { @@ -230,7 +225,12 @@ fn init_any_needed_empty_storage_tries<'a>( if !storage_tries.contains_key(h_addr) { let trie = state_accounts_with_no_accesses_but_storage_tries .get(h_addr) - .map(|s_root| HashedPartialTrie::new(Node::Hash(*s_root))) + .map(|s_root| { + let mut it = StorageTrie::default(); + it.insert_hash(TrieKey::default(), *s_root) + .expect("empty trie insert cannot fail"); + it + }) .unwrap_or_default(); storage_tries.insert(*h_addr, trie); @@ -250,19 +250,27 @@ fn create_minimal_partial_tries_needed_by_txn( delta_application_out .additional_state_trie_paths_to_not_hash .into_iter(), - )?; + )? + .as_hashed_partial_trie() + .clone(); - let txn_k = Nibbles::from_bytes_be(&rlp::encode(&txn_idx)).unwrap(); + let txn_k = TrieKey::from_txn_ix(txn_idx); - let transactions_trie = - create_trie_subset_wrapped(&curr_block_tries.txn, [txn_k], TrieType::Txn)?; + let transactions_trie = create_trie_subset_wrapped( + curr_block_tries.txn.as_hashed_partial_trie(), + [txn_k], + TrieType::Txn, + )?; - let receipts_trie = - create_trie_subset_wrapped(&curr_block_tries.receipt, [txn_k], TrieType::Receipt)?; + let receipts_trie = create_trie_subset_wrapped( + curr_block_tries.receipt.as_hashed_partial_trie(), + [txn_k], + TrieType::Receipt, + )?; let storage_tries = create_minimal_storage_partial_tries( &curr_block_tries.storage, - nodes_used_by_txn.storage_accesses.iter(), + &nodes_used_by_txn.storage_accesses, &delta_application_out.additional_storage_trie_paths_to_not_hash, )?; @@ -281,7 +289,7 @@ fn apply_deltas_to_trie_state( ) -> anyhow::Result { let mut out = TrieDeltaApplicationOutput::default(); - for (hashed_acc_addr, storage_writes) in deltas.storage_writes.iter() { + for (hashed_acc_addr, storage_writes) in &deltas.storage_writes { let storage_trie = trie_state .storage .get_mut(hashed_acc_addr) @@ -290,21 +298,21 @@ fn apply_deltas_to_trie_state( hashed_acc_addr ))?; - for (slot, val) in storage_writes - .iter() - .map(|(k, v)| (Nibbles::from_h256_be(hash(k.bytes_be())), v)) - { + for (key, val) in storage_writes { + let slot = TrieKey::from_hash(hash(key.into_nibbles().bytes_be())); // If we are writing a zero, then we actually need to perform a delete. match val == &ZERO_STORAGE_SLOT_VAL_RLPED { - false => storage_trie.insert(slot, val.clone()).context(format!( - "at slot {} with value {}", - CustomFmt(U512::from_big_endian(slot.bytes_be().as_slice())), - U512::from_big_endian(val.as_slice()) - ))?, + false => { + storage_trie.insert(slot, val.clone()).context(format!( + "at slot {:?} with value {}", + slot, + U512::from_big_endian(val.as_slice()) + ))?; + } true => { if let Some(remaining_slot_key) = delete_node_and_report_remaining_key_if_branch_collapsed( - storage_trie, + storage_trie.as_mut_hashed_partial_trie_unchecked(), &slot, )? { @@ -318,17 +326,12 @@ fn apply_deltas_to_trie_state( } } - for (hashed_acc_addr, s_trie_writes) in deltas.state_writes.iter() { - let val_k = Nibbles::from_h256_be(*hashed_acc_addr); + for (hashed_acc_addr, s_trie_writes) in &deltas.state_writes { + let val_k = TrieKey::from_hash(*hashed_acc_addr); // If the account was created, then it will not exist in the trie yet. let is_created = !trie_state.state.contains(val_k); - let mut account = trie_state - .state - .get(val_k) - .map(account_from_rlped_bytes) - .transpose()? - .unwrap_or_default(); + let mut account = trie_state.state.get_by_key(val_k).unwrap_or_default(); s_trie_writes.apply_writes_to_state_node( &mut account, @@ -336,7 +339,6 @@ fn apply_deltas_to_trie_state( &trie_state.storage, )?; - let updated_account_bytes = rlp::encode(&account); if is_created { // If the account did not exist prior this transaction, we // need to make sure the transaction didn't revert. @@ -347,14 +349,12 @@ fn apply_deltas_to_trie_state( if !receipt.status { // The transaction failed, hence any created account should be removed. - trie_state.state.delete(val_k)?; + trie_state.state.remove(val_k)?; trie_state.storage.remove(hashed_acc_addr); continue; } } - trie_state - .state - .insert(val_k, updated_account_bytes.to_vec())?; + trie_state.state.insert_by_key(val_k, account)?; } Ok(out) @@ -369,15 +369,15 @@ fn get_trie_trace(trie: &HashedPartialTrie, k: &Nibbles) -> TriePath { /// 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, - delete_k: &Nibbles, -) -> anyhow::Result> { - let old_trace = get_trie_trace(trie, delete_k); - trie.delete(*delete_k)?; - let new_trace = get_trie_trace(trie, delete_k); - - Ok(node_deletion_resulted_in_a_branch_collapse( - &old_trace, &new_trace, - )) + delete_k: &TrieKey, +) -> anyhow::Result> { + let old_trace = get_trie_trace(trie, &delete_k.into_nibbles()); + trie.delete(delete_k.into_nibbles())?; + let new_trace = get_trie_trace(trie, &delete_k.into_nibbles()); + Ok( + node_deletion_resulted_in_a_branch_collapse(&old_trace, &new_trace) + .map(TrieKey::from_nibbles), + ) } /// Comparing the path of the deleted key before and after the deletion, @@ -434,7 +434,7 @@ fn add_withdrawals_to_txns( let additional_paths = if last_inputs.txn_number_before == 0.into() { // We need to include the beacon roots contract as this payload is at the // start of the block execution. - vec![Nibbles::from_h256_be(H256( + vec![TrieKey::from_hash(H256( BEACON_ROOTS_CONTRACT_ADDRESS_HASHED, ))] } else { @@ -444,8 +444,10 @@ fn add_withdrawals_to_txns( last_inputs.tries.state_trie = create_minimal_state_partial_trie( &final_trie_state.state, withdrawal_addrs, - additional_paths.into_iter(), - )?; + additional_paths, + )? + .as_hashed_partial_trie() + .clone(); } update_trie_state_from_withdrawals( @@ -454,7 +456,7 @@ fn add_withdrawals_to_txns( )?; last_inputs.withdrawals = withdrawals; - last_inputs.trie_roots_after.state_root = final_trie_state.state.hash(); + last_inputs.trie_roots_after.state_root = final_trie_state.state.root(); Ok(()) } @@ -463,19 +465,20 @@ fn add_withdrawals_to_txns( /// our local trie state. fn update_trie_state_from_withdrawals<'a>( withdrawals: impl IntoIterator + 'a, - state: &mut HashedPartialTrie, + state: &mut StateTrie, ) -> anyhow::Result<()> { for (addr, h_addr, amt) in withdrawals { - let h_addr_nibs = Nibbles::from_h256_be(h_addr); - - let acc_bytes = state.get(h_addr_nibs).context(format!( + let mut acc_data = state.get_by_address(addr).context(format!( "No account present at {addr:x} (hashed: {h_addr:x}) to withdraw {amt} Gwei from!" ))?; - let mut acc_data = account_from_rlped_bytes(acc_bytes)?; acc_data.balance += amt; - state.insert(h_addr_nibs, rlp::encode(&acc_data).to_vec())?; + state + .insert_by_address(addr, acc_data) + // 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(()) @@ -552,9 +555,9 @@ fn process_txn_info( * for more info). */ tries, trie_roots_after: TrieRoots { - state_root: curr_block_tries.state.hash(), - transactions_root: curr_block_tries.txn.hash(), - receipts_root: curr_block_tries.receipt.hash(), + state_root: curr_block_tries.state.root(), + transactions_root: curr_block_tries.txn.root(), + receipts_root: curr_block_tries.receipt.root(), }, checkpoint_state_trie_root: extra_data.checkpoint_state_trie_root, contract_code: txn_info.contract_code_accessed, @@ -576,7 +579,7 @@ impl StateTrieWrites { &self, state_node: &mut AccountRlp, h_addr: &H256, - acc_storage_tries: &HashMap, + acc_storage_tries: &HashMap, ) -> anyhow::Result<()> { let storage_root_hash_change = match self.storage_trie_change { false => None, @@ -585,7 +588,7 @@ impl StateTrieWrites { .get(h_addr) .context(format!("missing account storage trie {:x}", h_addr))?; - Some(storage_trie.hash()) + Some(storage_trie.root()) } }; @@ -599,28 +602,30 @@ impl StateTrieWrites { } fn create_minimal_state_partial_trie( - state_trie: &HashedPartialTrie, - state_accesses: impl Iterator, - additional_state_trie_paths_to_not_hash: impl Iterator, -) -> anyhow::Result { + state_trie: &StateTrie, + state_accesses: impl IntoIterator, + additional_state_trie_paths_to_not_hash: impl IntoIterator, +) -> anyhow::Result { create_trie_subset_wrapped( - state_trie, + state_trie.as_hashed_partial_trie(), state_accesses .into_iter() - .map(Nibbles::from_h256_be) + .map(TrieKey::from_hash) .chain(additional_state_trie_paths_to_not_hash), TrieType::State, ) + .map(StateTrie::from_hashed_partial_trie_unchecked) } // TODO!!!: We really need to be appending the empty storage tries to the base // trie somewhere else! This is a big hack! fn create_minimal_storage_partial_tries<'a>( - storage_tries: &HashMap, - accesses_per_account: impl Iterator)>, - additional_storage_trie_paths_to_not_hash: &HashMap>, + storage_tries: &HashMap, + accesses_per_account: impl IntoIterator)>, + additional_storage_trie_paths_to_not_hash: &HashMap>, ) -> anyhow::Result> { accesses_per_account + .into_iter() .map(|(h_addr, mem_accesses)| { // Guaranteed to exist due to calling `init_any_needed_empty_storage_tries` // earlier on. @@ -634,7 +639,7 @@ fn create_minimal_storage_partial_tries<'a>( ); let partial_storage_trie = create_trie_subset_wrapped( - base_storage_trie, + base_storage_trie.as_hashed_partial_trie(), storage_slots_to_not_hash, TrieType::Storage, )?; @@ -646,15 +651,14 @@ fn create_minimal_storage_partial_tries<'a>( fn create_trie_subset_wrapped( trie: &HashedPartialTrie, - accesses: impl IntoIterator, + accesses: impl IntoIterator, trie_type: TrieType, ) -> anyhow::Result { - mpt_trie::trie_subsets::create_trie_subset(trie, accesses) - .context(format!("missing keys when creating {}", trie_type)) -} - -fn account_from_rlped_bytes(bytes: &[u8]) -> anyhow::Result { - Ok(rlp::decode(bytes)?) + mpt_trie::trie_subsets::create_trie_subset( + trie, + accesses.into_iter().map(TrieKey::into_nibbles), + ) + .context(format!("missing keys when creating {}", trie_type)) } impl TxnMetaState { @@ -680,22 +684,11 @@ fn eth_to_gwei(eth: U256) -> U256 { // This is just `rlp(0)`. const ZERO_STORAGE_SLOT_VAL_RLPED: [u8; 1] = [128]; -/// Formatting aid for error context. -struct CustomFmt(T); - -impl fmt::Display for CustomFmt { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let mut buf = [0u8; 64]; - self.0.to_big_endian(&mut buf); - f.write_fmt(format_args!("{} (hashed: 0x{:064X})", self.0, hash(buf))) - } -} - /// Aid for error context. /// Covers all Ethereum trie types (see for details). #[derive(Debug, strum::Display)] #[allow(missing_docs)] -pub enum TrieType { +enum TrieType { State, Storage, Receipt, diff --git a/trace_decoder/src/lib.rs b/trace_decoder/src/lib.rs index 2bd95cb1c..1e4f44496 100644 --- a/trace_decoder/src/lib.rs +++ b/trace_decoder/src/lib.rs @@ -17,7 +17,8 @@ //! - Performance - this won't be the bottleneck in any proving system. //! - Robustness - malicious or malformed input may crash this library. //! -//! TODO(0xaatif): refactor all the docs below +//! TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 +//! refactor all the docs below //! //! It might not be obvious why we need traces for each txn in order to generate //! proofs. While it's true that we could just run all the txns of a block in an @@ -83,11 +84,12 @@ const _DEVELOPER_DOCS: () = (); /// Defines the main functions used to generate the IR. mod decoding; -// TODO(0xaatif): add backend/prod support /// Defines functions that processes a [BlockTrace] so that it is easier to turn /// the block transactions into IRs. mod processed_block_trace; 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; @@ -104,7 +106,7 @@ use keccak_hash::H256; use mpt_trie::partial_trie::HashedPartialTrie; use processed_block_trace::ProcessedTxnInfo; use serde::{Deserialize, Serialize}; -use typed_mpt::{StateTrie, StorageTrie, TriePath}; +use typed_mpt::{StateTrie, StorageTrie, TrieKey}; /// Core payload needed to generate proof for a block. /// Additional data retrievable from the blockchain node (using standard ETH RPC @@ -320,17 +322,17 @@ pub fn entrypoint( state: state.items().try_fold( StateTrie::default(), |mut acc, (nibbles, hash_or_val)| { - let path = TriePath::from_nibbles(nibbles); + let path = TrieKey::from_nibbles(nibbles); match hash_or_val { mpt_trie::trie_ops::ValOrHash::Val(bytes) => { - acc.insert_by_path( + acc.insert_by_key( path, rlp::decode(&bytes) .context("invalid AccountRlp in direct state trie")?, )?; } mpt_trie::trie_ops::ValOrHash::Hash(h) => { - acc.insert_hash_by_path(path, h)?; + acc.insert_hash_by_key(path, h)?; } }; anyhow::Ok(acc) @@ -341,7 +343,7 @@ pub fn entrypoint( .map(|(k, SeparateTriePreImage::Direct(v))| { v.items() .try_fold(StorageTrie::default(), |mut acc, (nibbles, hash_or_val)| { - let path = TriePath::from_nibbles(nibbles); + let path = TrieKey::from_nibbles(nibbles); match hash_or_val { mpt_trie::trie_ops::ValOrHash::Val(value) => { acc.insert(path, value)?; diff --git a/trace_decoder/src/processed_block_trace.rs b/trace_decoder/src/processed_block_trace.rs index eb0a07042..3d99d64d0 100644 --- a/trace_decoder/src/processed_block_trace.rs +++ b/trace_decoder/src/processed_block_trace.rs @@ -4,10 +4,9 @@ use std::iter::once; use ethereum_types::{Address, H256, U256}; use evm_arithmetization::generation::mpt::{AccountRlp, LegacyReceiptRlp}; -use mpt_trie::nibbles::Nibbles; -use mpt_trie::partial_trie::PartialTrie; use crate::hash; +use crate::typed_mpt::TrieKey; use crate::PartialTriePreImages; use crate::{ContractCodeUsage, TxnInfo}; @@ -98,7 +97,7 @@ impl TxnInfo { nodes_used_by_txn.storage_accesses.push(( hashed_addr, storage_access_keys - .map(|H256(bytes)| Nibbles::from_h256_be(hash(bytes))) + .map(|H256(bytes)| TrieKey::from_hash(hash(bytes))) .collect(), )); @@ -124,7 +123,7 @@ impl TxnInfo { let storage_writes_vec = storage_writes .into_iter() - .map(|(k, v)| (Nibbles::from_h256_be(k), rlp::encode(&v).to_vec())) + .map(|(k, v)| (TrieKey::from_hash(k), rlp::encode(&v).to_vec())) .collect(); nodes_used_by_txn @@ -141,8 +140,7 @@ impl TxnInfo { if !is_precompile || tries .state - .as_hashed_partial_trie() - .get(Nibbles::from_h256_be(hashed_addr)) + .get_by_key(TrieKey::from_hash(hashed_addr)) .is_some() { nodes_used_by_txn.state_accesses.push(hashed_addr); @@ -225,32 +223,30 @@ fn create_empty_code_access_map() -> HashMap> { HashMap::from_iter(once((EMPTY_CODE_HASH, Vec::new()))) } -pub(crate) type StorageAccess = Vec; -pub(crate) type StorageWrite = Vec<(Nibbles, Vec)>; - /// Note that "*_accesses" includes writes. #[derive(Debug, Default)] pub(crate) struct NodesUsedByTxn { - pub(crate) state_accesses: Vec, - pub(crate) state_writes: Vec<(H256, StateTrieWrites)>, + pub state_accesses: Vec, + pub state_writes: Vec<(H256, StateTrieWrites)>, // Note: All entries in `storage_writes` also appear in `storage_accesses`. - pub(crate) storage_accesses: Vec<(H256, StorageAccess)>, - pub(crate) storage_writes: Vec<(H256, StorageWrite)>, - pub(crate) state_accounts_with_no_accesses_but_storage_tries: HashMap, + pub storage_accesses: Vec<(H256, Vec)>, + #[allow(clippy::type_complexity)] + pub storage_writes: Vec<(H256, Vec<(TrieKey, Vec)>)>, + pub state_accounts_with_no_accesses_but_storage_tries: HashMap, } #[derive(Debug)] pub(crate) struct StateTrieWrites { - pub(crate) balance: Option, - pub(crate) nonce: Option, - pub(crate) storage_trie_change: bool, - pub(crate) code_hash: Option, + pub balance: Option, + pub nonce: Option, + pub storage_trie_change: bool, + pub code_hash: Option, } #[derive(Debug, Default)] pub(crate) struct TxnMetaState { - pub(crate) txn_bytes: Option>, - pub(crate) receipt_node_bytes: Vec, - pub(crate) gas_used: u64, + pub txn_bytes: Option>, + pub receipt_node_bytes: Vec, + pub gas_used: u64, } diff --git a/trace_decoder/src/type1.rs b/trace_decoder/src/type1.rs index e7dbcf696..305dffa9b 100644 --- a/trace_decoder/src/type1.rs +++ b/trace_decoder/src/type1.rs @@ -10,7 +10,7 @@ use evm_arithmetization::generation::mpt::AccountRlp; use nunny::NonEmpty; use u4::U4; -use crate::typed_mpt::{StateTrie, StorageTrie, TriePath}; +use crate::typed_mpt::{StateTrie, StorageTrie, TrieKey}; use crate::wire::{Instruction, SmtLeaf}; #[derive(Debug, Default, Clone)] @@ -19,7 +19,7 @@ pub struct Frontend { pub code: BTreeSet>>, /// The key here matches the [`TriePath`] inside [`Self::state`] for /// accounts which had inline storage. - pub storage: BTreeMap, + pub storage: BTreeMap, } pub fn frontend(instructions: impl IntoIterator) -> anyhow::Result { @@ -54,10 +54,10 @@ fn visit( Node::Hash(Hash { raw_hash }) => { frontend .state - .insert_hash_by_path(TriePath::new(path.iter().copied())?, raw_hash.into())?; + .insert_hash_by_key(TrieKey::new(path.iter().copied())?, raw_hash.into())?; } Node::Leaf(Leaf { key, value }) => { - let path = TriePath::new(path.iter().copied().chain(key))?; + let path = TrieKey::new(path.iter().copied().chain(key))?; match value { Either::Left(Value { .. }) => bail!("unsupported value node at top level"), Either::Right(Account { @@ -91,7 +91,7 @@ fn visit( } }, }; - let clobbered = frontend.state.insert_by_path(path, account)?; + let clobbered = frontend.state.insert_by_key(path, account)?; ensure!(clobbered.is_none(), "duplicate account"); } } @@ -126,12 +126,12 @@ fn node2storagetrie(node: Node) -> anyhow::Result { ) -> anyhow::Result<()> { match node { Node::Hash(Hash { raw_hash }) => { - mpt.insert_hash(TriePath::new(path.iter().copied())?, raw_hash.into())?; + mpt.insert_hash(TrieKey::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))?, + TrieKey::new(path.iter().copied().chain(key))?, rlp::encode(&raw_value.as_slice()).to_vec(), )?, Either::Right(_) => bail!("unexpected account node in storage trie"), @@ -258,7 +258,8 @@ fn execute( has_storage, } => { // BUG: the spec sometimes writes Node::Account with 5 fields.. - // TODO(0xaatif): should these fields even be optional? + // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 + // should these fields even be optional? let nonce = nonce.unwrap_or_default(); let balance = balance.unwrap_or_default(); let account = match (has_code, has_storage) { diff --git a/trace_decoder/src/type2.rs b/trace_decoder/src/type2.rs index 648ce34c3..2d10edf40 100644 --- a/trace_decoder/src/type2.rs +++ b/trace_decoder/src/type2.rs @@ -184,7 +184,8 @@ fn node2trie( } SmtLeafType::Storage(it) => { ensure!(collated.storage_root.is_none(), "double write of field"); - // TODO(0xaatif): do we not do anything with the storage here? + // 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)) } SmtLeafType::CodeLength => smt_trie::keys::key_code_length(address), diff --git a/trace_decoder/src/typed_mpt.rs b/trace_decoder/src/typed_mpt.rs index f42d7fa37..4c3efaebe 100644 --- a/trace_decoder/src/typed_mpt.rs +++ b/trace_decoder/src/typed_mpt.rs @@ -1,5 +1,6 @@ //! Principled MPT types used in this library. +use core::fmt; use std::marker::PhantomData; use copyvec::CopyVec; @@ -11,7 +12,7 @@ use mpt_trie::{ }; use u4::{AsNibbles, U4}; -/// Map where keys are [up to 64 nibbles](TriePath), +/// Map where keys are [up to 64 nibbles](TrieKey), /// and values are [`rlp::Encodable`]/[`rlp::Decodable`]. /// /// See . @@ -24,6 +25,8 @@ struct TypedMpt { } impl TypedMpt { + const PANIC_MSG: &str = "T encoding/decoding should round-trip,\ + and only encoded `T`s are ever inserted"; fn new() -> Self { Self { inner: HashedPartialTrie::new(Node::Empty), @@ -31,37 +34,44 @@ impl TypedMpt { } } /// Insert a node which represents an out-of-band sub-trie. - fn insert_hash(&mut self, path: TriePath, hash: H256) -> Result<(), Error> { + fn insert_hash(&mut self, key: TrieKey, hash: H256) -> Result<(), Error> { self.inner - .insert(path.into_nibbles(), hash) + .insert(key.into_nibbles(), hash) .map_err(|source| Error { source }) } - /// Returns an [`Error`] if the `path` crosses into a part of the trie that + /// Returns an [`Error`] if the `key` crosses into a part of the trie that /// isn't hydrated. - fn insert(&mut self, path: TriePath, value: T) -> Result, Error> + fn insert(&mut self, key: TrieKey, value: T) -> Result, Error> where T: rlp::Encodable + rlp::Decodable, { - let prev = self.get(path); + let prev = self.get(key); self.inner - .insert(path.into_nibbles(), rlp::encode(&value).to_vec()) + .insert(key.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 + /// Note that this returns [`None`] if `key` 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 + fn get(&self, key: TrieKey) -> 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", - )) + let bytes = self.inner.get(key.into_nibbles())?; + Some(rlp::decode(bytes).expect(Self::PANIC_MSG)) + } + fn remove(&mut self, key: TrieKey) -> Result, Error> + where + T: rlp::Decodable, + { + match self.inner.delete(key.into_nibbles()) { + Ok(Some(it)) => Ok(Some(rlp::decode(&it).expect(Self::PANIC_MSG))), + Ok(None) => Ok(None), + Err(source) => Err(Error { source }), + } } fn as_hashed_partial_trie(&self) -> &HashedPartialTrie { &self.inner @@ -70,21 +80,15 @@ 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 = TriePath::from_nibbles(nib); + let path = TrieKey::from_nibbles(nib); Some((path, self.get(path)?)) }) } - /// This allows users to break the [`TypedMpt`] invariant. - /// If data that isn't an [`rlp::encode`]-ed `T` is inserted, - /// subsequent API calls may panic. - pub fn as_mut_hashed_partial_trie_unchecked(&mut self) -> &mut HashedPartialTrie { - &mut self.inner - } } impl Default for TypedMpt { @@ -97,7 +101,7 @@ impl<'a, T> IntoIterator for &'a TypedMpt where T: rlp::Decodable, { - type Item = (TriePath, T); + type Item = (TrieKey, T); type IntoIter = Box + 'a>; fn into_iter(self) -> Self::IntoIter { Box::new(self.iter()) @@ -115,11 +119,20 @@ pub struct Error { /// /// Semantically equivalent to [`mpt_trie::nibbles::Nibbles`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] -pub struct TriePath(CopyVec); +pub struct TrieKey(CopyVec); + +impl fmt::Display for TrieKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + for u in self.0 { + f.write_fmt(format_args!("{:x}", u))? + } + Ok(()) + } +} -impl TriePath { +impl TrieKey { pub fn new(components: impl IntoIterator) -> anyhow::Result { - Ok(TriePath(CopyVec::try_from_iter(components)?)) + Ok(TrieKey(CopyVec::try_from_iter(components)?)) } pub fn into_hash_left_padded(mut self) -> H256 { for _ in 0..self.0.spare_capacity_mut().len() { @@ -135,16 +148,16 @@ impl TriePath { pub fn from_hash(H256(bytes): H256) -> Self { Self::new(AsNibbles(bytes)).expect("32 bytes is 64 nibbles, which fits") } - #[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 - fn from_txn_ix(txn_ix: usize) -> Self { - TriePath::new(AsNibbles(rlp::encode(&txn_ix))).expect( + + pub fn from_txn_ix(txn_ix: usize) -> Self { + TrieKey::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 { + pub 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) @@ -164,7 +177,6 @@ impl TriePath { } } -#[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 /// Per-block, `txn_ix -> [u8]`. /// /// See @@ -173,15 +185,14 @@ pub struct TransactionTrie { untyped: HashedPartialTrie, } -#[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 impl TransactionTrie { pub fn insert(&mut self, txn_ix: usize, val: Vec) -> Result>, Error> { let prev = self .untyped - .get(TriePath::from_txn_ix(txn_ix).into_nibbles()) + .get(TrieKey::from_txn_ix(txn_ix).into_nibbles()) .map(Vec::from); self.untyped - .insert(TriePath::from_txn_ix(txn_ix).into_nibbles(), val) + .insert(TrieKey::from_txn_ix(txn_ix).into_nibbles(), val) .map_err(|source| Error { source })?; Ok(prev) } @@ -193,7 +204,6 @@ impl TransactionTrie { } } -#[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 /// Per-block, `txn_ix -> [u8]`. /// /// See @@ -202,15 +212,14 @@ pub struct ReceiptTrie { untyped: HashedPartialTrie, } -#[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 impl ReceiptTrie { pub fn insert(&mut self, txn_ix: usize, val: Vec) -> Result>, Error> { let prev = self .untyped - .get(TriePath::from_txn_ix(txn_ix).into_nibbles()) + .get(TrieKey::from_txn_ix(txn_ix).into_nibbles()) .map(Vec::from); self.untyped - .insert(TriePath::from_txn_ix(txn_ix).into_nibbles(), val) + .insert(TrieKey::from_txn_ix(txn_ix).into_nibbles(), val) .map_err(|source| Error { source })?; Ok(prev) } @@ -236,44 +245,58 @@ impl StateTrie { address: Address, account: AccountRlp, ) -> Result, Error> { - self.insert_by_path(TriePath::from_address(address), account) + self.insert_by_key(TrieKey::from_address(address), account) } - pub fn insert_by_path( + pub fn insert_by_key( &mut self, - path: TriePath, + key: TrieKey, account: AccountRlp, ) -> Result, Error> { - self.typed.insert(path, account) + self.typed.insert(key, account) } - pub fn insert_hash_by_path(&mut self, path: TriePath, hash: H256) -> Result<(), Error> { - self.typed.insert_hash(path, hash) + pub fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> Result<(), Error> { + self.typed.insert_hash(key, hash) } - pub fn get_by_path(&self, path: TriePath) -> Option { - self.typed.get(path) + pub fn get_by_key(&self, key: TrieKey) -> Option { + self.typed.get(key) } pub fn get_by_address(&self, address: Address) -> Option { - self.get_by_path(TriePath::from_hash(keccak_hash::keccak(address))) + self.get_by_key(TrieKey::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 { self.typed.as_hashed_partial_trie() } - + pub fn remove(&mut self, key: TrieKey) -> Result, Error> { + self.typed.remove(key) + } + pub fn contains(&self, key: TrieKey) -> bool { + self.typed + .as_hashed_partial_trie() + .contains(key.into_nibbles()) + } /// This allows users to break the [`TypedMpt`] invariant. - /// If data that isn't an [`rlp::encode`]-ed `T` is inserted, + /// If data that isn't a [`rlp::encode`]-ed [`AccountRlp`] is inserted, /// subsequent API calls may panic. - pub fn as_mut_hashed_partial_trie_unchecked(&mut self) -> &mut HashedPartialTrie { - self.typed.as_mut_hashed_partial_trie_unchecked() + pub fn from_hashed_partial_trie_unchecked( + src: mpt_trie::partial_trie::HashedPartialTrie, + ) -> Self { + Self { + typed: TypedMpt { + inner: src, + _ty: PhantomData, + }, + } } } impl<'a> IntoIterator for &'a StateTrie { - type Item = (TriePath, AccountRlp); + type Item = (TrieKey, AccountRlp); type IntoIter = Box + 'a>; @@ -290,31 +313,25 @@ pub struct StorageTrie { untyped: HashedPartialTrie, } impl StorageTrie { - pub fn insert(&mut self, path: TriePath, value: Vec) -> Result>, Error> { - let prev = self.untyped.get(path.into_nibbles()).map(Vec::from); + pub fn insert(&mut self, key: TrieKey, value: Vec) -> Result>, Error> { + let prev = self.untyped.get(key.into_nibbles()).map(Vec::from); self.untyped - .insert(path.into_nibbles(), value) + .insert(key.into_nibbles(), value) .map_err(|source| Error { source })?; Ok(prev) } - pub fn insert_hash(&mut self, path: TriePath, hash: H256) -> Result<(), Error> { + pub fn insert_hash(&mut self, key: TrieKey, hash: H256) -> Result<(), Error> { self.untyped - .insert(path.into_nibbles(), hash) + .insert(key.into_nibbles(), hash) .map_err(|source| Error { source }) } pub fn root(&self) -> H256 { self.untyped.hash() } - #[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 - pub fn remove(&mut self, path: TriePath) -> Result>, Error> { - self.untyped - .delete(path.into_nibbles()) - .map_err(|source| Error { source }) - } pub fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie { &self.untyped } - #[allow(unused)] // TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275 + pub fn as_mut_hashed_partial_trie_unchecked(&mut self) -> &mut HashedPartialTrie { &mut self.untyped }