chore(rln): further refactoring of interface (#261)

vacp2p · Jun 18, 2024 · 5540ddc · 5540ddc
1 parent d8f813b
commit 5540ddc
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 64 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rln/Cargo.toml b/rln/Cargo.toml
@@ -48,6 +48,7 @@ num-bigint = { version = "=0.4.3", default-features = false, features = [
 ] }
 num-traits = "=0.2.15"
 once_cell = "=1.17.1"
+lazy_static = "=1.4.0"
 rand = "=0.8.5"
 rand_chacha = "=0.3.1"
 tiny-keccak = { version = "=2.0.2", features = ["keccak"] }

diff --git a/rln/src/circuit.rs b/rln/src/circuit.rs
@@ -10,50 +10,49 @@ use ark_serialize::CanonicalDeserialize;
 use cfg_if::cfg_if;
 use color_eyre::{Report, Result};
 
-cfg_if! {
-    if #[cfg(not(target_arch = "wasm32"))] {
-        use ark_circom::{WitnessCalculator};
-        use once_cell::sync::{Lazy};
-        use std::sync::Mutex;
-        use wasmer::{Module, Store};
-        use std::sync::Arc;
-    }
-}
+#[cfg(not(target_arch = "wasm32"))]
+use {
+    ark_circom::WitnessCalculator,
+    lazy_static::lazy_static,
+    std::sync::{Arc, Mutex},
+    wasmer::{Module, Store},
+};
 
-cfg_if! {
-    if #[cfg(feature = "arkzkey")] {
-        use ark_zkey::read_arkzkey_from_bytes;
-        const ARKZKEY_BYTES: &[u8] = include_bytes!("tree_height_20/rln_final.arkzkey");
-    } else {
-        use std::io::Cursor;
-        use ark_circom::read_zkey;
-    }
-}
+#[cfg(feature = "arkzkey")]
+use ark_zkey::read_arkzkey_from_bytes;
+
+#[cfg(not(feature = "arkzkey"))]
+use {ark_circom::read_zkey, std::io::Cursor};
+
+#[cfg(feature = "arkzkey")]
+const ARKZKEY_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/rln_final.arkzkey");
 
 pub const ZKEY_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/rln_final.zkey");
 pub const VK_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/verification_key.arkvkey");
 const WASM_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/rln.wasm");
 
 #[cfg(not(target_arch = "wasm32"))]
-static ZKEY: Lazy<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> = Lazy::new(|| {
-    cfg_if! {
-        if #[cfg(feature = "arkzkey")] {
-            read_arkzkey_from_bytes(ARKZKEY_BYTES).expect("Failed to read arkzkey")
-        } else {
-            let mut reader = Cursor::new(ZKEY_BYTES);
-            read_zkey(&mut reader).expect("Failed to read zkey")
+lazy_static! {
+    #[cfg(not(target_arch = "wasm32"))]
+    static ref ZKEY: (ProvingKey<Curve>, ConstraintMatrices<Fr>) = {
+        cfg_if! {
+                if #[cfg(feature = "arkzkey")] {
+                    read_arkzkey_from_bytes(ARKZKEY_BYTES).expect("Failed to read arkzkey")
+                } else {
+                    let mut reader = Cursor::new(ZKEY_BYTES);
+                    read_zkey(&mut reader).expect("Failed to read zkey")
+                }
         }
-    }
-});
+    };
 
-#[cfg(not(target_arch = "wasm32"))]
-static VK: Lazy<VerifyingKey<Curve>> =
-    Lazy::new(|| vk_from_ark_serialized(VK_BYTES).expect("Failed to read vk"));
+    #[cfg(not(target_arch = "wasm32"))]
+    static ref VK: VerifyingKey<Curve> = vk_from_ark_serialized(VK_BYTES).expect("Failed to read vk");
 
-#[cfg(not(target_arch = "wasm32"))]
-static WITNESS_CALCULATOR: Lazy<Arc<Mutex<WitnessCalculator>>> = Lazy::new(|| {
-    circom_from_raw(WASM_BYTES.to_vec()).expect("Failed to create witness calculator")
-});
+    #[cfg(not(target_arch = "wasm32"))]
+    static ref WITNESS_CALCULATOR: Arc<Mutex<WitnessCalculator>> = {
+        circom_from_raw(WASM_BYTES).expect("Failed to create witness calculator")
+    };
+}
 
 pub const TEST_TREE_HEIGHT: usize = 20;
 
