feat: wrap up groth16_prover_zkey_file, build for macOS, iOS device, Android aarch64

.cargo/config.toml

@@ -0,0 +1,2 @@
+[env]
+RUST_WITNESS_LINK_TEST_WITNESS="1"

Cargo.toml

@@ -7,12 +7,11 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.95"
-ark-bn254 = "0.4.0"
-ark-circom = { git = "https://github.com/zkmopro/circom-compat.git", version = "0.1.0", branch = "wasm-delete" }
 num-bigint = "0.4.6"
 rust-witness = "0.1.2"
 serde = { version = "1.0.217", features = ["derive"] }
 serde_json = "1.0.135"
+num-traits = "0.2.19"
 
 [build-dependencies]
 cc = "1.0"

build.rs

@@ -1,8 +1,9 @@
-use std::path::PathBuf;
+use std::{env, fs, path::{Path, PathBuf}};
 
 fn main() {
-    // #[cfg(test)]
-    rust_witness::transpile::transpile_wasm("./test-vectors".to_string());
+    if let Ok(_) = std::env::var("RUST_WITNESS_LINK_TEST_WITNESS") {
+        rust_witness::transpile::transpile_wasm("./test-vectors".to_string());
+    }
 
     let target = std::env::var("TARGET").unwrap();
     let arch = target.split('-').next().unwrap();
@@ -39,4 +40,29 @@ fn main() {
     println!("cargo:rustc-link-lib=static=fr");
     println!("cargo:rustc-link-lib=static=fq");
     println!("cargo:rustc-link-lib=static=gmp");
+
+    // refer to https://github.com/bbqsrc/cargo-ndk to see how to link the libc++_shared.so file in Android
+    if env::var("CARGO_CFG_TARGET_OS").unwrap() == "android" {
+        android();
+    }
+}
+
+fn android() {
+    println!("cargo:rustc-link-lib=c++_shared");
+
+    if let Ok(output_path) = env::var("CARGO_NDK_OUTPUT_PATH") {
+        let sysroot_libs_path = PathBuf::from(env::var_os("CARGO_NDK_SYSROOT_LIBS_PATH").unwrap());
+        let lib_path = sysroot_libs_path.join("libc++_shared.so");
+        assert!(
+            lib_path.exists(),
+            "Error: Source file {:?} does not exist",
+            lib_path
+        );
+        let dest_dir = Path::new(&output_path).join(&env::var("CARGO_NDK_ANDROID_TARGET").unwrap());
+        println!("cargo:rerun-if-changed={}", dest_dir.display());
+        if !dest_dir.exists() {
+            fs::create_dir_all(&dest_dir).unwrap();
+        }
+        fs::copy(lib_path, Path::new(&dest_dir).join("libc++_shared.so")).unwrap();
+    }
 }

src/lib.rs

@@ -1,169 +1,141 @@
 use std::collections::HashMap;
-use std::ffi::CString;
-use std::fs::File;
-use std::os::raw::c_char;
-use std::os::raw::c_uint;
 use std::str::FromStr;
 
-use anyhow::Context;
 use anyhow::Result;
-use ark_bn254::Bn254;
-use ark_circom::read_proving_key;
-use ark_circom::ZkeyHeaderReader;
 use num_bigint::BigInt;
-use serde::Deserialize;
-use serde::Serialize;
-
 pub type WtnsFn = fn(HashMap<String, Vec<BigInt>>) -> Vec<BigInt>;
 
-// match what rapidsnark expects
-#[derive(Debug, Serialize, Deserialize)]
-#[allow(non_snake_case)]
-struct VerificationKey {
-    protocol: String,
-    curve: String,
-    nPublic: u32,
-    vk_alpha_1: [String; 3],
-    vk_beta_2: [[String; 2]; 3],
-    vk_gamma_2: [[String; 2]; 3],
-    vk_delta_2: [[String; 2]; 3],
-    IC: Vec<[String; 3]>,
-}
-
-#[repr(C)]
 pub struct ProofResult {
-    proof: *mut c_char,
-    public_signals: *mut c_char,
+    proof: String,
+    public_signals: String,
 }
 
 extern "C" {
-    fn groth16_api_prove(
-        zkeyFilename: *const c_char,
-        wtnsData: *mut u8,
-        wtnsDataLen: c_uint,
-    ) -> *mut ProofResult;
-    fn groth16_api_verify(proof: *mut ProofResult, key_json: *const c_char) -> bool;
-    fn free_proof_result(result: *mut ProofResult);
+    fn groth16_prover_zkey_file(
+        zkey_file_path: *const std::os::raw::c_char,
+        wtns_buffer: *const std::os::raw::c_void,
+        wtns_size: std::ffi::c_ulong,
+        proof_buffer: *mut std::os::raw::c_char,
+        proof_size: *mut std::ffi::c_ulong,
+        public_buffer: *mut std::os::raw::c_char,
+        public_size: *mut std::ffi::c_ulong,
+        error_msg: *mut std::os::raw::c_char,
+        error_msg_maxsize: std::ffi::c_ulong,
+    ) -> i32;
 }
 
-pub fn verify_proof(zkey_path: &str, proof: String) -> Result<bool> {
-    let mut header_reader = ZkeyHeaderReader::new(zkey_path);
-    header_reader.read();
-    let file = File::open(zkey_path)?;
-    let mut reader = std::io::BufReader::new(file);
-    let proving_key = read_proving_key::<_, Bn254>(&mut reader)?;
-    // convert out proving key to json so we can
-    // use it with rapidsnark
-    let vk = proving_key.vk;
-    // let v = proving_key.vk.alpha_g1.to_string();
-    let vkey = VerificationKey {
-        protocol: "groth16".to_string(),
-        curve: "bn128".to_string(),
-        nPublic: 0, // this is unused in the rapidsnark verifier
-        vk_alpha_1: [
-            vk.alpha_g1.x.to_string(),
-            vk.alpha_g1.y.to_string(),
-            "1".to_string(),
-        ],
-        vk_beta_2: [
-            [vk.beta_g2.x.c0.to_string(), vk.beta_g2.x.c1.to_string()],
-            [vk.beta_g2.y.c0.to_string(), vk.beta_g2.y.c1.to_string()],
-            ["1".to_string(), "0".to_string()],
-        ],
-        vk_gamma_2: [
-            [vk.gamma_g2.x.c0.to_string(), vk.gamma_g2.x.c1.to_string()],
-            [vk.gamma_g2.y.c0.to_string(), vk.gamma_g2.y.c1.to_string()],
-            ["1".to_string(), "0".to_string()],
-        ],
-        vk_delta_2: [
-            [vk.delta_g2.x.c0.to_string(), vk.delta_g2.x.c1.to_string()],
-            [vk.delta_g2.y.c0.to_string(), vk.delta_g2.y.c1.to_string()],
-            ["1".to_string(), "0".to_string()],
-        ],
-        IC: vk
-            .gamma_abc_g1
-            .iter()
-            .map(|p| [p.x.to_string(), p.y.to_string(), "1".to_string()])
-            .collect(),
-    };
-    let vkey_json = serde_json::to_string(&vkey)?;
-    let vkey_json_cstr = CString::new(vkey_json)?;
-    let v: serde_json::Value = serde_json::from_str(&proof)?;
-    let proof = v["proof"].to_string();
-    let signals = v["signals"].to_string();
-    unsafe {
-        let result = groth16_api_verify(
-            &mut ProofResult {
-                proof: CString::new(proof).unwrap().into_raw(),
-                public_signals: CString::new(signals).unwrap().into_raw(),
-            },
-            vkey_json_cstr.as_ptr(),
-        );
-        Ok(result)
+use num_traits::ops::bytes::ToBytes;
+use std::io::{self};
+
+/// Parse bigints to `wtns` format.<br/>
+/// Reference: [witnesscalc/src/witnesscalc.cpp](https://github.com/0xPolygonID/witnesscalc/blob/4a789880727aa0df50f1c4ef78ec295f5a30a15e/src/witnesscalc.cpp)
+pub fn parse_bigints_to_witness(bigints: Vec<BigInt>) -> io::Result<Vec<u8>> {
+    let mut buffer = Vec::new();
+    let version: u32 = 2;
+    let n_sections: u32 = 2;
+    let n8: u32 = 32;
+    let q = BigInt::from_str(
+        "21888242871839275222246405745257275088548364400416034343698204186575808495617",
+    )
+    .unwrap();
+    let n_witness_values: u32 = bigints.len() as u32;
+
+    // Write the format bytes (4 bytes)
+    let wtns_format = "wtns".as_bytes();
+    buffer.extend_from_slice(wtns_format);
+
+    // Write version (4 bytes)
+    buffer.extend_from_slice(&version.to_le_bytes());
+
+    // Write number of sections (4 bytes)
+    buffer.extend_from_slice(&n_sections.to_le_bytes());
+
+    // Iterate through sections to write the data
+    // Section 1 (Field parameters)
+    let section_id_1: u32 = 1;
+    let section_length_1: u64 = 8 + n8 as u64;
+    buffer.extend_from_slice(&section_id_1.to_le_bytes());
+    buffer.extend_from_slice(&section_length_1.to_le_bytes());
+
+    // Write n8 (4 bytes), q (32 bytes), and n_witness_values (4 bytes)
+    buffer.extend_from_slice(&n8.to_le_bytes());
+    buffer.extend_from_slice(&q.to_signed_bytes_le());
+    buffer.extend_from_slice(&n_witness_values.to_le_bytes());
+
+    // Section 2 (Witness data)
+    let section_id_2: u32 = 2;
+    let section_length_2: u64 = bigints.len() as u64 * n8 as u64; // Witness data size
+    buffer.extend_from_slice(&section_id_2.to_le_bytes());
+    buffer.extend_from_slice(&section_length_2.to_le_bytes());
+
+    // Write the witness data (each BigInt to n8 bytes)
+    for bigint in bigints {
+        let mut bytes = bigint.to_le_bytes();
+        bytes.resize(n8 as usize, 0); // Ensure each BigInt is padded to n8 bytes
+        buffer.extend_from_slice(&bytes);
     }
+
+    // Return the buffer containing the complete witness data
+    Ok(buffer)
 }
 
-pub fn generate_proof(
+/// Wrapper for `groth16_prover_zkey_file`
+pub fn groth16_prover_zkey_file_wrapper(
     zkey_path: &str,
-    inputs: std::collections::HashMap<String, Vec<String>>,
-    witness_fn: WtnsFn,
-) -> Result<String> {
-    // Form the inputs
-    let bigint_inputs = inputs
-        .into_iter()
-        .map(|(k, v)| {
-            (
-                k,
-                v.into_iter()
-                    .map(|i| BigInt::from_str(&i).unwrap())
-                    .collect(),
-            )
-        })
-        .collect();
+    wtns_buffer: Vec<u8>,
+) -> Result<ProofResult> {
+    let wtns_size = wtns_buffer.len() as u64;
 
-    let mut wtns = witness_fn(bigint_inputs)
-        .into_iter()
-        .map(|w| w.to_biguint().unwrap())
-        .flat_map(|v| {
-            let mut bytes = v.to_bytes_le();
-            bytes.resize(32, 0);
-            bytes
-        })
-        .collect::<Vec<_>>();
+    let mut proof_buffer = vec![0u8; 4 * 1024 * 1024]; // Adjust size as needed
+    let mut proof_size: u64 = 4 * 1024 * 1024;
+    let proof_ptr = proof_buffer.as_mut_ptr() as *mut std::ffi::c_char;
 
-    // Convert Rust strings to C strings
-    let zkey_cstr = CString::new(zkey_path).context("Failed to create CString for zkey path")?;
+    let mut public_buffer = vec![0u8; 4 * 1024 * 1024]; // Adjust size as needed
+    let mut public_size: u64 = 4 * 1024 * 1024;
+    let public_ptr = public_buffer.as_mut_ptr() as *mut std::ffi::c_char;
+
+    let mut error_msg = vec![0u8; 256]; // Error message buffer
+    let error_msg_ptr = error_msg.as_mut_ptr() as *mut std::ffi::c_char;
 
     unsafe {
-        let proof_ptr =
-            groth16_api_prove(zkey_cstr.as_ptr(), wtns.as_mut_ptr(), wtns.len() as c_uint);
-
-        if proof_ptr.is_null() {
-            return Err(anyhow::anyhow!("Proof generation failed"));
+        let result = groth16_prover_zkey_file(
+            zkey_path.as_ptr() as *const std::ffi::c_char,
+            wtns_buffer.as_ptr() as *const std::os::raw::c_void, // Witness buffer
+            wtns_size,
+            proof_ptr,
+            &mut proof_size,
+            public_ptr,
+            &mut public_size,
+            error_msg_ptr,
+            error_msg.len() as u64,
+        );
+        if result != 0 {
+            let error_string = std::ffi::CStr::from_ptr(error_msg_ptr)
+                .to_string_lossy()
+                .into_owned();
+            return Err(anyhow::anyhow!("Proof generation failed: {}", error_string));
         }
-
         // Convert both strings
-        let result = &*proof_ptr;
-        let proof = std::ffi::CStr::from_ptr(result.proof)
+        let proof = std::ffi::CStr::from_ptr(proof_ptr)
             .to_string_lossy()
             .into_owned();
-        let public_signals = std::ffi::CStr::from_ptr(result.public_signals)
+        let public_signals = std::ffi::CStr::from_ptr(public_ptr)
             .to_string_lossy()
             .into_owned();
-        free_proof_result(proof_ptr);
-        Ok(format!(
-            "{{ \"proof\": {proof},\"signals\": {public_signals}}}"
-        ))
+        Ok(ProofResult {
+            proof,
+            public_signals,
+        })
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use anyhow::bail;
     use anyhow::Result;
     use num_bigint::BigInt;
-    use std::collections::HashMap;
-    use std::str::FromStr;
+    use std::{collections::HashMap, str::FromStr};
+
+    use crate::{parse_bigints_to_witness, WtnsFn};
 
     rust_witness::witness!(multiplier2);
     rust_witness::witness!(keccak256256test);
@@ -190,6 +162,28 @@ mod tests {
         inputs
     }
 
+    fn compute_witness(
+        inputs: HashMap<String, Vec<String>>,
+        witness_fn: WtnsFn,
+    ) -> Result<Vec<u8>> {
+        // Form the inputs
+        let bigint_inputs = inputs
+            .into_iter()
+            .map(|(k, v)| {
+                (
+                    k,
+                    v.into_iter()
+                        .map(|i| BigInt::from_str(&i).unwrap())
+                        .collect(),
+                )
+            })
+            .collect();
+
+        let wtns: Vec<BigInt> = witness_fn(bigint_inputs);
+        let witnesscalc_wtns = parse_bigints_to_witness(wtns)?;
+        Ok(witnesscalc_wtns)
+    }
+
     #[test]
     fn test_prove_rapidsnark() -> Result<()> {
         // Create a new MoproCircom instance
@@ -205,11 +199,14 @@ mod tests {
         inputs.insert("a".to_string(), vec![a.to_string()]);
         inputs.insert("b".to_string(), vec![b.to_string()]);
 
-        let proof_json = super::generate_proof(&zkey_path, inputs, multiplier2_witness)?;
-        let valid = super::verify_proof(&zkey_path, proof_json)?;
-        if !valid {
-            bail!("Proof is invalid");
-        }
+        let wtns_buffer = compute_witness(inputs, multiplier2_witness)?;
+        let proof_result = super::groth16_prover_zkey_file_wrapper(&zkey_path, wtns_buffer)?;
+        println!("{}", proof_result.proof);
+        println!("{}", proof_result.public_signals);
+        // let valid = super::verify_proof(&zkey_path, proof_json)?;
+        // if !valid {
+        //     bail!("Proof is invalid");
+        // }
         Ok(())
     }
 
@@ -225,12 +222,19 @@ mod tests {
 
         let inputs = bytes_to_circuit_inputs(&input_vec);
 
+        // Generate Witness Buffer
+        let wtns_buffer = compute_witness(inputs, keccak256256test_witness)?;
+        let wtns_data = std::fs::read("./test-vectors/keccak256_256_test.wtns")?;
+        assert_eq!(wtns_buffer, wtns_data);
+
         // Generate Proof
-        let proof_json = super::generate_proof(&zkey_path, inputs, keccak256256test_witness)?;
-        let valid = super::verify_proof(&zkey_path, proof_json)?;
-        if !valid {
-            bail!("Proof is invalid");
-        }
+        let proof_result = super::groth16_prover_zkey_file_wrapper(&zkey_path, wtns_buffer)?;
+        println!("{}", proof_result.proof);
+        println!("{}", proof_result.public_signals);
+        // let valid = super::verify_proof(&zkey_path, proof_json)?;
+        // if !valid {
+        //     bail!("Proof is invalid");
+        // }
         Ok(())
     }
 }