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/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, AllData};
-#[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,19 +16,16 @@ 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!();
 
-#[cfg(feature = "test_only")]
 #[derive(Deserialize, Serialize, RemoteExecute)]
 pub struct BatchTestOnly {
     pub save_inputs_on_error: bool,
 }
 
-#[cfg(feature = "test_only")]
 impl Operation for BatchTestOnly {
     type Input = (GenerationInputs, usize);
     type Output = ();
@@ -50,7 +43,6 @@ pub struct SegmentProof {
     pub save_inputs_on_error: bool,
 }
 
-#[cfg(not(feature = "test_only"))]
 impl Operation for SegmentProof {
     type Input = AllData;
     type Output = proof_gen::proof_types::SegmentAggregatableProof;
@@ -88,8 +80,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 = AllData;
     type Output = ();
 
@@ -102,14 +98,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 {
@@ -169,7 +163,6 @@ impl SegmentProofSpan {
     }
 }
 
-#[cfg(not(feature = "test_only"))]
 impl Drop for SegmentProofSpan {
     fn drop(&mut self) {
         event!(

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();
@@ -209,28 +210,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)));
 

zero_bin/tools/prove_rpc.sh

@@ -17,18 +17,8 @@ export RUST_LOG=info
 # See also .cargo/config.toml.
 export RUSTFLAGS='-C target-cpu=native -Zlinker-features=-lld'
 
-if [[ $8 == "test_only" ]]; then
-  # Circuit sizes don't matter in test_only mode, so we keep them minimal.
-    export ARITHMETIC_CIRCUIT_SIZE="16..17"
-    export BYTE_PACKING_CIRCUIT_SIZE="9..10"
-    export CPU_CIRCUIT_SIZE="12..13"
-    export KECCAK_CIRCUIT_SIZE="4..5"
-    export KECCAK_SPONGE_CIRCUIT_SIZE="9..10"
-    export LOGIC_CIRCUIT_SIZE="12..13"
-    export MEMORY_CIRCUIT_SIZE="17..18"
-    export MEMORY_BEFORE_CIRCUIT_SIZE="7..8"
-    export MEMORY_AFTER_CIRCUIT_SIZE="7..8"
-else
+# Circuit sizes only matter in non test_only mode.
+if ! [[ $8 == "test_only" ]]; then
     export ARITHMETIC_CIRCUIT_SIZE="16..21"
     export BYTE_PACKING_CIRCUIT_SIZE="8..21"
     export CPU_CIRCUIT_SIZE="8..21"
@@ -112,7 +102,7 @@ fi
 if [[ $8 == "test_only" ]]; then
     # test only run
     echo "Proving blocks ${BLOCK_INTERVAL} in a test_only mode now... (Total: ${TOT_BLOCKS})"
-    command='cargo r --release --features test_only --bin leader -- --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
+    command='cargo r --release --bin leader -- --test-only --runtime in-memory --load-strategy on-demand rpc --rpc-type "$NODE_RPC_TYPE" --rpc-url "$NODE_RPC_URL" --block-interval $BLOCK_INTERVAL --proof-output-dir $PROOF_OUTPUT_DIR $PREV_PROOF_EXTRA_ARG --backoff "$BACKOFF" --max-retries "$RETRIES" '
     if [ "$OUTPUT_TO_TERMINAL" = true ]; then
         eval $command
         retVal=$?

zero_bin/tools/prove_stdio.sh

@@ -44,18 +44,8 @@ if [[ $INPUT_FILE == "" ]]; then
     exit 1
 fi
 
-if [[ $TEST_ONLY == "test_only" ]]; then
-    # Circuit sizes don't matter in test_only mode, so we keep them minimal.
-    export ARITHMETIC_CIRCUIT_SIZE="16..17"
-    export BYTE_PACKING_CIRCUIT_SIZE="9..10"
-    export CPU_CIRCUIT_SIZE="12..13"
-    export KECCAK_CIRCUIT_SIZE="4..5"
-    export KECCAK_SPONGE_CIRCUIT_SIZE="9..10"
-    export LOGIC_CIRCUIT_SIZE="12..13"
-    export MEMORY_CIRCUIT_SIZE="17..18"
-    export MEMORY_BEFORE_CIRCUIT_SIZE="7..8"
-    export MEMORY_AFTER_CIRCUIT_SIZE="7..8"
-else
+# Circuit sizes only matter in non test_only mode.
+if ! [[ $TEST_ONLY == "test_only" ]]; then
     if [[ $INPUT_FILE == *"witness_b19807080"* ]]; then
       # These sizes are configured specifically for block 19807080. Don't use this in other scenarios
         echo "Using specific circuit sizes for witness_b19807080.json"
@@ -98,7 +88,7 @@ fi
 # proof. This is useful for quickly testing decoding and all of the
 # other non-proving code.
 if [[ $TEST_ONLY == "test_only" ]]; then
-    cargo run --release --features test_only --bin leader -- --runtime in-memory --load-strategy on-demand stdio < $INPUT_FILE &> $TEST_OUT_PATH
+    cargo run --release --bin leader -- --test-only --runtime in-memory --load-strategy on-demand stdio < $INPUT_FILE &> $TEST_OUT_PATH
     if grep -q 'All proof witnesses have been generated successfully.' $TEST_OUT_PATH; then
         echo -e "\n\nSuccess - Note this was just a test, not a proof"
         rm $TEST_OUT_PATH