Aligned adapted prover with run_and_prove (#385)

starkware-libs · Jan 27, 2025 · e264df3 · e264df3
1 parent 8369484
commit e264df3
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 20 deletions.
diff --git a/stwo_cairo_prover/crates/adapted_prover/src/main.rs b/stwo_cairo_prover/crates/adapted_prover/src/main.rs
@@ -32,6 +32,7 @@ struct Args {
     track_relations: bool,
     #[structopt(long = "display_components")]
     display_components: bool,
+    /// Verify the generated proof.
     #[structopt(long = "verify")]
     verify: bool,
 }
@@ -40,16 +41,16 @@ struct Args {
 enum Error {
     #[error("Invalid arguments: {0}")]
     Cli(#[from] clap::Error),
-    #[error("VM import failed: {0}")]
-    VmImport(#[from] VmImportError),
+    #[error("IO failed: {0}")]
+    IO(#[from] std::io::Error),
     #[error("Proving failed: {0}")]
     Proving(#[from] ProvingError),
-    #[error("Verification failed: {0}")]
-    Verification(#[from] CairoVerificationError),
     #[error("Serialization failed: {0}")]
     Serde(#[from] serde_json::error::Error),
-    #[error("IO failed: {0}")]
-    IO(#[from] std::io::Error),
+    #[error("Verification failed: {0}")]
+    Verification(#[from] CairoVerificationError),
+    #[error("VM import failed: {0}")]
+    VmImport(#[from] VmImportError),
 }
 
 fn main() -> ExitCode {
@@ -68,7 +69,7 @@ fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
         .build();
 
     log::info!(
-        "Casm states by opcode:\n {},",
+        "Casm states by opcode:\n{}",
         vm_output.state_transitions.casm_states_by_opcode
     );
 

diff --git a/stwo_cairo_prover/crates/run_and_prove/src/main.rs b/stwo_cairo_prover/crates/run_and_prove/src/main.rs
@@ -2,17 +2,19 @@ use std::path::PathBuf;
 use std::process::ExitCode;
 
 use clap::Parser;
-use stwo_cairo_prover::cairo_air::air::CairoProof;
-use stwo_cairo_prover::cairo_air::{prove_cairo, ConfigBuilder};
+use stwo_cairo_prover::cairo_air::{
+    prove_cairo, verify_cairo, CairoVerificationError, ConfigBuilder,
+};
 use stwo_cairo_prover::input::plain::adapt_finished_runner;
 use stwo_cairo_prover::input::vm_import::VmImportError;
 use stwo_cairo_utils::binary_utils::run_binary;
 use stwo_cairo_utils::vm_utils::{run_vm, VmArgs, VmError};
 use stwo_prover::core::prover::ProvingError;
-use stwo_prover::core::vcs::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
+use stwo_prover::core::vcs::blake2_merkle::Blake2sMerkleChannel;
 use thiserror::Error;
 use tracing::{span, Level};
 
+// TODO(yuval): unite this and adapted_prover to a single binary, or at least share more code.
 /// Command line arguments for run_and_prove.
 /// Example command line (use absolute paths):
 ///     ```
@@ -31,29 +33,34 @@ struct Args {
     track_relations: bool,
     #[structopt(long = "display_components")]
     display_components: bool,
+    /// Verify the generated proof.
+    #[structopt(long = "verify")]
+    verify: bool,
 }
 
 #[derive(Debug, Error)]
 enum Error {
-    #[error("Invalid arguments")]
+    #[error("Invalid arguments: {0}")]
     Cli(#[from] clap::Error),
-    #[error("Failed to interact with the file system")]
+    #[error("IO failed: {0}")]
     IO(#[from] std::io::Error),
-    #[error("VM run failed: {0}")]
-    Vm(#[from] VmError),
+    #[error("Proving failed: {0}")]
+    Proving(#[from] ProvingError),
     #[error("Serialization failed: {0}")]
     Serde(#[from] serde_json::error::Error),
+    #[error("Verification failed: {0}")]
+    Verification(#[from] CairoVerificationError),
+    #[error("VM run failed: {0}")]
+    Vm(#[from] VmError),
     #[error("VM import failed: {0}")]
     VmImport(#[from] VmImportError),
-    #[error("Proving failed: {0}")]
-    Proving(#[from] ProvingError),
 }
 
 fn main() -> ExitCode {
     run_binary(run, "run_and_prove")
 }
 
-fn run(args: impl Iterator<Item = String>) -> Result<CairoProof<Blake2sMerkleHasher>, Error> {
+fn run(args: impl Iterator<Item = String>) -> Result<(), Error> {
     let _span = span!(Level::INFO, "run").entered();
     let args = Args::try_parse_from(args)?;
     let cairo_runner = run_vm(&args.vm_args)?;
@@ -63,12 +70,20 @@ fn run(args: impl Iterator<Item = String>) -> Result<CairoProof<Blake2sMerkleHas
         .display_components(args.display_components)
         .build();
 
-    let casm_states_by_opcode_count = &cairo_input.state_transitions.casm_states_by_opcode.counts();
-    log::info!("Casm states by opcode count: {casm_states_by_opcode_count:?}");
+    log::info!(
+        "Casm states by opcode:\n{}",
+        cairo_input.state_transitions.casm_states_by_opcode
+    );
 
+    // TODO(Ohad): Propagate hash from CLI args.
     let proof = prove_cairo::<Blake2sMerkleChannel>(cairo_input, prover_config)?;
 
     std::fs::write(args.proof_path, serde_json::to_string(&proof)?)?;
 
-    Ok(proof)
+    if args.verify {
+        verify_cairo::<Blake2sMerkleChannel>(proof)?;
+        log::info!("Proof verified successfully");
+    }
+
+    Ok(())
 }