From 5540ddc993adf831ad98a34ff836f6263e9a1f3b Mon Sep 17 00:00:00 2001 From: Aaryamann Challani <43716372+rymnc@users.noreply.github.com> Date: Tue, 18 Jun 2024 11:56:23 +0530 Subject: [PATCH] chore(rln): further refactoring of interface (#261) --- Cargo.lock | 1 + rln/Cargo.toml | 1 + rln/src/circuit.rs | 124 ++++++++++++++++++++++----------------------- rln/src/public.rs | 2 +- 4 files changed, 63 insertions(+), 65 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8c8ffbc..236fcf3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2284,6 +2284,7 @@ dependencies = [ "cfg-if", "color-eyre", "criterion 0.4.0", + "lazy_static 1.4.0", "num-bigint", "num-traits", "once_cell", diff --git a/rln/Cargo.toml b/rln/Cargo.toml index e676725..a268978 100644 --- 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 index 34f1d38..2cb7920 100644 --- 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, ConstraintMatrices)> = 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, ConstraintMatrices) = { + 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> = - Lazy::new(|| vk_from_ark_serialized(VK_BYTES).expect("Failed to read vk")); + #[cfg(not(target_arch = "wasm32"))] + static ref VK: VerifyingKey = vk_from_ark_serialized(VK_BYTES).expect("Failed to read vk"); -#[cfg(not(target_arch = "wasm32"))] -static WITNESS_CALCULATOR: Lazy>> = 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> = { + 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) -> Result<(ProvingKey, ConstraintMatrices)> { - 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, ConstraintMatrices)> { + 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, ConstraintMatrices } // Loads the verification key from a bytes vector -pub fn vk_from_raw(vk_data: &[u8], zkey_data: &Vec) -> Result> { - let verifying_key: VerifyingKey; - +pub fn vk_from_raw(vk_data: &[u8], zkey_data: &[u8]) -> Result> { if !vk_data.is_empty() { - verifying_key = vk_from_ark_serialized(vk_data)?; - Ok(verifying_key) - } else 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 vk_from_ark_serialized(vk_data); } + + if !zkey_data.is_empty() { + let (proving_key, _matrices) = zkey_from_raw(zkey_data)?; + 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 { // Initializes the witness calculator using a bytes vector #[cfg(not(target_arch = "wasm32"))] -pub fn circom_from_raw(wasm_buffer: Vec) -> Result>> { - let store = Store::default(); - let module = Module::new(&store, wasm_buffer)?; +pub fn circom_from_raw(wasm_buffer: &[u8]) -> Result>> { + 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 index 6d2dfb9..0ab5416 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -153,7 +153,7 @@ impl RLN { mut tree_config_input: R, ) -> Result { #[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)?;