feat: add witness gen CLI command

axiom-crypto · May 21, 2024 · 16a0fd7 · 16a0fd7
1 parent fe21d14
commit 16a0fd7
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 5 deletions.
diff --git a/circuit/src/run/inner.rs b/circuit/src/run/inner.rs
@@ -14,7 +14,7 @@ use log::info;
 
 use crate::{
     scaffold::{AxiomCircuit, AxiomCircuitScaffold},
-    types::{AxiomCircuitParams, AxiomCircuitPinning, AxiomV2CircuitOutput},
+    types::{AxiomCircuitParams, AxiomCircuitPinning, AxiomV2CircuitOutput, AxiomV2DataAndResults},
     utils::{
         build_axiom_v2_compute_query, check_compute_proof_format, check_compute_query_format,
         get_query_schema_from_compute_query, verify_snark, DK,
@@ -134,3 +134,10 @@ pub fn run<P: JsonRpcClient + Clone, S: AxiomCircuitScaffold<P, Fr>>(
 
     circuit_output
 }
+
+pub fn witness_gen<P: JsonRpcClient + Clone, S: AxiomCircuitScaffold<P, Fr>>(
+    circuit: &mut AxiomCircuit<Fr, P, S>,
+) -> AxiomV2DataAndResults {
+    let output = circuit.scaffold_output();
+    output
+}
diff --git a/sdk/src/run/mod.rs b/sdk/src/run/mod.rs
@@ -16,7 +16,7 @@ use axiom_circuit::{
     },
     run::{
         aggregation::{agg_circuit_keygen, agg_circuit_run},
-        inner::{keygen, mock, run},
+        inner::{keygen, mock, run, witness_gen},
     },
     scaffold::{AxiomCircuit, AxiomCircuitScaffold},
     types::AxiomCircuitParams,
@@ -34,7 +34,7 @@ use crate::{
     utils::{
         io::{
             read_agg_pk_and_pinning, read_metadata, read_pk_and_pinning, write_agg_keygen_output,
-            write_keygen_output, write_metadata, write_output,
+            write_keygen_output, write_metadata, write_output, write_witness_gen_output,
         },
         read_srs_from_dir_or_install,
     },
@@ -49,7 +49,7 @@ pub fn run_cli_on_scaffold<
     cli: AxiomCircuitRunnerOptions,
 ) {
     match cli.command {
-        SnarkCmd::Mock | SnarkCmd::Prove => {
+        SnarkCmd::Mock | SnarkCmd::Prove | SnarkCmd::WitnessGen => {
             if cli.input_path.is_none() {
                 panic!("The `input_path` argument is required for the selected command.");
             }
@@ -234,6 +234,10 @@ pub fn run_cli_on_scaffold<
                 cli.data_path.join(PathBuf::from("output.json")),
             );
         }
+        SnarkCmd::WitnessGen => {
+            let output = witness_gen(&mut runner);
+            write_witness_gen_output(output, cli.data_path.join(PathBuf::from("compute.json")));
+        }
     }
 }
 

diff --git a/sdk/src/run/types.rs b/sdk/src/run/types.rs
@@ -14,6 +14,8 @@ pub enum SnarkCmd {
     Keygen,
     /// Generate an Axiom compute query
     Prove,
+    /// Perform witness generation only, for axiom-std
+    WitnessGen,
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug)]

diff --git a/sdk/src/utils/io.rs b/sdk/src/utils/io.rs
@@ -17,7 +17,7 @@ use axiom_circuit::{
     scaffold::{AxiomCircuit, AxiomCircuitScaffold},
     types::{
         AggregationCircuitPinning, AxiomCircuitPinning, AxiomClientCircuitMetadata,
-        AxiomV2CircuitOutput,
+        AxiomV2CircuitOutput, AxiomV2DataAndResults
     },
 };
 use ethers::providers::Http;
@@ -205,3 +205,13 @@ pub fn write_output(
         .unwrap_or_else(|_| panic!("Could not create file at {json_output_path:?}"));
     serde_json::to_writer_pretty(&f, &output).expect("Writing output should not fail");
 }
+
+pub fn write_witness_gen_output(
+    output: AxiomV2DataAndResults,
+    json_output_path: PathBuf,
+) {
+    info!("Writing JSON output to {:?}", &json_output_path);
+    let f = File::create(&json_output_path)
+        .unwrap_or_else(|_| panic!("Could not create file at {json_output_path:?}"));
+    serde_json::to_writer_pretty(&f, &output.compute_results).expect("Writing output should not fail");
+}