feat: trie diff tool

evm_arithmetization/src/fixed_recursive_verifier.rs

@@ -37,7 +37,7 @@ use starky::stark::Stark;
 
 use crate::all_stark::{all_cross_table_lookups, AllStark, Table, NUM_TABLES};
 use crate::cpu::kernel::aggregator::KERNEL;
-use crate::generation::segments::{GenerationSegmentData, SegmentDataIterator, SegmentError};
+use crate::generation::segments::{GenerationSegmentData, SegmentDataIterator};
 use crate::generation::{GenerationInputs, TrimmedGenerationInputs};
 use crate::get_challenges::observe_public_values_target;
 use crate::proof::{
@@ -1889,8 +1889,7 @@ where
         let mut proofs = vec![];
 
         for segment_run in segment_iterator {
-            let (_, mut next_data) =
-                segment_run.map_err(|e: SegmentError| anyhow::format_err!(e))?;
+            let (_, mut next_data) = segment_run?;
             let proof = self.prove_segment(
                 all_stark,
                 config,

evm_arithmetization/src/generation/mod.rs

@@ -1,9 +1,10 @@
 use std::collections::HashMap;
+use std::fmt::Display;
 
 use anyhow::anyhow;
 use ethereum_types::{Address, BigEndianHash, H256, U256};
 use keccak_hash::keccak;
-use log::log_enabled;
+use log::error;
 use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
 use plonky2::field::extension::Extendable;
 use plonky2::field::polynomial::PolynomialValues;
@@ -51,6 +52,29 @@ pub const NUM_EXTRA_CYCLES_BEFORE: usize = 64;
 /// Memory values used to initialize `MemBefore`.
 pub type MemBeforeValues = Vec<(MemoryAddress, U256)>;
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ErrorWithTries<E = anyhow::Error> {
+    pub inner: E,
+    pub tries: Option<DebugOutputTries>,
+}
+impl<E: Display> Display for ErrorWithTries<E> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.inner.fmt(f)
+    }
+}
+
+impl<E: std::error::Error> std::error::Error for ErrorWithTries<E> {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        self.inner.source()
+    }
+}
+
+impl<E> ErrorWithTries<E> {
+    pub fn new(inner: E, tries: Option<DebugOutputTries>) -> Self {
+        Self { inner, tries }
+    }
+}
+
 /// Inputs needed for trace generation.
 #[derive(Clone, Debug, Deserialize, Serialize, Default)]
 #[serde(bound = "")]
@@ -234,6 +258,15 @@ impl<F: RichField> GenerationInputs<F> {
     }
 }
 
+/// Post transaction execution tries retrieved from the prover's memory.
+/// Used primarily for error debugging in case of a failed execution.
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct DebugOutputTries {
+    pub state_trie: HashedPartialTrie,
+    pub transaction_trie: HashedPartialTrie,
+    pub receipt_trie: HashedPartialTrie,
+}
+
 fn apply_metadata_and_tries_memops<F: RichField + Extendable<D>, const D: usize>(
     state: &mut GenerationState<F>,
     inputs: &TrimmedGenerationInputs<F>,
@@ -492,10 +525,8 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         "simulate CPU",
         simulate_cpu(&mut state, *max_cpu_len_log)
     );
-    if cpu_res.is_err() {
-        output_debug_tries(&state)?;
-        cpu_res?;
-    };
+
+    cpu_res?;
 
     let trace_lengths = state.traces.get_lengths();
 
@@ -595,59 +626,50 @@ fn simulate_cpu<F: RichField>(
     Ok((final_registers, mem_after))
 }
 
-/// Outputs the tries that have been obtained post transaction execution, as
+/// Collects the tries that have been obtained post transaction execution, as
 /// they are represented in the prover's memory.
