save tries to disk

evm_arithmetization/src/generation/segments.rs

@@ -91,6 +91,8 @@ pub type SegmentRunResult = Option<Box<(GenerationSegmentData, Option<Generation
 #[error("{}", .message)]
 pub struct SegmentError {
  pub message: String,
+ pub block: u64,
+ pub segment_index: usize,
  pub tries: Option<DebugTrieOutputs>,
 }
 
@@ -170,6 +172,8 @@ impl<F: RichField> SegmentDataIterator<F> {
  // In case of the error, return tries as part of the error for easier debugging
  Err(SegmentError {
  message: s,
+ block: block.as_u64(),
+ segment_index,
  tries: collect_debug_tries(self.interpreter.get_generation_state()).ok(),
  })
  }

trace_decoder/src/core.rs

@@ -347,7 +347,6 @@ fn middle<StateTrieT: StateTrie + Clone>(
  }
 
  for txn in batch {
- println!(">>>>>>>>>>>>>> Txn index: {}", txn_ix);
  let do_increment_txn_ix = txn.is_some();
  let TxnInfo {
  traces,
@@ -429,11 +428,10 @@ fn middle<StateTrieT: StateTrie + Clone>(
  acct.balance = balance.unwrap_or(acct.balance);
  //*******************************************************************
  // >>>>>>>>>>> DEBUG: Introduce error in the trie diff simulation
- println!("Address: {addr:x}, txn_ix: {txn_ix}");
- let addressdebug = "0x".to_string()+ &hex::encode(addr.0);
+ let addressdebug = "0x".to_string() + &hex::encode(addr.0);
  if addressdebug.eq("0x71f755898886f79efa73334450f05a824faeae02") {
- println!(">>>>>>>> I have found address and making a change!");
- acct.balance = acct.balance+U256::from(1);
+ println!(">>>>>>>> I have found address and making a change txn index {txn_ix} address {addressdebug}!");
+ acct.balance += U256::from(1);
  }
  //*******************************************************************
  acct.nonce = nonce.unwrap_or(acct.nonce);

zero/Cargo.toml

@@ -29,6 +29,7 @@ once_cell = { workspace = true }
 paladin-core = { workspace = true }
 plonky2 = { workspace = true }
 plonky2_maybe_rayon = { workspace = true }
+regex = "1.5.4"
 ruint = { workspace = true, features = ["num-traits", "primitive-types"] }
 serde = { workspace = true }
 serde_json = { workspace = true }

zero/src/bin/trie_diff.rs

@@ -8,9 +8,9 @@ use clap::{Parser, ValueHint};
 use futures::{future, TryStreamExt};
 use paladin::directive::{Directive, IndexedStream};
 use paladin::runtime::Runtime;
-use tracing::info;
-use mpt_trie::partial_trie::PartialTrie;
+use regex::Regex;
 use trace_decoder::observer::TriesObserver;
+use tracing::{error, info};
 use zero::ops::register;
 use zero::prover::{cli::CliProverConfig, BlockProverInput, ProverConfig};
 use zero::{prover_state::persistence::CIRCUIT_VERSION, version};
@@ -62,18 +62,19 @@ async fn main() -> Result<()> {
 
  let seg_ops = zero::ops::SegmentProofTestOnly {
  save_inputs_on_error: prover_config.save_inputs_on_error,
+ save_tries_on_error: true,
  };
 
  let des = &mut serde_json::Deserializer::from_str(&buffer);
  let block_prover_inputs = serde_path_to_error::deserialize::<_, Vec<BlockProverInput>>(des)?
  .into_iter()
  .collect::<Vec<_>>();
 
- for block in block_prover_inputs {
+ for block_prover_input in block_prover_inputs {
  let mut observer = TriesObserver::new();
  let block_generation_inputs = trace_decoder::entrypoint(
- block.block_trace,
- block.other_data,
+ block_prover_input.block_trace.clone(),
+ block_prover_input.other_data.clone(),
  prover_config.batch_size,
  &mut observer,
  )?;
@@ -83,24 +84,51 @@ async fn main() -> Result<()> {
  let simulation = Directive::map(
  IndexedStream::from(
  block_generation_inputs
+ .clone()
  .into_iter()
- .zip(repeat(prover_config.max_cpu_len_log)),
+ .enumerate()
+ .zip(repeat(prover_config.max_cpu_len_log))
+ .map(|((batch_index, inputs), max_cpu_len_log)| {
+ (inputs, max_cpu_len_log, batch_index)
+ }),
  ),
  &seg_ops,
  );
 
  if let Err(e2) = simulation
  .run(&runtime)
  .await
- .inspect_err(|e1| println!("This is error 1! {e1}"))?
+ .inspect_err(|e1| error!("Failed to run simulation, error: {e1}"))?
  .try_for_each(|_| future::ok(()))
  .await
  {
- println!("This is error 2 {e2}");
+ // Try to parse block and batch index from error message.
+ let error_message = e2.to_string();
+ let re = Regex::new(r"block:(\d+) batch:(\d+)")?;
+ if let Some(cap) = re.captures(&error_message) {
+ let block_number: u64 = cap[1].parse()?;
+ let batch_index: usize = cap[2].parse()?;
+
+ info!("Loading tries from the batch error file...");
+ let prover_tries =
+ zero::debug_utils::load_tries_from_disk(block_number, batch_index)?;
+
+ info!("Performing trie diff...");
+ zero::trie_diff::compare_tries(
+ &block_prover_input,
+ &observer.data[prover_tries.batch_index],
+ &prover_tries.tries,
+ )?;
+ } else {
+ error!(
+ "Failed to extract block and batch numbers from error message: {}",
+ error_message
+ );
+ return Err(e2);
+ }
  }
 
  info!("Trie diff finished, no problems found.")
- //trie_diff::diff::compare_tries(block_generation_inputs, ... )
  }
 
  Ok(())

zero/src/debug_utils.rs

@@ -2,9 +2,9 @@ use std::fs::{self, File};
 use std::io::{self, Write};
 use std::path::{Path, PathBuf};
 
-use serde::Serialize;
-use serde_json::Error as SerdeError;
-use thiserror::Error;
+use anyhow::Context;
+use evm_arithmetization::generation::DebugTrieOutputs;
+use serde::{Deserialize, Serialize};
 
 const DEBUG_FOLDER: &str = "./debug";
 
@@ -46,22 +46,6 @@ fn ensure_directory_exists(folder_path: &Path) -> io::Result<()> {
  }
 }
 
-/// An error type for save debug input information.
-#[derive(Error, Debug)]
-pub enum SaveInputError {
- #[error("failed to create directory '{0}'")]
- CreateDirectoryError(PathBuf, #[source] io::Error),
-
- #[error("failed to create file '{0}'")]
- CreateFileError(PathBuf, #[source] io::Error),
-
- #[error("failed to serialize inputs")]
- SerializationError(#[source] SerdeError),
-
- #[error("failed to write to file '{0}'")]
- WriteToFileError(PathBuf, #[source] io::Error),
-}
-
 /// Serializes a collection of inputs to a pretty-printed JSON format and saves
 /// them to a file.
 ///
@@ -76,27 +60,77 @@ pub enum SaveInputError {
 ///
 /// This function returns a `Result<(), std::io::Error>` indicating the
 /// operation's success or failure.
-pub fn save_inputs_to_disk<T: Serialize>(
- file_name: String,
- inputs: T,
-) -> Result<(), SaveInputError> {
+pub fn save_inputs_to_disk<T: Serialize>(file_name: String, inputs: T) -> anyhow::Result<()> {
  let debug_folder = Path::new(DEBUG_FOLDER);
+
+ // Check if output directory exists, and create one if it doesn't.
+ if !debug_folder.exists() {
+ fs::create_dir(debug_folder)?;
+ }
+
  let input_file_path = debug_folder.join(file_name);
 
  // Ensure the DEBUG_FOLDER exists
- ensure_directory_exists(debug_folder)
- .map_err(|e| SaveInputError::CreateDirectoryError(debug_folder.to_path_buf(), e))?;
+ ensure_directory_exists(debug_folder)?;
 
- let mut file = File::create(&input_file_path)
- .map_err(|e| SaveInputError::CreateFileError(input_file_path.clone(), e))?;
+ let mut file = File::create(&input_file_path)?;
 
  // Serialize the entire collection to a pretty JSON string
- let all_inputs_str =
- serde_json::to_string_pretty(&inputs).map_err(SaveInputError::SerializationError)?;
+ let all_inputs_str = serde_json::to_string_pretty(&inputs)?;
 
  // Write the serialized data to the file
- file.write_all(all_inputs_str.as_bytes())
- .map_err(|e| SaveInputError::WriteToFileError(input_file_path, e))?;
+ file.write_all(all_inputs_str.as_bytes())?;
 
  Ok(())
 }
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ErrorTriesFile {
+ pub error: String,
+ pub block_number: u64,
+ pub batch_index: usize,
+ pub tries: DebugTrieOutputs,
+}
+
+pub fn generate_tries_debug_file_name(block_number: u64, batch_index: usize) -> String {
+ format!("b{}_batch{}_error_tries.data", block_number, batch_index)
+}
+
+pub fn save_tries_to_disk(
+ err: &str,
+ block_number: u64,
+ batch_index: usize,
+ tries: &DebugTrieOutputs,
+) -> anyhow::Result<()> {
+ let output_dir = PathBuf::from(DEBUG_FOLDER);
+
+ // Check if output directory exists, and create one if it doesn't.
+ if !output_dir.exists() {
+ fs::create_dir(output_dir.clone())?;
+ }
+
+ let mut tries_debug_file_path = output_dir;
+ tries_debug_file_path.push(generate_tries_debug_file_name(block_number, batch_index));
+
+ let simulation_error_str = serde_json::to_string(&ErrorTriesFile {
+ error: err.to_string(),
+ block_number,
+ batch_index,
+ tries: tries.clone(),
+ })
+ .context("unable to serialize simulation error to save tries")?;
+ fs::write(tries_debug_file_path, simulation_error_str)
+ .expect("unable to write simulation error to file");
+ Ok(())
+}
+
+pub fn load_tries_from_disk(
+ block_number: u64,
+ batch_index: usize,
+) -> anyhow::Result<ErrorTriesFile> {
+ let mut tries_debug_file_path = PathBuf::from(DEBUG_FOLDER);
+ tries_debug_file_path.push(generate_tries_debug_file_name(block_number, batch_index));
+ let file = File::open(tries_debug_file_path)?;
+ let data: ErrorTriesFile = serde_json::from_reader(file)?;
+ Ok(data)
+}

zero/src/lib.rs

@@ -12,6 +12,7 @@ pub mod prover_state;
 pub mod provider;
 pub mod rpc;
 pub mod tracing;
+pub mod trie_diff;
 pub mod version;
 
 /// Size of the channel used to send block prover inputs to the per block

zero/src/ops.rs

@@ -2,6 +2,7 @@ zk_evm_common::check_chain_features!();
 
 use std::time::Instant;
 
+use anyhow::anyhow;
 use evm_arithmetization::generation::TrimmedGenerationInputs;
 use evm_arithmetization::proof::PublicValues;
 use evm_arithmetization::{prover::testing::simulate_execution_all_segments, GenerationInputs};
@@ -20,6 +21,7 @@ use serde::{Deserialize, Serialize};
 use tracing::error;
 use tracing::{event, info_span, Level};
 
+use crate::debug_utils::save_tries_to_disk;
 use crate::{debug_utils::save_inputs_to_disk, prover_state::p_state};
 
 registry!();
@@ -72,28 +74,54 @@ impl Operation for SegmentProof {
 #[derive(Deserialize, Serialize, RemoteExecute)]
 pub struct SegmentProofTestOnly {
  pub save_inputs_on_error: bool,
+ pub save_tries_on_error: bool,
 }
 
 impl Operation for SegmentProofTestOnly {
- type Input = (GenerationInputs<Field>, usize);
+ // The input is a tuple of the batch generation inputs, max_cpu_len_log and
+ // batch index.
+ type Input = (GenerationInputs<Field>, usize, usize);
  type Output = ();
 
  fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
- if self.save_inputs_on_error {
+ if self.save_inputs_on_error || self.save_tries_on_error {
  simulate_execution_all_segments::<Field>(inputs.0.clone(), inputs.1).map_err(|err| {
- if let Err(write_err) = save_inputs_to_disk(
- format!(
- "b{}_txns_{}..{}_input.json",
- inputs.0.block_metadata.block_number,
- inputs.0.txn_number_before,
- inputs.0.txn_number_before + inputs.0.signed_txns.len(),
- ),
- inputs.0,
- ) {
- error!("Failed to save txn proof input to disk: {:?}", write_err);
+ let block_number = inputs.0.block_metadata.block_number.as_u64();
+ let batch_index = inputs.2;
+
+ let err = if self.save_tries_on_error {
+ if let Some(ref tries) = err.tries {
+ if let Err(write_err) =
+ save_tries_to_disk(&err.to_string(), block_number, batch_index, tries)
+ {
+ error!("Failed to save tries to disk: {:?}", write_err);
+ }
+ }
+ anyhow!(
+ "block:{} batch:{} error: {}",
+ block_number,
+ batch_index,
+ err.to_string()
+ )
+ } else {
+ err.into()
+ };
+
+ if self.save_inputs_on_error {
+ if let Err(write_err) = save_inputs_to_disk(
+ format!(
+ "b{}_txns_{}..{}_input.json",
+ block_number,
+ inputs.0.txn_number_before,
+ inputs.0.txn_number_before + inputs.0.signed_txns.len(),
+ ),
+ inputs.0,
+ ) {
+ error!("Failed to save txn proof input to disk: {:?}", write_err);
+ }
  }
 
- FatalError::from_anyhow(err.into(), FatalStrategy::Terminate)
+ FatalError::from_anyhow(err, FatalStrategy::Terminate)
  })?
  } else {
  simulate_execution_all_segments::<Field>(inputs.0, inputs.1)