Make test-only a CLI argument (#521)

* Make test-only a CLI argument

* Delete unused struct

* Update cargo files

* Apply offline comments

* Send dummy proof

zero_bin/leader/Cargo.toml

@@ -34,7 +34,6 @@ zero_bin_common = { workspace = true }
 
 [features]
 default = []
-test_only = ["ops/test_only", "prover/test_only"]
 
 [build-dependencies]
 cargo_metadata = { workspace = true }

zero_bin/leader/src/client.rs

@@ -63,47 +63,49 @@ pub(crate) async fn client_main(
     runtime.close().await?;
     let proved_blocks = proved_blocks?;
 
-    if cfg!(feature = "test_only") {
+    if params.prover_config.test_only {
         info!("All proof witnesses have been generated successfully.");
     } else {
         info!("All proofs have been generated successfully.");
     }
 
-    if params.keep_intermediate_proofs {
-        if params.proof_output_dir.is_some() {
-            // All proof files (including intermediary) are written to disk and kept
-            warn!("Skipping cleanup, intermediate proof files are kept");
+    if !params.prover_config.test_only {
+        if params.keep_intermediate_proofs {
+            if params.proof_output_dir.is_some() {
+                // All proof files (including intermediary) are written to disk and kept
+                warn!("Skipping cleanup, intermediate proof files are kept");
+            } else {
+                // Output all proofs to stdout
+                std::io::stdout().write_all(&serde_json::to_vec(
+                    &proved_blocks
+                        .into_iter()
+                        .filter_map(|(_, block)| block)
+                        .collect::<Vec<_>>(),
+                )?)?;
+            }
+        } else if let Some(proof_output_dir) = params.proof_output_dir.as_ref() {
+            // Remove intermediary proof files
+            proved_blocks
+                .into_iter()
+                .rev()
+                .skip(1)
+                .map(|(block_number, _)| {
+                    generate_block_proof_file_name(&proof_output_dir.to_str(), block_number)
+                })
+                .for_each(|path| {
+                    if let Err(e) = std::fs::remove_file(path) {
+                        error!("Failed to remove intermediate proof file: {e}");
+                    }
+                });
         } else {
-            // Output all proofs to stdout
-            std::io::stdout().write_all(&serde_json::to_vec(
-                &proved_blocks
-                    .into_iter()
-                    .filter_map(|(_, block)| block)
-                    .collect::<Vec<_>>(),
-            )?)?;
-        }
-    } else if let Some(proof_output_dir) = params.proof_output_dir.as_ref() {
-        // Remove intermediary proof files
-        proved_blocks
-            .into_iter()
-            .rev()
-            .skip(1)
-            .map(|(block_number, _)| {
-                generate_block_proof_file_name(&proof_output_dir.to_str(), block_number)
-            })
-            .for_each(|path| {
-                if let Err(e) = std::fs::remove_file(path) {
-                    error!("Failed to remove intermediate proof file: {e}");
-                }
-            });
-    } else {
-        // Output only last proof to stdout
-        if let Some(last_block) = proved_blocks
-            .into_iter()
-            .filter_map(|(_, block)| block)
-            .last()
-        {
-            std::io::stdout().write_all(&serde_json::to_vec(&last_block)?)?;
+            // Output only last proof to stdout
+            if let Some(last_block) = proved_blocks
+                .into_iter()
+                .filter_map(|(_, block)| block)
+                .last()
+            {
+                std::io::stdout().write_all(&serde_json::to_vec(&last_block)?)?;
+            }
         }
     }
 

zero_bin/leader/src/http.rs

@@ -71,15 +71,27 @@ async fn prove(
 
     let block_number = payload.prover_input.get_block_number();
 
-    match payload
-        .prover_input
-        .prove(
-            &runtime,
-            payload.previous.map(futures::future::ok),
-            prover_config,
-        )
-        .await
-    {
+    let proof_res = if prover_config.test_only {
+        payload
+            .prover_input
+            .prove_test(
+                &runtime,
+                payload.previous.map(futures::future::ok),
+                prover_config,
+            )
+            .await
+    } else {
+        payload
+            .prover_input
+            .prove(
+                &runtime,
+                payload.previous.map(futures::future::ok),
+                prover_config,
+            )
+            .await
+    };
+
+    match proof_res {
         Ok(b_proof) => match write_to_file(output_dir, block_number, &b_proof) {
             Ok(file) => {
                 info!("Successfully wrote proof to {}", file.display());

zero_bin/leader/src/main.rs

@@ -9,6 +9,7 @@ use dotenvy::dotenv;
 use ops::register;
 use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
+use prover::ProverConfig;
 use tracing::{info, warn};
 use zero_bin_common::{
     block_interval::BlockInterval, prover_state::persistence::set_circuit_cache_dir_env_if_not_set,
@@ -53,22 +54,25 @@ async fn main() -> Result<()> {
     }
 
     let args = cli::Cli::parse();
-    if let paladin::config::Runtime::InMemory = args.paladin.runtime {
-        // If running in emulation mode, we'll need to initialize the prover
-        // state here.
-        args.prover_state_config
-            .into_prover_state_manager()
-            .initialize()?;
-    }
 
     let runtime = Runtime::from_config(&args.paladin, register()).await?;
 
-    let cli_prover_config = args.prover_config;
+    let prover_config: ProverConfig = args.prover_config.into();
+
+    // If not in test_only mode and running in emulation mode, we'll need to
+    // initialize the prover state here.
+    if !prover_config.test_only {
+        if let paladin::config::Runtime::InMemory = args.paladin.runtime {
+            args.prover_state_config
+                .into_prover_state_manager()
+                .initialize()?;
+        }
+    }
 
     match args.command {
         Command::Stdio { previous_proof } => {
             let previous_proof = get_previous_proof(previous_proof)?;
-            stdio::stdio_main(runtime, previous_proof, cli_prover_config.into()).await?;
+            stdio::stdio_main(runtime, previous_proof, prover_config).await?;
         }
         Command::Http { port, output_dir } => {
             // check if output_dir exists, is a directory, and is writable
@@ -80,7 +84,7 @@ async fn main() -> Result<()> {
                 panic!("output-dir is not a writable directory");
             }
 
-            http::http_main(runtime, port, output_dir, cli_prover_config.into()).await?;
+            http::http_main(runtime, port, output_dir, prover_config).await?;
         }
         Command::Rpc {
             rpc_url,
@@ -120,7 +124,7 @@ async fn main() -> Result<()> {
                     checkpoint_block_number,
                     previous_proof,
                     proof_output_dir,
-                    prover_config: cli_prover_config.into(),
+                    prover_config,
                     keep_intermediate_proofs,
                 },
             )

zero_bin/leader/src/stdio.rs

@@ -26,7 +26,7 @@ pub(crate) async fn stdio_main(
     runtime.close().await?;
     let proved_blocks = proved_blocks?;
 
-    if cfg!(feature = "test_only") {
+    if prover_config.test_only {
         info!("All proof witnesses have been generated successfully.");
     } else {
         info!("All proofs have been generated successfully.");

zero_bin/ops/Cargo.toml

@@ -21,4 +21,3 @@ zero_bin_common = { path = "../common" }
 
 [features]
 default = []
-test_only = []

zero_bin/ops/src/lib.rs

@@ -1,16 +1,12 @@
-#[cfg(not(feature = "test_only"))]
 use std::time::Instant;
 
-#[cfg(not(feature = "test_only"))]
 use evm_arithmetization::generation::TrimmedGenerationInputs;
 use evm_arithmetization::proof::PublicValues;
-#[cfg(feature = "test_only")]
 use evm_arithmetization::{prover::testing::simulate_execution_all_segments, GenerationInputs};
 use paladin::{
     operation::{FatalError, FatalStrategy, Monoid, Operation, Result},
     registry, RemoteExecute,
 };
-#[cfg(feature = "test_only")]
 use proof_gen::types::Field;
 use proof_gen::{
     proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof},
@@ -20,7 +16,6 @@ use proof_gen::{
 };
 use serde::{Deserialize, Serialize};
 use tracing::error;
-#[cfg(not(feature = "test_only"))]
 use tracing::{event, info_span, Level};
 use zero_bin_common::{debug_utils::save_inputs_to_disk, prover_state::p_state};
 
@@ -31,7 +26,6 @@ pub struct SegmentProof {
     pub save_inputs_on_error: bool,
 }
 
-#[cfg(not(feature = "test_only"))]
 impl Operation for SegmentProof {
     type Input = evm_arithmetization::AllData;
     type Output = proof_gen::proof_types::SegmentAggregatableProof;
@@ -72,8 +66,12 @@ impl Operation for SegmentProof {
     }
 }
 
-#[cfg(feature = "test_only")]
-impl Operation for SegmentProof {
+#[derive(Deserialize, Serialize, RemoteExecute)]
+pub struct SegmentProofTestOnly {
+    pub save_inputs_on_error: bool,
+}
+
+impl Operation for SegmentProofTestOnly {
     type Input = (GenerationInputs, usize);
     type Output = ();
 
@@ -107,14 +105,12 @@ impl Operation for SegmentProof {
 ///
 /// - When created, it starts a span with the transaction proof id.
 /// - When dropped, it logs the time taken by the transaction proof.
-#[cfg(not(feature = "test_only"))]
 struct SegmentProofSpan {
     _span: tracing::span::EnteredSpan,
     start: Instant,
     descriptor: String,
 }
 
-#[cfg(not(feature = "test_only"))]
 impl SegmentProofSpan {
     /// Get a unique id for the transaction proof.
     fn get_id(ir: &TrimmedGenerationInputs, segment_index: usize) -> String {
@@ -174,7 +170,6 @@ impl SegmentProofSpan {
     }
 }
 
-#[cfg(not(feature = "test_only"))]
 impl Drop for SegmentProofSpan {
     fn drop(&mut self) {
         event!(

zero_bin/prover/Cargo.toml

@@ -30,4 +30,3 @@ clap = {workspace = true}
 
 [features]
 default = []
-test_only = ["ops/test_only"]

zero_bin/prover/src/cli.rs

@@ -14,6 +14,10 @@ pub struct CliProverConfig {
     /// If true, save the public inputs to disk on error.
     #[arg(short='i', long, help_heading = HELP_HEADING, default_value_t = false)]
     save_inputs_on_error: bool,
+    /// If true, only test the trace decoder and witness generation without
+    /// generating a proof.
+    #[arg(long, help_heading = HELP_HEADING, default_value_t = false)]
+    test_only: bool,
 }
 
 impl From<CliProverConfig> for crate::ProverConfig {
@@ -22,6 +26,7 @@ impl From<CliProverConfig> for crate::ProverConfig {
             batch_size: cli.batch_size,
             max_cpu_len_log: cli.max_cpu_len_log,
             save_inputs_on_error: cli.save_inputs_on_error,
+            test_only: cli.test_only,
         }
     }
 }