-/// This will do nothing if the CPU execution failed outside of the final trie
-/// root checks.
-pub(crate) fn output_debug_tries<F: RichField>(state: &GenerationState<F>) -> anyhow::Result<()> {
-    if !log_enabled!(log::Level::Debug) {
-        return Ok(());
-    }
-
-    // Retrieve previous PC (before jumping to KernelPanic), to see if we reached
-    // `perform_final_checks`. We will output debugging information on the final
-    // tries only if we got a root mismatch.
-    let previous_pc = state.get_registers().program_counter;
-
-    let label = KERNEL.offset_name(previous_pc);
-
-    if label.contains("check_state_trie")
-        || label.contains("check_txn_trie")
-        || label.contains("check_receipt_trie")
-    {
-        let state_trie_ptr = u256_to_usize(
-            state
-                .memory
-                .read_global_metadata(GlobalMetadata::StateTrieRoot),
-        )
-        .map_err(|_| anyhow!("State trie pointer is too large to fit in a usize."))?;
-        log::debug!(
-            "Computed state trie: {:?}",
-            get_state_trie::<HashedPartialTrie>(&state.memory, state_trie_ptr)
-        );
-
-        let txn_trie_ptr = u256_to_usize(
-            state
-                .memory
-                .read_global_metadata(GlobalMetadata::TransactionTrieRoot),
-        )
-        .map_err(|_| anyhow!("Transactions trie pointer is too large to fit in a usize."))?;
-        log::debug!(
-            "Computed transactions trie: {:?}",
-            get_txn_trie::<HashedPartialTrie>(&state.memory, txn_trie_ptr)
-        );
-
-        let receipt_trie_ptr = u256_to_usize(
-            state
-                .memory
-                .read_global_metadata(GlobalMetadata::ReceiptTrieRoot),
-        )
-        .map_err(|_| anyhow!("Receipts trie pointer is too large to fit in a usize."))?;
-        log::debug!(
-            "Computed receipts trie: {:?}",
-            get_receipt_trie::<HashedPartialTrie>(&state.memory, receipt_trie_ptr)
-        );
-    }
-
-    Ok(())
+pub(crate) fn collect_debug_tries<F: RichField>(
+    state: &GenerationState<F>,
+) -> Option<DebugOutputTries> {
+    let state_trie_ptr = u256_to_usize(
+        state
+            .memory
+            .read_global_metadata(GlobalMetadata::StateTrieRoot),
+    )
+    .inspect_err(|e| error!("failed to retrieve state trie pointer: {e:?}"))
+    .ok()?;
+
+    let state_trie = get_state_trie::<HashedPartialTrie>(&state.memory, state_trie_ptr)
+        .inspect_err(|e| error!("unable to retrieve state trie for debugging purposes: {e:?}"))
+        .ok()?;
+
+    let txn_trie_ptr = u256_to_usize(
+        state
+            .memory
+            .read_global_metadata(GlobalMetadata::TransactionTrieRoot),
+    )
+    .inspect_err(|e| error!("failed to retrieve transactions trie pointer: {e:?}"))
+    .ok()?;
+    let transaction_trie = get_txn_trie::<HashedPartialTrie>(&state.memory, txn_trie_ptr)
+        .inspect_err(|e| {
+            error!("unable to retrieve transaction trie for debugging purposes: {e:?}",)
+        })
+        .ok()?;
+
+    let receipt_trie_ptr = u256_to_usize(
+        state
+            .memory
+            .read_global_metadata(GlobalMetadata::ReceiptTrieRoot),
+    )
+    .inspect_err(|e| error!("failed to retrieve receipts trie pointer: {e:?}"))
+    .ok()?;
+    let receipt_trie = get_receipt_trie::<HashedPartialTrie>(&state.memory, receipt_trie_ptr)
+        .inspect_err(|e| error!("unable to retrieve receipt trie for debugging purposes: {e:?}"))
+        .ok()?;
+
+    Some(DebugOutputTries {
+        state_trie,
+        transaction_trie,
+        receipt_trie,
+    })
 }

evm_arithmetization/src/generation/segments.rs

@@ -9,7 +9,7 @@ use super::TrimmedGenerationInputs;
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::interpreter::{set_registers_and_run, ExtraSegmentData, Interpreter};
 use crate::generation::state::State;
-use crate::generation::{debug_inputs, GenerationInputs};
+use crate::generation::{collect_debug_tries, debug_inputs, ErrorWithTries, GenerationInputs};
 use crate::witness::memory::MemoryState;
 use crate::witness::state::RegistersState;
 
@@ -88,8 +88,10 @@ pub struct SegmentDataIterator<F: RichField> {
 pub type SegmentRunResult = Option<Box<(GenerationSegmentData, Option<GenerationSegmentData>)>>;
 
 #[derive(thiserror::Error, Debug, Serialize, Deserialize)]
-#[error("{}", .0)]
-pub struct SegmentError(pub String);
+#[error("{}", .message)]
+pub struct SegmentError {
+    pub message: String,
+}
 
 impl<F: RichField> SegmentDataIterator<F> {
     pub fn new(inputs: &GenerationInputs<F>, max_cpu_len_log: Option<usize>) -> Self {
@@ -113,7 +115,7 @@ impl<F: RichField> SegmentDataIterator<F> {
     fn generate_next_segment(
         &mut self,
         partial_segment_data: Option<GenerationSegmentData>,
-    ) -> Result<SegmentRunResult, SegmentError> {
+    ) -> Result<SegmentRunResult, ErrorWithTries<SegmentError>> {
         // Get the (partial) current segment data, if it is provided. Otherwise,
         // initialize it.
         let mut segment_data = if let Some(partial) = partial_segment_data {
@@ -133,8 +135,9 @@ impl<F: RichField> SegmentDataIterator<F> {
 
         // Run the interpreter to get `registers_after` and the partial data for the
         // next segment.
-        let run = set_registers_and_run(segment_data.registers_after, &mut self.interpreter);
-        if let Ok((updated_registers, mem_after)) = run {
+        let execution_result =
+            set_registers_and_run(segment_data.registers_after, &mut self.interpreter);
+        if let Ok((updated_registers, mem_after)) = execution_result {
             let partial_segment_data = Some(build_segment_data(
                 segment_index + 1,
                 Some(updated_registers),
@@ -157,21 +160,28 @@ impl<F: RichField> SegmentDataIterator<F> {
                     inputs.txn_number_before + inputs.txn_hashes.len()
                 ),
             };
-            let s = format!(
-                "Segment generation {:?} for block {:?} ({}) failed with error {:?}",
-                segment_index,
-                block,
-                txn_range,
-                run.unwrap_err()
-            );
-            Err(SegmentError(s))
+            // In case of the error, return tries as part of the error for easier debugging.
+            Err(ErrorWithTries::new(
+                SegmentError {
+                    message: format!(
+                        "Segment generation {:?} for block:{} batch:{} tx_range:({}) failed with error {:?}",
+                        segment_index,
+                        block.low_u64(),
+                        segment_index,
+                        txn_range,
+                        execution_result.unwrap_err()
+                    ),
+                },
+                collect_debug_tries(self.interpreter.get_generation_state()),
+            ))
         }
     }
 }
 
 /// Returned type from a `SegmentDataIterator`, needed to prove all segments in
 /// a transaction batch.
-pub type AllData<F> = Result<(TrimmedGenerationInputs<F>, GenerationSegmentData), SegmentError>;
+pub type AllData<F> =
+    Result<(TrimmedGenerationInputs<F>, GenerationSegmentData), ErrorWithTries<SegmentError>>;
 
 impl<F: RichField> Iterator for SegmentDataIterator<F> {
     type Item = AllData<F>;

evm_arithmetization/src/prover.rs

@@ -370,13 +370,10 @@ pub(crate) fn features_check<F: RichField>(inputs: &TrimmedGenerationInputs<F>)
 /// A utility module designed to test witness generation externally.
 pub mod testing {
     use super::*;
+    use crate::generation::ErrorWithTries;
     use crate::{
         cpu::kernel::interpreter::Interpreter,
-        generation::{
-            output_debug_tries,
-            segments::{SegmentDataIterator, SegmentError},
-            state::State,
-        },
+        generation::segments::{SegmentDataIterator, SegmentError},
     };
 
     /// Simulates the zkEVM CPU execution.
@@ -388,13 +385,7 @@ pub mod testing {
         let initial_offset = KERNEL.global_labels["init"];
         let mut interpreter: Interpreter<F> =
             Interpreter::new_with_generation_inputs(initial_offset, initial_stack, &inputs, None);
-        let result = interpreter.run();
-
-        if result.is_err() {
-            output_debug_tries(interpreter.get_generation_state())?;
-        }
-
-        result?;
+        interpreter.run()?;
         Ok(())
     }
 
@@ -415,8 +406,7 @@ pub mod testing {
         let mut proofs = vec![];
 
         for segment_run in segment_data_iterator {
-            let (_, mut next_data) =
-                segment_run.map_err(|e: SegmentError| anyhow::format_err!(e))?;
+            let (_, mut next_data) = segment_run?;
             let proof = prove(
                 all_stark,
                 config,
@@ -434,16 +424,14 @@ pub mod testing {
     pub fn simulate_execution_all_segments<F>(
         inputs: GenerationInputs<F>,
         max_cpu_len_log: usize,
-    ) -> Result<()>
+    ) -> Result<(), ErrorWithTries<SegmentError>>
     where
         F: RichField,
     {
         features_check(&inputs.clone().trim());
 
         for segment in SegmentDataIterator::<F>::new(&inputs, Some(max_cpu_len_log)) {
-            if let Err(e) = segment {
-                return Err(anyhow::format_err!(e));
-            }
+            segment?;
         }
 
         Ok(())

evm_arithmetization/src/public_types.rs

@@ -39,7 +39,10 @@ pub type ProofWithPublicInputs =
 /// proofs.
 pub type PublicValues = crate::proof::PublicValues<Field>;
 
-pub type AllData = Result<(TrimmedGenerationInputs, GenerationSegmentData), SegmentError>;
+pub type AllData = Result<
+    (TrimmedGenerationInputs, GenerationSegmentData),
+    crate::generation::ErrorWithTries<SegmentError>,
+>;
 
 /// Returned type from the zkEVM STARK prover, before recursive verification.
 pub type AllProof = crate::proof::AllProof<Field, RecursionConfig, EXTENSION_DEGREE>;

mpt_trie/src/debug_tools/diff.rs

@@ -136,13 +136,15 @@ impl Display for DiffPoint {
 /// Meta information for a node in a trie.
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 pub struct NodeInfo {
-    key: Nibbles,
-
+    /// Mpt trie node key.
+    pub key: Nibbles,
     /// The direct value associated with the node (only applicable to `Leaf` &
     /// `Branch` nodes).
-    value: Option<Vec<u8>>,
-    node_type: TrieNodeType,
-    hash: H256,
+    pub value: Option<Vec<u8>>,
+    /// Type of this node.
+    pub node_type: TrieNodeType,
+    /// Node hash.
+    pub hash: H256,
 }
 
 impl Display for NodeInfo {