feat: reorganize cli params (#432)

* feat: reorganize cli params

* fix: conficts

* fix: review

* fix: review 2

zero_bin/leader/src/cli.rs

@@ -2,6 +2,7 @@ use std::path::PathBuf;
 
 use alloy::transports::http::reqwest::Url;
 use clap::{Parser, Subcommand, ValueHint};
+use prover::cli::CliProverConfig;
 use rpc::RpcType;
 use zero_bin_common::prover_state::cli::CliProverStateConfig;
 
@@ -14,6 +15,9 @@ pub(crate) struct Cli {
     #[clap(flatten)]
     pub(crate) paladin: paladin::config::Config,
 
+    #[clap(flatten)]
+    pub(crate) prover_config: CliProverConfig,
+
     // Note this is only relevant for the leader when running in in-memory
     // mode.
     #[clap(flatten)]
@@ -28,17 +32,6 @@ pub(crate) enum Command {
         /// The previous proof output.
         #[arg(long, short = 'f', value_hint = ValueHint::FilePath)]
         previous_proof: Option<PathBuf>,
-        #[arg(short, long, default_value_t = 20)]
-        max_cpu_len_log: usize,
-        /// Number of transactions in a batch to process at once
-        #[arg(short, long, default_value_t = 1)]
-        batch_size: usize,
-        /// Number of segments to keep in memory and prove in parallel
-        #[arg(long, default_value_t = 64)]
-        segment_chunk_size: usize,
-        /// If true, save the public inputs to disk on error.
-        #[arg(short, long, default_value_t = false)]
-        save_inputs_on_error: bool,
     },
     /// Reads input from a node rpc and writes output to stdout.
     Rpc {
@@ -61,21 +54,9 @@ pub(crate) enum Command {
         /// stdout.
         #[arg(long, short = 'o', value_hint = ValueHint::FilePath)]
         proof_output_dir: Option<PathBuf>,
-        /// The log of the max number of CPU cycles per proof.
-        #[arg(short, long, default_value_t = 20)]
-        max_cpu_len_log: usize,
-        /// Number of transactions in a batch to process at once
-        #[arg(short, long, default_value_t = 1)]
-        batch_size: usize,
-        /// Number of segments to keep in memory and prove in parallel
-        #[arg(long, default_value_t = 64)]
-        segment_chunk_size: usize,
-        /// If true, save the public inputs to disk on error.
-        #[arg(short, long, default_value_t = false)]
-        save_inputs_on_error: bool,
         /// Network block time in milliseconds. This value is used
         /// to determine the blockchain node polling interval.
-        #[arg(long, env = "ZERO_BIN_BLOCK_TIME", default_value_t = 2000)]
+        #[arg(short, long, env = "ZERO_BIN_BLOCK_TIME", default_value_t = 2000)]
         block_time: u64,
         /// Keep intermediate proofs. Default action is to
         /// delete them after the final proof is generated.
@@ -101,16 +82,5 @@ pub(crate) enum Command {
         /// The directory to which output should be written.
         #[arg(short, long, value_hint = ValueHint::DirPath)]
         output_dir: PathBuf,
-        #[arg(short, long, default_value_t = 20)]
-        max_cpu_len_log: usize,
-        /// Number of transactions in a batch to process at once
-        #[arg(short, long, default_value_t = 1)]
-        batch_size: usize,
-        /// Number of segments to keep in memory and prove in parallel
-        #[arg(long, default_value_t = 64)]
-        segment_chunk_size: usize,
-        /// If true, save the public inputs to disk on error.
-        #[arg(short, long, default_value_t = false)]
-        save_inputs_on_error: bool,
     },
 }

zero_bin/leader/src/client.rs

@@ -5,7 +5,7 @@ use alloy::transports::http::reqwest::Url;
 use anyhow::Result;
 use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
-use prover::ProverParams;
+use prover::ProverConfig;
 use rpc::{retry::build_http_retry_provider, RpcType};
 use tracing::{error, info, warn};
 use zero_bin_common::block_interval::BlockInterval;
@@ -19,15 +19,12 @@ pub struct RpcParams {
     pub max_retries: u32,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct ProofParams {
     pub checkpoint_block_number: u64,
     pub previous_proof: Option<GeneratedBlockProof>,
     pub proof_output_dir: Option<PathBuf>,
-    pub max_cpu_len_log: usize,
-    pub batch_size: usize,
-    pub segment_chunk_size: usize,
-    pub save_inputs_on_error: bool,
+    pub prover_config: ProverConfig,
     pub keep_intermediate_proofs: bool,
 }
 
@@ -58,13 +55,8 @@ pub(crate) async fn client_main(
     let proved_blocks = prover_input
         .prove(
             &runtime,
-            params.max_cpu_len_log,
             params.previous_proof.take(),
-            ProverParams {
-                batch_size: params.batch_size,
-                segment_chunk_size: params.segment_chunk_size,
-            },
-            params.save_inputs_on_error,
+            params.prover_config,
             params.proof_output_dir.clone(),
         )
         .await;

zero_bin/leader/src/http.rs

@@ -5,7 +5,7 @@ use anyhow::{bail, Result};
 use axum::{http::StatusCode, routing::post, Json, Router};
 use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
-use prover::{BlockProverInput, ProverParams};
+use prover::{BlockProverInput, ProverConfig};
 use serde::{Deserialize, Serialize};
 use serde_json::to_writer;
 use tracing::{debug, error, info};
@@ -15,9 +15,7 @@ pub(crate) async fn http_main(
     runtime: Runtime,
     port: u16,
     output_dir: PathBuf,
-    max_cpu_len_log: usize,
-    prover_params: ProverParams,
-    save_inputs_on_error: bool,
+    prover_config: ProverConfig,
 ) -> Result<()> {
     let addr = SocketAddr::from(([0, 0, 0, 0], port));
     debug!("listening on {}", addr);
@@ -27,16 +25,7 @@ pub(crate) async fn http_main(
         "/prove",
         post({
             let runtime = runtime.clone();
-            move |body| {
-                prove(
-                    body,
-                    runtime,
-                    output_dir.clone(),
-                    max_cpu_len_log,
-                    prover_params,
-                    save_inputs_on_error,
-                )
-            }
+            move |body| prove(body, runtime, output_dir.clone(), prover_config)
         }),
     );
     let listener = tokio::net::TcpListener::bind(&addr).await?;
@@ -76,9 +65,7 @@ async fn prove(
     Json(payload): Json<HttpProverInput>,
     runtime: Arc<Runtime>,
     output_dir: PathBuf,
-    max_cpu_len_log: usize,
-    prover_params: ProverParams,
-    save_inputs_on_error: bool,
+    prover_config: ProverConfig,
 ) -> StatusCode {
     debug!("Received payload: {:#?}", payload);
 
@@ -88,10 +75,8 @@ async fn prove(
         .prover_input
         .prove(
             &runtime,
-            max_cpu_len_log,
             payload.previous.map(futures::future::ok),
-            prover_params,
-            save_inputs_on_error,
+            prover_config,
         )
         .await
     {

zero_bin/leader/src/main.rs

@@ -62,35 +62,14 @@ async fn main() -> Result<()> {
 
     let runtime = Runtime::from_config(&args.paladin, register()).await?;
 
+    let cli_prover_config = args.prover_config;
+
     match args.command {
-        Command::Stdio {
-            previous_proof,
-            max_cpu_len_log,
-            batch_size,
-            segment_chunk_size,
-            save_inputs_on_error,
-        } => {
+        Command::Stdio { previous_proof } => {
             let previous_proof = get_previous_proof(previous_proof)?;
-            stdio::stdio_main(
-                runtime,
-                max_cpu_len_log,
-                previous_proof,
-                prover::ProverParams {
-                    batch_size,
-                    segment_chunk_size,
-                },
-                save_inputs_on_error,
-            )
-            .await?;
+            stdio::stdio_main(runtime, previous_proof, cli_prover_config.into()).await?;
         }
-        Command::Http {
-            port,
-            output_dir,
-            max_cpu_len_log,
-            batch_size,
-            segment_chunk_size,
-            save_inputs_on_error,
-        } => {
+        Command::Http { port, output_dir } => {
             // check if output_dir exists, is a directory, and is writable
             let output_dir_metadata = std::fs::metadata(&output_dir);
             if output_dir_metadata.is_err() {
@@ -100,18 +79,7 @@ async fn main() -> Result<()> {
                 panic!("output-dir is not a writable directory");
             }
 
-            http::http_main(
-                runtime,
-                port,
-                output_dir,
-                max_cpu_len_log,
-                prover::ProverParams {
-                    batch_size,
-                    segment_chunk_size,
-                },
-                save_inputs_on_error,
-            )
-            .await?;
+            http::http_main(runtime, port, output_dir, cli_prover_config.into()).await?;
         }
         Command::Rpc {
             rpc_url,
@@ -120,10 +88,6 @@ async fn main() -> Result<()> {
             checkpoint_block_number,
             previous_proof,
             proof_output_dir,
-            max_cpu_len_log,
-            batch_size,
-            segment_chunk_size,
-            save_inputs_on_error,
             block_time,
             keep_intermediate_proofs,
             backoff,
@@ -154,10 +118,7 @@ async fn main() -> Result<()> {
                     checkpoint_block_number,
                     previous_proof,
                     proof_output_dir,
-                    max_cpu_len_log,
-                    batch_size,
-                    segment_chunk_size,
-                    save_inputs_on_error,
+                    prover_config: cli_prover_config.into(),
                     keep_intermediate_proofs,
                 },
             )

zero_bin/leader/src/stdio.rs

@@ -3,16 +3,14 @@ use std::io::{Read, Write};
 use anyhow::Result;
 use paladin::runtime::Runtime;
 use proof_gen::proof_types::GeneratedBlockProof;
-use prover::{ProverInput, ProverParams};
+use prover::{ProverConfig, ProverInput};
 use tracing::info;
 
 /// The main function for the stdio mode.
 pub(crate) async fn stdio_main(
     runtime: Runtime,
-    max_cpu_len_log: usize,
     previous: Option<GeneratedBlockProof>,
-    prover_params: ProverParams,
-    save_inputs_on_error: bool,
+    prover_config: ProverConfig,
 ) -> Result<()> {
     let mut buffer = String::new();
     std::io::stdin().read_to_string(&mut buffer)?;
@@ -23,14 +21,7 @@ pub(crate) async fn stdio_main(
     };
 
     let proved_blocks = prover_input
-        .prove(
-            &runtime,
-            max_cpu_len_log,
-            previous,
-            prover_params,
-            save_inputs_on_error,
-            None,
-        )
+        .prove(&runtime, previous, prover_config, None)
         .await;
     runtime.close().await?;
     let proved_blocks = proved_blocks?;

zero_bin/prover/Cargo.toml

@@ -25,6 +25,7 @@ ruint = { workspace = true, features = ["num-traits", "primitive-types"] }
 ops = { workspace = true }
 zero_bin_common ={ workspace = true }
 num-traits = { workspace = true }
+clap = {workspace = true}
 
 [features]
 default = []

zero_bin/prover/src/cli.rs

@@ -0,0 +1,31 @@
+use clap::Args;
+
+const HELP_HEADING: &str = "Prover options";
+
+/// Represents the main configuration structure for the runtime.
+#[derive(Args, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Default)]
+pub struct CliProverConfig {
+    /// The log of the max number of CPU cycles per proof.
+    #[arg(short, long, help_heading = HELP_HEADING, default_value_t = 20)]
+    max_cpu_len_log: usize,
+    /// Number of transactions in a batch to process at once.
+    #[arg(short, long, help_heading = HELP_HEADING, default_value_t = 1)]
+    batch_size: usize,
+    /// Number of segments to keep in memory and prove in parallel.
+    #[arg(short='c', long, help_heading = HELP_HEADING, default_value_t = 64)]
+    segment_chunk_size: usize,
+    /// 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,
+}
+
+impl From<CliProverConfig> for crate::ProverConfig {
+    fn from(cli: CliProverConfig) -> Self {
+        Self {
+            segment_chunk_size: cli.segment_chunk_size,
+            batch_size: cli.batch_size,
+            max_cpu_len_log: cli.max_cpu_len_log,
+            save_inputs_on_error: cli.save_inputs_on_error,
+        }
+    }
+}