feat: add witness-gen CLI command

circuit/src/run/inner.rs

@@ -16,7 +16,7 @@ use ethers::{
 
 use crate::{
     scaffold::{AxiomCircuit, AxiomCircuitScaffold},
-    types::{AxiomCircuitParams, AxiomCircuitPinning, AxiomV2CircuitOutput},
+    types::{AxiomCircuitParams, AxiomCircuitPinning, AxiomV2CircuitOutput, AxiomV2DataAndResults},
     utils::build_axiom_v2_compute_query,
 };
 
@@ -114,3 +114,13 @@ pub fn run<P: JsonRpcClient + Clone, S: AxiomCircuitScaffold<P, Fr>>(
     };
     output
 }
+
+pub fn witness_gen<P: JsonRpcClient + Clone, S: AxiomCircuitScaffold<P, Fr>>(
+    provider: Provider<P>,
+    pinning: AxiomCircuitPinning,
+    inputs: Option<S::InputValue>,
+) -> AxiomV2DataAndResults {
+    let runner = AxiomCircuit::<_, _, S>::prover(provider, pinning.clone()).use_inputs(inputs);
+    let output = runner.scaffold_output();
+    output
+}

circuit/src/types.rs

@@ -44,7 +44,7 @@ pub struct AxiomCircuitPinning {
 #[serde(rename_all = "camelCase")]
 pub struct AxiomV2DataAndResults {
     pub(crate) data_query: Vec<Subquery>,
-    pub(crate) compute_results: Vec<H256>,
+    pub compute_results: Vec<H256>,
 }
 
 #[derive(Debug, Serialize, Clone)]

sdk/src/cmd.rs

@@ -39,6 +39,8 @@ pub enum SnarkCmd {
     Prove,
     /// Generate an Axiom compute query
     Run,
+    /// Perform witness generation only, for axiom-std
+    WitnessGen,
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug)]
@@ -60,6 +62,7 @@ impl std::fmt::Display for SnarkCmd {
             Self::Keygen => write!(f, "keygen"),
             Self::Prove => write!(f, "prove"),
             Self::Run => write!(f, "run"),
+            Self::WitnessGen => write!(f, "witness-gen"),
         }
     }
 }
@@ -109,7 +112,7 @@ where
 {
     let cli = Cli::parse();
     match cli.command {
-        SnarkCmd::Mock | SnarkCmd::Prove | SnarkCmd::Run => {
+        SnarkCmd::Mock | SnarkCmd::Prove | SnarkCmd::Run | SnarkCmd::WitnessGen => {
             if cli.input_path.is_none() {
                 panic!("The `input_path` argument is required for the selected command.");
             }
@@ -262,5 +265,18 @@ where
                 .unwrap_or_else(|_| panic!("Could not create file at {output_json_path:?}"));
             serde_json::to_writer_pretty(&f, &output.data).expect("Writing output should not fail");
         }
+        SnarkCmd::WitnessGen => {
+            let circuit = AxiomCompute::<A>::new()
+                .use_params(params.clone())
+                .use_provider(provider.clone());
+            let (_, _, pinning) = circuit.keygen();
+            let results = circuit.use_pinning(pinning).use_inputs(input).witness_gen();
+
+            let output_path = data_path.join(PathBuf::from("compute.json"));
+            let f = File::create(&output_path)
+                .unwrap_or_else(|_| panic!("Could not create file at {output_path:?}"));
+            serde_json::to_writer_pretty(f, &results.compute_results)
+                .expect("Writing compute results should not fail");
+        }
     }
 }

sdk/src/compute.rs

@@ -13,10 +13,10 @@ use axiom_circuit::{
         utils::hilo::HiLo,
     },
     input::flatten::InputFlatten,
-    run::inner::{keygen, mock, prove, run},
+    run::inner::{keygen, mock, prove, run, witness_gen},
     scaffold::{AxiomCircuit, AxiomCircuitScaffold},
     subquery::caller::SubqueryCaller,
-    types::{AxiomCircuitParams, AxiomCircuitPinning, AxiomV2CircuitOutput},
+    types::{AxiomCircuitParams, AxiomCircuitPinning, AxiomV2CircuitOutput, AxiomV2DataAndResults},
     utils::to_hi_lo,
 };
 use ethers::providers::{Http, Provider};
@@ -227,6 +227,14 @@ where
         run::<Http, Self>(provider, self.pinning.clone().unwrap(), converted_input, pk)
     }
 
+    /// Perform witness gen only
+    pub fn witness_gen(&self) -> AxiomV2DataAndResults {
+        self.check_all_set();
+        let provider = self.provider.clone().unwrap();
+        let converted_input = self.input.clone().map(|input| input.into());
+        witness_gen::<Http, Self>(provider, self.pinning.clone().unwrap(), converted_input)
+    }
+
     /// Returns an [AxiomCircuit] instance, for functions that expect the halo2 circuit trait
     pub fn circuit(&self) -> AxiomCircuit<Fr, Http, Self> {
         self.check_provider_and_params_set();