Merge pull request #15 from worldcoin/philsippl/init_path

Configurable path to dylib

Cargo.toml

@@ -17,6 +17,7 @@ license-file = "mit-license.md"
 default = []
 bench = [ "criterion", "proptest" ]
 mimc = [ "zkp-u256" ]
+dylib = [ "wasmer/dylib", "wasmer-engine-dylib", "wasmer-compiler-cranelift" ]
 
 [[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
@@ -67,9 +68,9 @@ tracing-test = "0.2"
 [build-dependencies]
 color-eyre = "0.5"
 enumset = "1.0.8"
-wasmer = { version = "2.0", features = [ "dylib" ] }
-wasmer-engine-dylib = "2.2.1"
-wasmer-compiler-cranelift = "2.2.1"
+wasmer = { version = "2.0" }
+wasmer-engine-dylib = { version = "2.2.1", optional = true }
+wasmer-compiler-cranelift = { version = "2.2.1", optional = true }
 
 [profile.release]
 codegen-units = 1

build.rs

@@ -2,12 +2,7 @@ use color_eyre::eyre::{eyre, Result};
 use std::{
     path::{Component, Path, PathBuf},
     process::Command,
-    str::FromStr, env,
 };
-use wasmer::{Module, Store, Target, Triple};
-use wasmer_compiler_cranelift::Cranelift;
-use wasmer_engine_dylib::Dylib;
-use enumset::enum_set;
 
 const ZKEY_FILE: &str = "./semaphore/build/snark/semaphore_final.zkey";
 const WASM_FILE: &str = "./semaphore/build/snark/semaphore.wasm";
@@ -26,7 +21,7 @@ fn absolute(path: &str) -> Result<PathBuf> {
             Component::ParentDir => {
                 absolute.pop();
             }
-            component @ _ => absolute.push(component.as_os_str()),
+            component => absolute.push(component.as_os_str()),
         }
     }
     Ok(absolute)
@@ -67,14 +62,24 @@ fn build_circuit() -> Result<()> {
     Ok(())
 }
 
+#[cfg(feature = "dylib")]
 fn build_dylib() -> Result<()> {
+    use enumset::enum_set;
+    use std::{env, str::FromStr};
+    use wasmer::{Module, Store, Target, Triple};
+    use wasmer_compiler_cranelift::Cranelift;
+    use wasmer_engine_dylib::Dylib;
+
     let wasm_file = absolute(WASM_FILE)?;
     assert!(wasm_file.exists());
 
     let out_dir = env::var("OUT_DIR")?;
     let out_dir = Path::new(&out_dir).to_path_buf();
     let dylib_file = out_dir.join("semaphore.dylib");
-    println!("cargo:rustc-env=BUILD_DYLIB_FILE={}", dylib_file.display());
+    println!(
+        "cargo:rustc-env=CIRCUIT_WASM_DYLIB={}",
+        dylib_file.display()
+    );
 
     if dylib_file.exists() {
         return Ok(());
@@ -85,9 +90,7 @@ fn build_dylib() -> Result<()> {
     let cpu_features = enum_set!();
     let target = Target::new(triple, cpu_features);
     let compiler_config = Cranelift::default();
-    let engine = Dylib::new(compiler_config)
-        .target(target)
-        .engine();
+    let engine = Dylib::new(compiler_config).target(target).engine();
 
     // Compile the WASM module
     let store = Store::new(&engine);
@@ -101,6 +104,7 @@ fn build_dylib() -> Result<()> {
 
 fn main() -> Result<()> {
     build_circuit()?;
+    #[cfg(feature = "dylib")]
     build_dylib()?;
     Ok(())
 }

src/circuit.rs

@@ -3,33 +3,76 @@ use ark_circom::{read_zkey, WitnessCalculator};
 use ark_groth16::ProvingKey;
 use ark_relations::r1cs::ConstraintMatrices;
 use core::include_bytes;
-use once_cell::sync::Lazy;
+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::{env, path::Path};
+#[cfg(feature = "dylib")]
+use wasmer::Dylib;
 
 const ZKEY_BYTES: &[u8] = include_bytes!(env!("BUILD_RS_ZKEY_FILE"));
+
+#[cfg(not(feature = "dylib"))]
 const WASM: &[u8] = include_bytes!(env!("BUILD_RS_WASM_FILE"));
 
-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: Lazy<Mutex<WitnessCalculator>> = Lazy::new(|| {
-    // Create Wasm module
-    let module = if let Some(path) = option_env!("CIRCUIT_WASM_DYLIB") {
-        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")
-    };
+static WITNESS_CALCULATOR: OnceCell<Mutex<WitnessCalculator>> = OnceCell::new();
 
-    // Create witness calculator
+/// Initialize the library.
+#[cfg(feature = "dylib")]
+pub fn initialize(dylib_path: &Path) {
+    WITNESS_CALCULATOR
+        .set(from_dylib(dylib_path))
+        .expect("Failed to initialize witness calculator");
+
+    // Force init of ZKEY
+    Lazy::force(&ZKEY);
+}
+
+#[cfg(feature = "dylib")]
+fn from_dylib(path: &Path) -> Mutex<WitnessCalculator> {
+    let store = Store::new(&Dylib::headless().engine());
+    // The module must be exported using [`Module::serialize`].
+    let module = unsafe {
+        Module::deserialize_from_file(&store, path).expect("Failed to load wasm dylib module")
+    };
     let result =
         WitnessCalculator::from_module(module).expect("Failed to create witness calculator");
     Mutex::new(result)
-});
+}
+
+#[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_or_init(|| {
+        let path = env::var("CIRCUIT_WASM_DYLIB").expect(
+            "Semaphore-rs is not initialized. The library needs to be initialized before use when \
+             build with the `cdylib` feature. You can initialize by calling `initialize` or \
+             seting the `CIRCUIT_WASM_DYLIB` environment variable.",
+        );
+        from_dylib(Path::new(&path))
+    })
+}
+
+#[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(),