Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: start optimizing wasm blob #226

Draft
wants to merge 4 commits into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

24 changes: 11 additions & 13 deletions rln-wasm/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use std::vec::Vec;

use js_sys::{BigInt as JsBigInt, Object, Uint8Array};
use num_bigint::BigInt;
use rln::circuit::{default_vk, default_zkey};
use rln::public::{hash, poseidon_hash, RLN};
use wasm_bindgen::prelude::*;

Expand Down Expand Up @@ -181,13 +182,10 @@ impl<'a> ProcessArg for &'a [u8] {

#[allow(clippy::not_unsafe_ptr_arg_deref)]
#[wasm_bindgen(js_name = newRLN)]
pub fn wasm_new(
tree_height: usize,
zkey: Uint8Array,
vk: Uint8Array,
) -> Result<*mut RLNWrapper, String> {
let instance = RLN::new_with_params(tree_height, zkey.to_vec(), vk.to_vec())
.map_err(|err| format!("{:#?}", err))?;
pub fn wasm_new() -> Result<*mut RLNWrapper, String> {
let zkey = default_zkey().map_err(|err| format!("{:#?}", err))?;
let vk = default_vk().map_err(|err| format!("{:#?}", err))?;
let instance = RLN::new_with_params(zkey, vk).map_err(|err| format!("{:#?}", err))?;
let wrapper = RLNWrapper { instance };
Ok(Box::into_raw(Box::new(wrapper)))
}
Expand Down Expand Up @@ -276,6 +274,12 @@ pub fn rln_witness_to_json(
Object::from_entries(&js_value).map_err(|err| format!("{:#?}", err))
}

#[allow(clippy::not_unsafe_ptr_arg_deref)]
#[wasm_bindgen(js_name = generateMembershipKey)]
pub fn wasm_key_gen(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
call_with_output_and_error_msg!(ctx, key_gen, "could not generate membership keys")
}

#[allow(clippy::not_unsafe_ptr_arg_deref)]
#[wasm_bindgen]
pub fn generate_rln_proof_with_witness(
Expand Down Expand Up @@ -305,12 +309,6 @@ pub fn generate_rln_proof_with_witness(
)
}

#[allow(clippy::not_unsafe_ptr_arg_deref)]
#[wasm_bindgen(js_name = generateMembershipKey)]
pub fn wasm_key_gen(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
call_with_output_and_error_msg!(ctx, key_gen, "could not generate membership keys")
}

#[allow(clippy::not_unsafe_ptr_arg_deref)]
#[wasm_bindgen(js_name = generateExtendedMembershipKey)]
pub fn wasm_extended_key_gen(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
Expand Down
23 changes: 4 additions & 19 deletions rln-wasm/tests/rln-wasm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
#[cfg(test)]
mod tests {
use js_sys::{BigInt as JsBigInt, Object, Uint8Array};
use rln::circuit::{Fr, TEST_TREE_HEIGHT};
use rln::circuit::Fr;
use rln::hashers::{hash_to_field, poseidon_hash};
use rln::utils::{bytes_le_to_fr, fr_to_bytes_le, normalize_usize};
use rln_wasm::*;
Expand All @@ -21,16 +21,9 @@ mod tests {

#[wasm_bindgen_test]
pub async fn test_basic_flow() {
let tree_height = TEST_TREE_HEIGHT;
let circom_path = format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/rln.wasm");
let zkey_path = format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/rln_final.zkey");
let vk_path =
format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/verification_key.json");
let zkey = read_file(&zkey_path).unwrap();
let vk = read_file(&vk_path).unwrap();

// Creating an instance of RLN
let rln_instance = wasm_new(tree_height, zkey, vk).unwrap();
let circom_path = format!("../rln/resources/tree_height_20/rln.wasm");
let rln_instance = wasm_new().unwrap();

// Creating membership key
let mem_keys = wasm_key_gen(rln_instance).unwrap();
Expand Down Expand Up @@ -123,18 +116,10 @@ mod tests {
let is_proof_valid = wasm_verify_with_roots(rln_instance, proof_with_signal, roots);
assert!(is_proof_valid.unwrap(), "verifying proof with roots failed");
}

#[wasm_bindgen_test]
fn test_metadata() {
let tree_height = TEST_TREE_HEIGHT;
let zkey_path = format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/rln_final.zkey");
let vk_path =
format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/verification_key.json");
let zkey = read_file(&zkey_path).unwrap();
let vk = read_file(&vk_path).unwrap();

// Creating an instance of RLN
let rln_instance = wasm_new(tree_height, zkey, vk).unwrap();
let rln_instance = wasm_new().unwrap();

let test_metadata = Uint8Array::new(&JsValue::from_str("test"));
// Inserting random metadata
Expand Down
Binary file removed rln/resources/tree_height_32/rln.wasm
Binary file not shown.
Binary file removed rln/resources/tree_height_32/rln_final.zkey
Binary file not shown.
114 changes: 0 additions & 114 deletions rln/resources/tree_height_32/verification_key.json

This file was deleted.

91 changes: 31 additions & 60 deletions rln/src/circuit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,39 +5,31 @@ use ark_bn254::{
G1Projective as ArkG1Projective, G2Affine as ArkG2Affine, G2Projective as ArkG2Projective,
};
use ark_circom::read_zkey;
use ark_circom::WitnessCalculator;
use ark_groth16::{ProvingKey, VerifyingKey};
use ark_relations::r1cs::ConstraintMatrices;
use cfg_if::cfg_if;
use color_eyre::{Report, Result};
use num_bigint::BigUint;
use serde_json::Value;
use std::io::Cursor;
use std::str::FromStr;
use wasmer::{Module, Store};

use cfg_if::cfg_if;

cfg_if! {
if #[cfg(not(target_arch = "wasm32"))] {
use ark_circom::{WitnessCalculator};
use once_cell::sync::OnceCell;
use std::sync::Mutex;
use wasmer::{Module, Store};
use include_dir::{include_dir, Dir};
use std::path::Path;
}
}

const ZKEY_FILENAME: &str = "rln_final.zkey";
const VK_FILENAME: &str = "verification_key.json";
const WASM_FILENAME: &str = "rln.wasm";

// These parameters are used for tests
// Note that the circuit and keys in TEST_RESOURCES_FOLDER are compiled for Merkle trees of height 20 & 32
// Changing these parameters to other values than these defaults will cause zkSNARK proof verification to fail
pub const TEST_PARAMETERS_INDEX: usize = 0;
pub const TEST_TREE_HEIGHT: usize = [20, 32][TEST_PARAMETERS_INDEX];
pub const TEST_RESOURCES_FOLDER: &str = ["tree_height_20", "tree_height_32"][TEST_PARAMETERS_INDEX];
const ZKEY_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/rln_final.zkey");

#[cfg(not(target_arch = "wasm32"))]
static RESOURCES_DIR: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/resources");
const VK_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/verification_key.json");

pub const WASM_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/rln.wasm");

// The following types define the pairing friendly elliptic curve, the underlying finite fields and groups default to this module
// Note that proofs are serialized assuming Fr to be 4x8 = 32 bytes in size. Hence, changing to a curve with different encoding will make proof verification to fail
Expand All @@ -62,18 +54,10 @@ pub fn zkey_from_raw(zkey_data: &Vec<u8>) -> Result<(ProvingKey<Curve>, Constrai
}

// Loads the proving key
#[cfg(not(target_arch = "wasm32"))]
pub fn zkey_from_folder(
resources_folder: &str,
) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> {
let zkey = RESOURCES_DIR.get_file(Path::new(resources_folder).join(ZKEY_FILENAME));
if let Some(zkey) = zkey {
let mut c = Cursor::new(zkey.contents());
let proving_key_and_matrices = read_zkey(&mut c)?;
Ok(proving_key_and_matrices)
} else {
Err(Report::msg("No proving key found!"))
}
pub fn default_zkey() -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> {
let mut c = Cursor::new(ZKEY_BYTES);
let proving_key_and_matrices = read_zkey(&mut c)?;
Ok(proving_key_and_matrices)
}

// Loads the verification key from a bytes vector
Expand All @@ -93,24 +77,8 @@ pub fn vk_from_raw(vk_data: &Vec<u8>, zkey_data: &Vec<u8>) -> Result<VerifyingKe
}

// Loads the verification key
#[cfg(not(target_arch = "wasm32"))]
pub fn vk_from_folder(resources_folder: &str) -> Result<VerifyingKey<Curve>> {
let vk = RESOURCES_DIR.get_file(Path::new(resources_folder).join(VK_FILENAME));
let zkey = RESOURCES_DIR.get_file(Path::new(resources_folder).join(ZKEY_FILENAME));

let verifying_key: VerifyingKey<Curve>;
if let Some(vk) = vk {
verifying_key = vk_from_json(vk.contents_utf8().ok_or(Report::msg(
"Could not read verification key from JSON file!",
))?)?;
Ok(verifying_key)
} else if let Some(_zkey) = zkey {
let (proving_key, _matrices) = zkey_from_folder(resources_folder)?;
verifying_key = proving_key.vk;
Ok(verifying_key)
} else {
Err(Report::msg("No proving/verification key found!"))
}
pub fn default_vk() -> Result<VerifyingKey<Curve>> {
vk_from_vector(VK_BYTES)
}

#[cfg(not(target_arch = "wasm32"))]
Expand All @@ -127,18 +95,25 @@ pub fn circom_from_raw(wasm_buffer: Vec<u8>) -> Result<&'static Mutex<WitnessCal
})
}

#[cfg(target_arch = "wasm32")]
pub fn circom_from_raw(wasm_buffer: Vec<u8>) -> Result<WitnessCalculator> {
let store = Store::default();
let module = Module::new(&store, wasm_buffer)?;
let witness_calculator = WitnessCalculator::from_module(module)?;
Ok(witness_calculator)
}

// Initializes the witness calculator
#[cfg(not(target_arch = "wasm32"))]
pub fn circom_from_folder(resources_folder: &str) -> Result<&'static Mutex<WitnessCalculator>> {
pub fn default_circom() -> Result<&'static Mutex<WitnessCalculator>> {
// We read the wasm file
let wasm = RESOURCES_DIR.get_file(Path::new(resources_folder).join(WASM_FILENAME));
circom_from_raw(WASM_BYTES.into())
}

if let Some(wasm) = wasm {
let wasm_buffer = wasm.contents();
circom_from_raw(wasm_buffer.to_vec())
} else {
Err(Report::msg("No wasm file found!"))
}
#[cfg(target_arch = "wasm32")]
pub fn default_circom() -> Result<WitnessCalculator> {
// We read the wasm file
circom_from_raw(WASM_BYTES.into())
}

// The following function implementations are taken/adapted from https://github.com/gakonst/ark-circom/blob/1732e15d6313fe176b0b1abb858ac9e095d0dbd7/src/zkey.rs
Expand Down Expand Up @@ -253,12 +228,8 @@ fn vk_from_vector(vk: &[u8]) -> Result<VerifyingKey<Curve>> {
}

// Checks verification key to be correct with respect to proving key
#[cfg(not(target_arch = "wasm32"))]
pub fn check_vk_from_zkey(
resources_folder: &str,
verifying_key: VerifyingKey<Curve>,
) -> Result<()> {
let (proving_key, _matrices) = zkey_from_folder(resources_folder)?;
pub fn check_vk_from_zkey(verifying_key: VerifyingKey<Curve>) -> Result<()> {
let (proving_key, _matrices) = default_zkey()?;
if proving_key.vk == verifying_key {
Ok(())
} else {
Expand Down
Loading
Loading