Fix trie output when debugging final root mismatch (#86)

* Fix receipt trie computation for multiple logs

* Have the simulation also returns tries upon root mismatch failure

* Nit

* Add trie specific labels

* Add doc

* Changelog

* Address comments

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Create subtries without ever hashing leaves ([#76](https://github.com/0xPolygonZero/zk_evm/pull/76))
 - Fix generation inputs logging pre-transaction execution ([#89](https://github.com/0xPolygonZero/zk_evm/pull/89))
 - Reduce state trie size for dummy payloads ([#88](https://github.com/0xPolygonZero/zk_evm/pull/88))
+- Fix post-txn trie debugging output for multi-logs receipts ([#86](https://github.com/0xPolygonZero/zk_evm/pull/86))
 
 ## [0.1.1] - 2024-03-01
 

evm_arithmetization/src/cpu/kernel/asm/main.asm

@@ -85,15 +85,18 @@ global execute_withdrawals:
 execute_withdrawals_post_stack_op:
     %withdrawals
 
-global hash_final_tries:
+global perform_final_checks:
     // stack: cum_gas, txn_counter, num_nibbles, txn_nb
     // Check that we end up with the correct `cum_gas`, `txn_nb` and bloom filter.
     %mload_global_metadata(@GLOBAL_METADATA_BLOCK_GAS_USED_AFTER) %assert_eq
     DUP3 %mload_global_metadata(@GLOBAL_METADATA_TXN_NUMBER_AFTER) %assert_eq
     %pop3
-    PUSH 1 // initial trie data length 
+    PUSH 1 // initial trie data length
+global check_state_trie:
     %mpt_hash_state_trie   %mload_global_metadata(@GLOBAL_METADATA_STATE_TRIE_DIGEST_AFTER)     %assert_eq
+global check_txn_trie:
     %mpt_hash_txn_trie     %mload_global_metadata(@GLOBAL_METADATA_TXN_TRIE_DIGEST_AFTER)       %assert_eq
+global check_receipt_trie:
     %mpt_hash_receipt_trie %mload_global_metadata(@GLOBAL_METADATA_RECEIPT_TRIE_DIGEST_AFTER)   %assert_eq
     // We don't need the trie data length here.
     POP

evm_arithmetization/src/generation/mod.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use anyhow::anyhow;
 use ethereum_types::{Address, BigEndianHash, H256, U256};
+use log::log_enabled;
 use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
 use plonky2::field::extension::Extendable;
 use plonky2::field::polynomial::PolynomialValues;
@@ -229,51 +230,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
 
     let cpu_res = timed!(timing, "simulate CPU", simulate_cpu(&mut state));
     if cpu_res.is_err() {
-        // Retrieve previous PC (before jumping to KernelPanic), to see if we reached
-        // `hash_final_tries`. We will output debugging information on the final
-        // tries only if we got a root mismatch.
-        let previous_pc = state
-            .traces
-            .cpu
-            .last()
-            .expect("We should have CPU rows")
-            .program_counter
-            .to_canonical_u64() as usize;
-
-        if KERNEL.offset_name(previous_pc).contains("hash_final_tries") {
-            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)
-            );
-        }
+        output_debug_tries(&state);
 
         cpu_res?;
     }
@@ -348,3 +305,60 @@ fn simulate_cpu<F: Field>(state: &mut GenerationState<F>) -> anyhow::Result<()>
 
     Ok(())
 }
+
+/// Outputs 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(())
+}

evm_arithmetization/src/generation/state.rs

@@ -175,8 +175,6 @@ pub(crate) trait State<F: Field> {
 
             self.transition()?;
         }
-
-        Ok(())
     }
 
     fn handle_error(&mut self, err: ProgramError) -> anyhow::Result<()>

evm_arithmetization/src/generation/trie_extractor.rs

@@ -159,22 +159,31 @@ pub(crate) fn read_logs(
     (0..num_logs)
         .map(|_| {
             let address = u256_to_h160(slice[offset].unwrap_or_default())?;
-            let num_topics = u256_to_usize(slice[offset + 1].unwrap_or_default())?;
+            offset += 1;
+
+            let num_topics = u256_to_usize(slice[offset].unwrap_or_default())?;
+            offset += 1;
 
             let topics = (0..num_topics)
-                .map(|i| H256::from_uint(&slice[offset + 2 + i].unwrap_or_default()))
+                .map(|i| H256::from_uint(&slice[offset + i].unwrap_or_default()))
                 .collect();
+            offset += num_topics;
+
+            let data_len = u256_to_usize(slice[offset].unwrap_or_default())?;
+            offset += 1;
+
+            let data = slice[offset..offset + data_len]
+                .iter()
+                .map(|&x| u256_to_u8(x.unwrap_or_default()))
+                .collect::<Result<_, _>>()?;
+            offset += data_len + 1; // We need to skip one extra element before looping.
 
-            let data_len = u256_to_usize(slice[offset + 2 + num_topics].unwrap_or_default())?;
             let log = LogRlp {
                 address,
                 topics,
-                data: slice[offset + 2 + num_topics + 1..offset + 2 + num_topics + 1 + data_len]
-                    .iter()
-                    .map(|&x| u256_to_u8(x.unwrap_or_default()))
-                    .collect::<Result<_, _>>()?,
+                data,
             };
-            offset += 2 + num_topics + 1 + data_len;
+
             Ok(log)
         })
         .collect()

evm_arithmetization/src/prover.rs

@@ -368,7 +368,10 @@ pub fn check_abort_signal(abort_signal: Option<Arc<AtomicBool>>) -> Result<()> {
 /// A utility module designed to test witness generation externally.
 pub mod testing {
     use super::*;
-    use crate::cpu::kernel::interpreter::Interpreter;
+    use crate::{
+        cpu::kernel::interpreter::Interpreter,
+        generation::{output_debug_tries, state::State},
+    };
 
     /// Simulates the zkEVM CPU execution.
     /// It does not generate any trace or proof of correct state transition.
@@ -377,6 +380,11 @@ pub mod testing {
         let initial_offset = KERNEL.global_labels["main"];
         let mut interpreter: Interpreter<F> =
             Interpreter::new_with_generation_inputs(initial_offset, initial_stack, inputs);
-        interpreter.run()
+        let result = interpreter.run();
+        if result.is_err() {
+            output_debug_tries(interpreter.get_generation_state())?;
+        }
+
+        Ok(())
     }
 }