@@ -69,21 +68,22 @@ pub type G2Affine = ArkG2Affine;
 pub type G2Projective = ArkG2Projective;
 
 // Loads the proving key using a bytes vector
-pub fn zkey_from_raw(zkey_data: &Vec<u8>) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> {
-    if !zkey_data.is_empty() {
-        let proving_key_and_matrices = match () {
-            #[cfg(feature = "arkzkey")]
-            () => read_arkzkey_from_bytes(zkey_data.as_slice())?,
-            #[cfg(not(feature = "arkzkey"))]
-            () => {
-                let mut c = Cursor::new(zkey_data);
-                read_zkey(&mut c)?
-            }
-        };
-        Ok(proving_key_and_matrices)
-    } else {
-        Err(Report::msg("No proving key found!"))
+pub fn zkey_from_raw(zkey_data: &[u8]) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> {
+    if zkey_data.is_empty() {
+        return Err(Report::msg("No proving key found!"));
     }
+
+    let proving_key_and_matrices = match () {
+        #[cfg(feature = "arkzkey")]
+        () => read_arkzkey_from_bytes(zkey_data)?,
+        #[cfg(not(feature = "arkzkey"))]
+        () => {
+            let mut reader = Cursor::new(zkey_data);
+            read_zkey(&mut reader)?
+        }
+    };
+
+    Ok(proving_key_and_matrices)
 }
 
 // Loads the proving key
@@ -93,19 +93,17 @@ pub fn zkey_from_folder() -> &'static (ProvingKey<Curve>, ConstraintMatrices<Fr>
 }
 
 // Loads the verification key from a bytes vector
-pub fn vk_from_raw(vk_data: &[u8], zkey_data: &Vec<u8>) -> Result<VerifyingKey<Curve>> {
-    let verifying_key: VerifyingKey<Curve>;
-
+pub fn vk_from_raw(vk_data: &[u8], zkey_data: &[u8]) -> Result<VerifyingKey<Curve>> {
     if !vk_data.is_empty() {
-        verifying_key = vk_from_ark_serialized(vk_data)?;
-        Ok(verifying_key)
-    } else if !zkey_data.is_empty() {
+        return vk_from_ark_serialized(vk_data);
+    }
+
+    if !zkey_data.is_empty() {
         let (proving_key, _matrices) = zkey_from_raw(zkey_data)?;
-        verifying_key = proving_key.vk;
-        Ok(verifying_key)
-    } else {
-        Err(Report::msg("No proving/verification key found!"))
+        return Ok(proving_key.vk);
     }
+
+    Err(Report::msg("No proving/verification key found!"))
 }
 
 // Loads the verification key
@@ -116,12 +114,10 @@ pub fn vk_from_folder() -> &'static VerifyingKey<Curve> {
 
 // Initializes the witness calculator using a bytes vector
 #[cfg(not(target_arch = "wasm32"))]
-pub fn circom_from_raw(wasm_buffer: Vec<u8>) -> Result<Arc<Mutex<WitnessCalculator>>> {
-    let store = Store::default();
-    let module = Module::new(&store, wasm_buffer)?;
+pub fn circom_from_raw(wasm_buffer: &[u8]) -> Result<Arc<Mutex<WitnessCalculator>>> {
+    let module = Module::new(&Store::default(), wasm_buffer)?;
     let result = WitnessCalculator::from_module(module)?;
-    let wrapped = Mutex::new(result);
-    Ok(Arc::new(wrapped))
+    Ok(Arc::new(Mutex::new(result)))
 }
 
 // Initializes the witness calculator

diff --git a/rln/src/public.rs b/rln/src/public.rs
@@ -153,7 +153,7 @@ impl RLN {
         mut tree_config_input: R,
     ) -> Result<RLN> {
         #[cfg(not(target_arch = "wasm32"))]
-        let witness_calculator = circom_from_raw(circom_vec)?;
+        let witness_calculator = circom_from_raw(&circom_vec)?;
 
         let proving_key = zkey_from_raw(&zkey_vec)?;
         let verification_key = vk_from_raw(&vk_vec, &zkey_vec)?;