diff --git a/evm/src/cpu/kernel/constants/trie_type.rs b/evm/src/cpu/kernel/constants/trie_type.rs index 08fd8748..30f4802b 100644 --- a/evm/src/cpu/kernel/constants/trie_type.rs +++ b/evm/src/cpu/kernel/constants/trie_type.rs @@ -1,5 +1,6 @@ use eth_trie_utils::partial_trie::PartialTrie; +#[derive(Copy, Clone)] pub(crate) enum PartialTrieType { Empty = 0, Hash = 1, diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs index d36e2a66..9bde0106 100644 --- a/evm/src/generation/mod.rs +++ b/evm/src/generation/mod.rs @@ -18,7 +18,10 @@ use crate::config::StarkConfig; use crate::cpu::bootstrap_kernel::generate_bootstrap_kernel; use crate::cpu::kernel::aggregator::KERNEL; use crate::cpu::kernel::constants::global_metadata::GlobalMetadata; +use crate::cpu::kernel::constants::global_metadata::GlobalMetadata::StateTrieRoot; +use crate::generation::mpt::AccountRlp; use crate::generation::state::GenerationState; +use crate::generation::trie_extractor::read_state_trie_value; use crate::memory::segments::Segment; use crate::proof::{BlockMetadata, PublicValues, TrieRoots}; use crate::witness::memory::MemoryAddress; @@ -28,6 +31,8 @@ pub mod mpt; pub(crate) mod prover_input; pub(crate) mod rlp; pub(crate) mod state; +mod trie_extractor; +use crate::generation::trie_extractor::read_trie; #[derive(Clone, Debug, Deserialize, Serialize, Default)] /// Inputs needed for trace generation. @@ -87,6 +92,15 @@ pub(crate) fn generate_traces, const D: usize>( )) }; + log::debug!( + "Updated state trie:\n{:#?}", + read_trie::( + &state.memory, + read_metadata(StateTrieRoot).as_usize(), + read_state_trie_value + ) + ); + let trie_roots_before = TrieRoots { state_root: H256::from_uint(&read_metadata(StateTrieRootDigestBefore)), transactions_root: H256::from_uint(&read_metadata(TransactionTrieRootDigestBefore)), diff --git a/evm/src/generation/trie_extractor.rs b/evm/src/generation/trie_extractor.rs new file mode 100644 index 00000000..d35d67eb --- /dev/null +++ b/evm/src/generation/trie_extractor.rs @@ -0,0 +1,98 @@ +use std::collections::HashMap; + +use eth_trie_utils::partial_trie::Nibbles; +use ethereum_types::{BigEndianHash, H256, U256}; +use plonky2::field::extension::Extendable; +use plonky2::hash::hash_types::RichField; + +use crate::cpu::kernel::constants::trie_type::PartialTrieType; +use crate::generation::mpt::AccountRlp; +use crate::memory::segments::Segment; +use crate::witness::memory::{MemoryAddress, MemoryState}; + +pub(crate) fn read_state_trie_value(slice: &[U256]) -> AccountRlp { + AccountRlp { + nonce: slice[0], + balance: slice[1], + storage_root: H256::from_uint(&slice[2]), + code_hash: H256::from_uint(&slice[3]), + } +} + +pub(crate) fn read_trie( + memory: &MemoryState, + ptr: usize, + read_value: fn(&[U256]) -> V, +) -> HashMap +where + F: RichField + Extendable, +{ + let mut res = HashMap::new(); + let empty_nibbles = Nibbles { + count: 0, + packed: U256::zero(), + }; + read_trie_helper::(memory, ptr, read_value, empty_nibbles, &mut res); + res +} + +pub(crate) fn read_trie_helper( + memory: &MemoryState, + ptr: usize, + read_value: fn(&[U256]) -> V, + prefix: Nibbles, + res: &mut HashMap, +) where + F: RichField + Extendable, +{ + let load = |offset| memory.get(MemoryAddress::new(0, Segment::TrieData, offset)); + let load_slice_from = |init_offset| { + &memory.contexts[0].segments[Segment::TrieData as usize].content[init_offset..] + }; + + let trie_type = PartialTrieType::all()[load(ptr).as_usize()]; + match trie_type { + PartialTrieType::Empty => {} + PartialTrieType::Hash => {} + PartialTrieType::Branch => { + let ptr_payload = ptr + 1; + for i in 0u8..16 { + let child_ptr = load(ptr_payload + i as usize).as_usize(); + read_trie_helper::( + memory, + child_ptr, + read_value, + prefix.merge_nibble(i), + res, + ); + } + let value_ptr = load(ptr_payload + 16).as_usize(); + if value_ptr != 0 { + res.insert(prefix, read_value(load_slice_from(value_ptr))); + }; + } + PartialTrieType::Extension => { + let count = load(ptr + 1).as_usize(); + let packed = load(ptr + 2); + let nibbles = Nibbles { count, packed }; + let child_ptr = load(ptr + 3).as_usize(); + read_trie_helper::( + memory, + child_ptr, + read_value, + prefix.merge_nibbles(&nibbles), + res, + ); + } + PartialTrieType::Leaf => { + let count = load(ptr + 1).as_usize(); + let packed = load(ptr + 2); + let nibbles = Nibbles { count, packed }; + let value_ptr = load(ptr + 3).as_usize(); + res.insert( + prefix.merge_nibbles(&nibbles), + read_value(load_slice_from(value_ptr)), + ); + } + } +}