feat: add support for axiom-std (#30)

* chore: resolve merge * chore: merge --------- Co-authored-by: Yu Jiang Tham <[email protected]>
axiom-crypto · Jun 7, 2024 · 45ddc3e · 45ddc3e
1 parent 79563fb
commit 45ddc3e
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 11 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/cli/mod.rs b/sdk/src/cli/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,15 @@ pub fn run_cli_on_scaffold<
     cli: AxiomCircuitRunnerOptions,
 ) {
     match cli.command {
-        SnarkCmd::Mock | SnarkCmd::Run => {
+        SnarkCmd::WitnessGen => {}
+        _ => {
+            if cli.to_stdout {
+                panic!("The `to_stdout` argument is only valid for the `witness-gen` command.");
+            }
+        }
+    }
+    match cli.command {
+        SnarkCmd::Mock | SnarkCmd::Run | SnarkCmd::WitnessGen => {
             if cli.input_path.is_none() {
                 panic!("The `input_path` argument is required for the selected command.");
             }
@@ -237,6 +245,14 @@ 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")),
+                cli.to_stdout,
+            );
+        }
     }
 }
 

diff --git a/sdk/src/cli/readme.md b/sdk/src/cli/readme.md
@@ -18,9 +18,10 @@ Options:
   -n, --name <NAME>              Name of the output metadata file [default: circuit]
   -d, --data-path <DATA_PATH>    For saving build artifacts [default: data]
   -c, --config <CONFIG>          For specifying custom circuit parameters
-      --srs <SRS>                For specifying custom KZG params directory [default: params]
-      --aggregate                Whether to aggregate the output (defaults to false)
-      --auto-config-aggregation  Whether to aggregate the output (defaults to false)
+      --srs <SRS>                For specifying custom KZG params directory
+      --aggregate                Whether to aggregate the output [default: false]
+      --auto-config-aggregation  Whether to aggregate the output [default: false]
+  -s, --to-stdout                (witness-gen only) Output to stdout (true) or json (false) [default: false]
   -h, --help                     Print help
   -V, --version                  Print version
 ```
diff --git a/sdk/src/cli/types.rs b/sdk/src/cli/types.rs
@@ -14,6 +14,8 @@ pub enum SnarkCmd {
     Keygen,
     /// Generate an Axiom compute query
     Run,
+    /// Perform witness generation only, for axiom-std
+    WitnessGen,
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug)]
@@ -90,17 +92,26 @@ pub struct AxiomCircuitRunnerOptions {
 
     #[arg(
         long = "aggregate",
-        help = "Whether to aggregate the output (defaults to false)",
+        help = "Whether to aggregate the output [default: false]",
         action
     )]
     /// Whether to aggregate the output
     pub should_aggregate: bool,
 
     #[arg(
         long = "auto-config-aggregation",
-        help = "Whether to aggregate the output (defaults to false)",
+        help = "Whether to aggregate the output [default: false]",
         action
     )]
     /// Whether to auto calculate the aggregation params
     pub should_auto_config_aggregation_circuit: bool,
+
+    #[arg(
+        short = 's',
+        long = "to-stdout",
+        help = "(witness-gen only) Output to stdout (true) or json (false) [default: false]",
+        action
+    )]
+    /// Whether to output to stdout (true) or json file (false)
+    pub to_stdout: bool,
 }
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;
@@ -203,5 +203,21 @@ pub fn write_output(
     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).expect("Writing output should not fail");
+
+        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,
+    to_stdout: bool,
+) {
+    if to_stdout {
+        serde_json::to_writer_pretty(&std::io::stdout(), &output.compute_results).expect("Writing output should not fail");
+    } else {
+        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");
+    }
+}