Some tooling for debugging tests where the updated tries are not correct

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,

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<F: RichField + Extendable<D>, const D: usize>(
         ))
     };
 
+    log::debug!(
+        "Updated state trie:\n{:#?}",
+        read_trie::<F, AccountRlp, D>(
+            &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)),

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<F, V, const D: usize>(
+    memory: &MemoryState,
+    ptr: usize,
+    read_value: fn(&[U256]) -> V,
+) -> HashMap<Nibbles, V>
+where
+    F: RichField + Extendable<D>,
+{
+    let mut res = HashMap::new();
+    let empty_nibbles = Nibbles {
+        count: 0,
+        packed: U256::zero(),
+    };
+    read_trie_helper::<F, V, D>(memory, ptr, read_value, empty_nibbles, &mut res);
+    res
+}
+
+pub(crate) fn read_trie_helper<F, V, const D: usize>(
+    memory: &MemoryState,
+    ptr: usize,
+    read_value: fn(&[U256]) -> V,
+    prefix: Nibbles,
+    res: &mut HashMap<Nibbles, V>,
+) where
+    F: RichField + Extendable<D>,
+{
+    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::<F, V, D>(
+                    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::<F, V, D>(
+                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)),
+            );
+        }
+    }
+}