cli: feat: add commit scheme option (#297)

Co-authored-by: Jason <[email protected]>
DelphinusLab · Oct 10, 2024 · 97734e2 · 97734e2
1 parent c6db01a
commit 97734e2
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 8 deletions.
diff --git a/crates/cli/src/app_builder.rs b/crates/cli/src/app_builder.rs
@@ -11,6 +11,7 @@ use clap::Command;
 use clap::ValueHint;
 
 use crate::args::HostMode;
+use crate::args::Scheme;
 use crate::command::DryRunArg;
 use crate::command::ProveArg;
 use crate::command::RunningArg;
@@ -197,6 +198,14 @@ fn setup_command() -> Command<'static> {
             ).takes_value(true)
             .value_delimiter(',')
             .required(false)
+        )
+        .arg(
+            arg!(
+                --scheme <SCHEME> "Specify polynomial commitment scheme"
+            )
+            .default_value("shplonk")
+            .value_parser(value_parser!(Scheme))
+            .required(false),
         );
 
     let command = if cfg!(not(feature = "uniform-circuit")) {
@@ -273,6 +282,7 @@ impl From<&ArgMatches> for SetupArg {
                 .map(|v| v.to_string())
                 .collect::<Vec<_>>(),
             wasm_image: WasmImageArg::parse(val),
+            scheme: *val.get_one::<Scheme>("scheme").unwrap(),
         }
     }
 }

diff --git a/crates/cli/src/args.rs b/crates/cli/src/args.rs
@@ -1,12 +1,30 @@
+use circuits_batcher::args::OpenSchema;
 use serde::Deserialize;
 use serde::Serialize;
 
 #[derive(clap::ArgEnum, Copy, Clone, Debug, Default, Serialize, Deserialize)]
