Use initialize function

Cargo.toml

@@ -17,6 +17,7 @@ license-file = "mit-license.md"
 default = []
 bench = [ "criterion", "proptest" ]
 mimc = [ "zkp-u256" ]
+dylib = [ "wasmer/dylib" ]
 
 [[bench]]
 name = "criterion"
@@ -48,7 +49,7 @@ serde = "1.0"
 sha2 = "0.10.1"
 thiserror = "1.0.0"
 tiny-keccak = { version = "2.0.2" }
-wasmer = { version = "2.0", features = [ "dylib" ] }
+wasmer = { version = "2.0" }
 zkp-u256 = { version = "0.2", optional = true } # TODO: Remove
 
 # Use the same `ethers-core` version as ark-circom

src/circuit.rs

@@ -2,38 +2,68 @@ use ark_bn254::{Bn254, Fr};
 use ark_circom::{read_zkey, WitnessCalculator};
 use ark_groth16::ProvingKey;
 use ark_relations::r1cs::ConstraintMatrices;
-use color_eyre::Result;
 use core::include_bytes;
 use once_cell::sync::{Lazy, OnceCell};
 use std::{io::Cursor, sync::Mutex};
-use wasmer::{Dylib, Module, Store};
+use wasmer::{Module, Store};
+
+#[cfg(feature = "dylib")]
+use std::path::Path;
+#[cfg(feature = "dylib")]
+use wasmer::{Dylib};
 
 const ZKEY_BYTES: &[u8] = include_bytes!("../semaphore/build/snark/semaphore_final.zkey");
 
+#[cfg(not(feature = "dylib"))]
 const WASM: &[u8] = include_bytes!("../semaphore/build/snark/semaphore.wasm");
 
-pub static ZKEY: Lazy<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> = Lazy::new(|| {
+static ZKEY: Lazy<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> = Lazy::new(|| {
     let mut reader = Cursor::new(ZKEY_BYTES);
     read_zkey(&mut reader).expect("zkey should be valid")
 });
 
-pub static WITNESS_CALCULATOR_DYLIB: OnceCell<String> = OnceCell::new();
+static WITNESS_CALCULATOR: OnceCell<Mutex<WitnessCalculator>> = OnceCell::new();
 
-pub static WITNESS_CALCULATOR: Lazy<Mutex<WitnessCalculator>> = Lazy::new(|| {
-    // Create Wasm module
-    let module = if let Some(path) = WITNESS_CALCULATOR_DYLIB.get() {
+/// Initialize the library.
+#[cfg(feature = "dylib")]
+pub fn initialize(dylib_path: &Path) {
     let store = Store::new(&Dylib::headless().engine());
     // The module must be exported using [`Module::serialize`].
-        unsafe {
-            Module::deserialize_from_file(&store, path).expect("Failed to load wasm dylib module")
-        }
-    } else {
-        let store = Store::default();
-        Module::from_binary(&store, WASM).expect("wasm should be valid")
+    let module = unsafe {
+        Module::deserialize_from_file(&store, dylib_path).expect("Failed to load wasm dylib module")
     };
+    let result =
+        WitnessCalculator::from_module(module).expect("Failed to create witness calculator");
+    WITNESS_CALCULATOR
+        .set(Mutex::new(result))
+        .expect("Failed to initialize witness calculator");
 
-    // Create witness calculator
+    // Force init of ZKEY
+    Lazy::force(&ZKEY);
+}
+
+#[must_use]
+pub fn zkey() -> &'static (ProvingKey<Bn254>, ConstraintMatrices<Fr>) {
+    &*ZKEY
+}
+
+#[cfg(feature = "dylib")]
+#[must_use]
+pub fn witness_calculator() -> &'static Mutex<WitnessCalculator> {
+    WITNESS_CALCULATOR.get().expect(
+        "Semaphore-rs not initialized. The library needs to be initialized before use when build \
+         with the `cdylib` feature.",
+    )
+}
+
+#[cfg(not(feature = "dylib"))]
+#[must_use]
+pub fn witness_calculator() -> &'static Mutex<WitnessCalculator> {
+    WITNESS_CALCULATOR.get_or_init(|| {
+        let store = Store::default();
+        let module = Module::from_binary(&store, WASM).expect("wasm should be valid");
         let result =
             WitnessCalculator::from_module(module).expect("Failed to create witness calculator");
         Mutex::new(result)
-});
+    })
+}

src/lib.rs

@@ -27,6 +27,9 @@ pub use crate::{
     poseidon_hash::poseidon_hash,
 };
 
+#[cfg(feature = "dylib")]
+pub use circuit::initialize;
+
 pub type Groth16Proof = ark_groth16::Proof<Bn<Parameters>>;
 pub type EthereumGroth16Proof = ark_circom::ethereum::Proof;
 

src/protocol.rs

@@ -1,5 +1,5 @@
 use crate::{
-    circuit::{WITNESS_CALCULATOR, ZKEY},
+    circuit::{witness_calculator, zkey},
     identity::Identity,
     merkle_tree::{self, Branch},
     poseidon_hash,
@@ -109,6 +109,11 @@ pub fn generate_proof(
     )
 }
 
+/// Generates a semaphore proof from entropy
+///
+/// # Errors
+///
+/// Returns a [`ProofError`] if proving fails.
 pub fn generate_proof_rng(
     identity: &Identity,
     merkle_proof: &merkle_tree::Proof<PoseidonHash>,
@@ -151,7 +156,7 @@ fn generate_proof_rs(
 
     let now = Instant::now();
 
-    let full_assignment = WITNESS_CALCULATOR
+    let full_assignment = witness_calculator()
         .lock()
         .expect("witness_calculator mutex should not get poisoned")
         .calculate_witness_element::<Bn254, _>(inputs, false)
@@ -160,13 +165,14 @@ fn generate_proof_rs(
     println!("witness generation took: {:.2?}", now.elapsed());
 
     let now = Instant::now();
+    let zkey = zkey();
     let ark_proof = create_proof_with_reduction_and_matrices::<_, CircomReduction>(
-        &ZKEY.0,
+        &zkey.0,
         r,
         s,
-        &ZKEY.1,
-        ZKEY.1.num_instance_variables,
-        ZKEY.1.num_constraints,
+        &zkey.1,
+        zkey.1.num_instance_variables,
+        zkey.1.num_constraints,
         full_assignment.as_slice(),
     )?;
     let proof = ark_proof.into();
@@ -188,7 +194,8 @@ pub fn verify_proof(
     external_nullifier_hash: Field,
     proof: &Proof,
 ) -> Result<bool, ProofError> {
-    let pvk = prepare_verifying_key(&ZKEY.0.vk);
+    let zkey = zkey();
+    let pvk = prepare_verifying_key(&zkey.0.vk);
 
     let public_inputs = [
         root.into(),