Make test-only a CLI argument

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,7 +63,7 @@ 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.");

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,18 +16,33 @@ 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};
 
 registry!();
 
+#[derive(Deserialize, Serialize, RemoteExecute)]
+pub struct BatchTestOnly {
+    pub save_inputs_on_error: bool,
+}
+
+impl Operation for BatchTestOnly {
+    type Input = (GenerationInputs, usize);
+    type Output = ();
+
+    fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
+        simulate_execution_all_segments::<Field>(inputs.0, inputs.1)
+            .map_err(|err| FatalError::from_anyhow(err, FatalStrategy::Terminate))?;
+
+        Ok(())
+    }
+}
+
 #[derive(Deserialize, Serialize, RemoteExecute)]
 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 +83,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 +122,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 +187,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,
         }
     }
 }

zero_bin/prover/src/lib.rs

@@ -21,6 +21,7 @@ pub struct ProverConfig {
     pub batch_size: usize,
     pub max_cpu_len_log: usize,
     pub save_inputs_on_error: bool,
+    pub test_only: bool,
 }
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -34,7 +35,6 @@ impl BlockProverInput {
         self.other_data.b_data.b_meta.block_number.into()
     }
 
-    #[cfg(not(feature = "test_only"))]
     pub async fn prove(
         self,
         runtime: &Runtime,
@@ -50,6 +50,7 @@ impl BlockProverInput {
             max_cpu_len_log,
             batch_size,
             save_inputs_on_error,
+            test_only: _,
         } = prover_config;
 
         let block_number = self.get_block_number();
@@ -123,8 +124,7 @@ impl BlockProverInput {
         }
     }
 
-    #[cfg(feature = "test_only")]
-    pub async fn prove(
+    pub async fn prove_test(
         self,
         runtime: &Runtime,
         previous: Option<impl Future<Output = Result<GeneratedBlockProof>>>,
@@ -139,6 +139,7 @@ impl BlockProverInput {
             max_cpu_len_log,
             batch_size,
             save_inputs_on_error,
+            test_only: _,
         } = prover_config;
 
         let block_number = self.get_block_number();
@@ -147,7 +148,7 @@ impl BlockProverInput {
         let block_generation_inputs =
             trace_decoder::entrypoint(self.block_trace, self.other_data, batch_size)?;
 
-        let batch_ops = ops::SegmentProof {
+        let seg_ops = ops::SegmentProofTestOnly {
             save_inputs_on_error,
         };
 
@@ -157,7 +158,7 @@ impl BlockProverInput {
                     .into_iter()
                     .zip(repeat(max_cpu_len_log)),
             ),
-            &batch_ops,
+            &seg_ops,
         );
 
         simulation
@@ -211,28 +212,45 @@ impl ProverInput {
 
                 // Prove the block
                 let proof_output_dir = proof_output_dir.clone();
-                let fut = block
-                    .prove(runtime, prev.take(), prover_config)
-                    .then(move |proof| async move {
-                        let proof = proof?;
-                        let block_number = proof.b_height;
-
-                        // Write latest generated proof to disk if proof_output_dir is provided
-                        let return_proof: Option<GeneratedBlockProof> =
-                            if proof_output_dir.is_some() {
-                                ProverInput::write_proof(proof_output_dir, &proof).await?;
-                                None
-                            } else {
-                                Some(proof.clone())
-                            };
-
-                        if tx.send(proof).is_err() {
-                            anyhow::bail!("Failed to send proof");
-                        }
-
-                        Ok((block_number, return_proof))
-                    })
-                    .boxed();
+                let fut = if prover_config.test_only {
+                    block
+                        .prove_test(runtime, prev.take(), prover_config)
+                        .then(move |proof| async move {
+                            let proof = proof?;
+                            let block_number = proof.b_height;
+
+                            if tx.send(proof).is_err() {
+                                anyhow::bail!("Bad dummy proof?");
+                            }
+
+                            // We ignore the returned dummy proof in test-only mode.
+                            Ok((block_number, None))
+                        })
+                        .boxed()
+                } else {
+                    block
+                        .prove(runtime, prev.take(), prover_config)
+                        .then(move |proof| async move {
+                            let proof = proof?;
+                            let block_number = proof.b_height;
+
+                            // Write latest generated proof to disk if proof_output_dir is provided.
+                            let return_proof: Option<GeneratedBlockProof> =
+                                if proof_output_dir.is_some() {
+                                    ProverInput::write_proof(proof_output_dir, &proof).await?;
+                                    None
+                                } else {
+                                    Some(proof.clone())
+                                };
+
+                            if tx.send(proof).is_err() {
+                                anyhow::bail!("Failed to send proof");
+                            }
+
+                            Ok((block_number, return_proof))
+                        })
+                        .boxed()
+                };
 
                 prev = Some(Box::pin(rx.map_err(anyhow::Error::new)));