-pub enum HostMode {
+pub(crate) enum HostMode {
     /// Trivial Wasm Host Environment
     #[default]
     Default,
 
     /// Wasm Host Environment with more Zk plugins
     Standard,
 }
+
+#[derive(clap::ArgEnum, Copy, Clone, Debug, Default, Serialize, Deserialize)]
+pub(crate) enum Scheme {
+    Gwc,
+
+    #[default]
+    Shplonk,
+}
+
+impl From<Scheme> for OpenSchema {
+    fn from(val: Scheme) -> OpenSchema {
+        match val {
+            Scheme::Gwc => OpenSchema::GWC,
+            Scheme::Shplonk => OpenSchema::Shplonk,
+        }
+    }
+}
diff --git a/crates/cli/src/command.rs b/crates/cli/src/command.rs
@@ -5,6 +5,7 @@ use std::io::Read;
 use std::path::Path;
 use std::path::PathBuf;
 
+use crate::args::Scheme;
 use crate::config::CircuitDataConfig;
 use crate::config::CircuitDataMd5;
 use crate::TRIVIAL_WASM;
@@ -38,6 +39,7 @@ pub(crate) struct SetupArg {
     pub(crate) host_mode: HostMode,
     pub(crate) phantom_functions: Vec<String>,
     pub(crate) wasm_image: Option<PathBuf>,
+    pub(crate) scheme: Scheme,
 }
 
 impl SetupArg {
@@ -215,6 +217,8 @@ impl SetupArg {
                 checksum,
                 phantom_functions: self.phantom_functions.clone(),
                 host_mode: self.host_mode,
+
+                scheme: self.scheme,
             };
             config.write(&mut File::create(&config_path)?)?;
 

diff --git a/crates/cli/src/config.rs b/crates/cli/src/config.rs
@@ -7,7 +7,6 @@ use std::path::PathBuf;
 
 use anyhow::Result;
 use circuits_batcher::args::HashType;
-use circuits_batcher::args::OpenSchema;
 use circuits_batcher::proof::ProofGenerationInfo;
 use circuits_batcher::proof::ProofInfo;
 use circuits_batcher::proof::ProofPieceInfo;
@@ -31,6 +30,7 @@ use serde::Serialize;
 use specs::slice_backend::SliceBackendBuilder;
 
 use crate::args::HostMode;
+use crate::args::Scheme;
 use crate::names::name_of_circuit_data;
 use crate::names::name_of_etable_slice;
 use crate::names::name_of_external_host_call_table_slice;
@@ -74,6 +74,8 @@ pub(crate) struct Config {
     pub(crate) checksum: (String, String),
     pub(crate) phantom_functions: Vec<String>,
     pub(crate) host_mode: HostMode,
+
+    pub(crate) scheme: Scheme,
 }
 
 impl Config {
@@ -413,7 +415,7 @@ impl Config {
                     &params,
                     &cached_proving_key.as_ref().unwrap().1,
                     proof_load_info.hashtype,
-                    OpenSchema::Shplonk,
+                    self.scheme.into(),
                 ),
                 ZkWasmCircuit::LastSliceCircuit(circuit) => proof_piece_info
                     .create_proof::<Bn256, _>(
@@ -422,7 +424,7 @@ impl Config {
                         &params,
                         &cached_proving_key.as_ref().unwrap().1,
                         proof_load_info.hashtype,
-                        OpenSchema::Shplonk,
+                        self.scheme.into(),
                     ),
             };
 
@@ -509,7 +511,7 @@ impl Config {
             };
 
             proof
-                .verify_proof(&params_verifier, OpenSchema::Shplonk)
+                .verify_proof(&params_verifier, self.scheme.into())
                 .unwrap();
 
             // TODO: handle checksum sanity check

diff --git a/test_cli.sh b/test_cli.sh
@@ -6,11 +6,12 @@ set -e
 set -x
 
 CUDA="--features cuda"
+SCHEME="--scheme shplonk"
 
 test_default_cli() {
     cargo build --release $CUDA
     rm -rf params/*.data params/*.config output
-    $CLI --params ./params wasm_output setup --wasm ./crates/zkwasm/wasm/wasm_output.wasm
+    $CLI --params ./params wasm_output setup --wasm ./crates/zkwasm/wasm/wasm_output.wasm $SCHEME
     $CLI --params ./params wasm_output dry-run --wasm crates/zkwasm/wasm/wasm_output.wasm --public 133:i64 --public 2:i64 --output ./output
     $CLI --params ./params wasm_output prove --wasm crates/zkwasm/wasm/wasm_output.wasm --public 133:i64 --public 2:i64 --output ./output
     $CLI --params ./params wasm_output verify --output ./output
@@ -19,7 +20,7 @@ test_default_cli() {
 test_uniform_circuit_cli() {
     cargo build --release --features uniform-circuit $CUDA
     rm -rf params/*.data params/*.config output
-    $CLI --params ./params wasm_output setup
+    $CLI --params ./params wasm_output setup $SCHEME
     $CLI --params ./params wasm_output dry-run --wasm crates/zkwasm/wasm/wasm_output.wasm --public 133:i64 --public 2:i64 --output ./output
     $CLI --params ./params wasm_output prove --wasm crates/zkwasm/wasm/wasm_output.wasm --public 133:i64 --public 2:i64 --output ./output
     $CLI --params ./params wasm_output verify --output ./output
@@ -28,7 +29,7 @@ test_uniform_circuit_cli() {
 test_continuation_cli() {
     cargo build --release --features continuation $CUDA
     rm -rf params/*.data params/*.config output
-    $CLI --params ./params fibonacci setup
+    $CLI --params ./params fibonacci setup $SCHEME
     $CLI --params ./params fibonacci dry-run --wasm crates/zkwasm/wasm/fibonacci.wasm --public 25:i64 --output ./output
     $CLI --params ./params fibonacci prove --wasm crates/zkwasm/wasm/fibonacci.wasm --public 25:i64 --output ./output
     $CLI --params ./params fibonacci verify --output ./output