chore(rln): use ark serialized verification key for faster serde (#259)

* chore(rln): use ark serialized verification key for faster serde * fix: unused imports * fix: rm verification_key.json * fix: s/vk_from_slice/vk_from_ark_serialized/g
vacp2p · Jun 14, 2024 · c6493bd · c6493bd
1 parent dd5edd6
commit c6493bd
Show file tree

Hide file tree

Showing 10 changed files with 21 additions and 264 deletions.
diff --git a/rln-cli/src/main.rs b/rln-cli/src/main.rs
@@ -38,9 +38,9 @@ fn main() -> Result<()> {
         }) => {
             let mut resources: Vec<Vec<u8>> = Vec::new();
             #[cfg(feature = "arkzkey")]
-            let filenames = ["rln.wasm", "rln_final.arkzkey", "verification_key.json"];
+            let filenames = ["rln.wasm", "rln_final.arkzkey", "verification_key.arkvkey"];
             #[cfg(not(feature = "arkzkey"))]
-            let filenames = ["rln.wasm", "rln_final.zkey", "verification_key.json"];
+            let filenames = ["rln.wasm", "rln_final.zkey", "verification_key.arkvkey"];
             for filename in filenames {
                 let fullpath = config.join(Path::new(filename));
                 let mut file = File::open(&fullpath)?;

diff --git a/rln-wasm/tests/rln-wasm.rs b/rln-wasm/tests/rln-wasm.rs
@@ -25,7 +25,7 @@ mod tests {
         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");
+            format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/verification_key.arkvkey");
         let zkey = read_file(&zkey_path).unwrap();
         let vk = read_file(&vk_path).unwrap();
 
@@ -129,7 +129,7 @@ mod tests {
         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");
+            format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/verification_key.arkvkey");
         let zkey = read_file(&zkey_path).unwrap();
         let vk = read_file(&vk_path).unwrap();
 

diff --git a/rln/README.md b/rln/README.md
@@ -74,7 +74,7 @@ rln = { git = "https://github.com/vacp2p/zerokit" }
 
 First, we need to create a RLN object for a chosen input Merkle tree size.
 
-Note that we need to pass to RLN object constructor the path where the circuit (`rln.wasm`, built for the input tree size), the corresponding proving key (`rln_final.zkey`) or (`rln_final.arkzkey`) and verification key (`verification_key.json`, optional) are found.
+Note that we need to pass to RLN object constructor the path where the circuit (`rln.wasm`, built for the input tree size), the corresponding proving key (`rln_final.zkey`) or (`rln_final.arkzkey`) and verification key (`verification_key.arkvkey`, optional) are found.
 
 In the following we will use [cursors](https://doc.rust-lang.org/std/io/struct.Cursor.html) as readers/writers for interfacing with RLN public APIs.
 

diff --git a/rln/benches/circuit_deser_benchmark.rs b/rln/benches/circuit_deser_benchmark.rs
@@ -1,19 +1,17 @@
 use criterion::{criterion_group, criterion_main, Criterion};
-use rln::circuit::{to_verifying_key, RESOURCES_DIR, VK_FILENAME};
-use serde_json::Value;
+use rln::circuit::{vk_from_ark_serialized, RESOURCES_DIR, VK_FILENAME};
 use std::path::Path;
 
 // Here we benchmark how long the deserialization of the
 // verifying_key takes, only testing the json => verifying_key conversion,
 // and skipping conversion from bytes => string => serde_json::Value
 pub fn vk_deserialize_benchmark(c: &mut Criterion) {
     let vk = RESOURCES_DIR.get_file(Path::new(VK_FILENAME)).unwrap();
-    let vk = vk.contents_utf8().unwrap();
-    let json: Value = serde_json::from_str(vk).unwrap();
+    let vk = vk.contents();
 
     c.bench_function("circuit::to_verifying_key", |b| {
         b.iter(|| {
-            let _ = to_verifying_key(&json);
+            let _ = vk_from_ark_serialized(&vk);
         })
     });
 }

diff --git a/rln/resources/tree_height_20/verification_key.arkvkey b/rln/resources/tree_height_20/verification_key.arkvkey
diff --git a/rln/resources/tree_height_20/verification_key.json b/rln/resources/tree_height_20/verification_key.json
diff --git a/rln/src/circuit.rs b/rln/src/circuit.rs
@@ -6,11 +6,9 @@ use ark_bn254::{
 };
 use ark_groth16::{ProvingKey, VerifyingKey};
 use ark_relations::r1cs::ConstraintMatrices;
+use ark_serialize::CanonicalDeserialize;
 use cfg_if::cfg_if;
 use color_eyre::{Report, Result};
-use num_bigint::BigUint;
-use serde_json::Value;
-use std::str::FromStr;
 
 cfg_if! {
     if #[cfg(not(target_arch = "wasm32"))] {
@@ -35,7 +33,7 @@ cfg_if! {
 }
 
 const ZKEY_FILENAME: &str = "tree_height_20/rln_final.zkey";
-pub const VK_FILENAME: &str = "tree_height_20/verification_key.json";
+pub const VK_FILENAME: &str = "tree_height_20/verification_key.arkvkey";
 const WASM_FILENAME: &str = "tree_height_20/rln.wasm";
 
 pub const TEST_TREE_HEIGHT: usize = 20;
@@ -101,7 +99,7 @@ pub fn vk_from_raw(vk_data: &[u8], zkey_data: &Vec<u8>) -> Result<VerifyingKey<C
     let verifying_key: VerifyingKey<Curve>;
 
     if !vk_data.is_empty() {
-        verifying_key = vk_from_vector(vk_data)?;
+        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)?;
@@ -120,9 +118,7 @@ pub fn vk_from_folder() -> Result<VerifyingKey<Curve>> {
 
     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!",
-        ))?)?;
+        verifying_key = vk_from_ark_serialized(vk.contents())?;
         Ok(verifying_key)
     } else if let Some(_zkey) = zkey {
         let (proving_key, _matrices) = zkey_from_folder()?;
@@ -161,115 +157,11 @@ pub fn circom_from_folder() -> Result<&'static Mutex<WitnessCalculator>> {
     }
 }
 
-// The following function implementations are taken/adapted from https://github.com/gakonst/ark-circom/blob/1732e15d6313fe176b0b1abb858ac9e095d0dbd7/src/zkey.rs
-
-// Utilities to convert a json verification key in a groth16::VerificationKey
-fn fq_from_str(s: &str) -> Result<Fq> {
-    Ok(Fq::from(BigUint::from_str(s)?))
-}
-
-// Extracts the element in G1 corresponding to its JSON serialization
-fn json_to_g1(json: &Value, key: &str) -> Result<G1Affine> {
-    let els: Vec<String> = json
-        .get(key)
-        .ok_or(Report::msg("no json value"))?
-        .as_array()
-        .ok_or(Report::msg("value not an array"))?
-        .iter()
-        .map(|i| i.as_str().ok_or(Report::msg("element is not a string")))
-        .map(|x| x.map(|v| v.to_owned()))
-        .collect::<Result<Vec<String>>>()?;
-
-    Ok(G1Affine::from(G1Projective::new(
-        fq_from_str(&els[0])?,
-        fq_from_str(&els[1])?,
-        fq_from_str(&els[2])?,
-    )))
-}
-
-// Extracts the vector of G1 elements corresponding to its JSON serialization
-fn json_to_g1_vec(json: &Value, key: &str) -> Result<Vec<G1Affine>> {
-    let els: Vec<Vec<String>> = json
-        .get(key)
-        .ok_or(Report::msg("no json value"))?
-        .as_array()
-        .ok_or(Report::msg("value not an array"))?
-        .iter()
-        .map(|i| {
-            i.as_array()
-                .ok_or(Report::msg("element is not an array"))
-                .and_then(|array| {
-                    array
-                        .iter()
-                        .map(|x| x.as_str().ok_or(Report::msg("element is not a string")))
-                        .map(|x| x.map(|v| v.to_owned()))
-                        .collect::<Result<Vec<String>>>()
-                })
-        })
-        .collect::<Result<Vec<Vec<String>>>>()?;
-
-    let mut res = vec![];
-    for coords in els {
-        res.push(G1Affine::from(G1Projective::new(
-            fq_from_str(&coords[0])?,
-            fq_from_str(&coords[1])?,
-            fq_from_str(&coords[2])?,
-        )))
-    }
-
-    Ok(res)
-}
-
-// Extracts the element in G2 corresponding to its JSON serialization
-fn json_to_g2(json: &Value, key: &str) -> Result<G2Affine> {
-    let els: Vec<Vec<String>> = json
-        .get(key)
-        .ok_or(Report::msg("no json value"))?
-        .as_array()
-        .ok_or(Report::msg("value not an array"))?
-        .iter()
-        .map(|i| {
-            i.as_array()
-                .ok_or(Report::msg("element is not an array"))
-                .and_then(|array| {
-                    array
-                        .iter()
-                        .map(|x| x.as_str().ok_or(Report::msg("element is not a string")))
-                        .map(|x| x.map(|v| v.to_owned()))
-                        .collect::<Result<Vec<String>>>()
-                })
-        })
-        .collect::<Result<Vec<Vec<String>>>>()?;
-
-    let x = Fq2::new(fq_from_str(&els[0][0])?, fq_from_str(&els[0][1])?);
-    let y = Fq2::new(fq_from_str(&els[1][0])?, fq_from_str(&els[1][1])?);
-    let z = Fq2::new(fq_from_str(&els[2][0])?, fq_from_str(&els[2][1])?);
-    Ok(G2Affine::from(G2Projective::new(x, y, z)))
-}
-
-// Converts JSON to a VerifyingKey
-pub fn to_verifying_key(json: &serde_json::Value) -> Result<VerifyingKey<Curve>> {
-    Ok(VerifyingKey {
-        alpha_g1: json_to_g1(json, "vk_alpha_1")?,
-        beta_g2: json_to_g2(json, "vk_beta_2")?,
-        gamma_g2: json_to_g2(json, "vk_gamma_2")?,
-        delta_g2: json_to_g2(json, "vk_delta_2")?,
-        gamma_abc_g1: json_to_g1_vec(json, "IC")?,
-    })
-}
-
-// Computes the verification key from its JSON serialization
-fn vk_from_json(vk: &str) -> Result<VerifyingKey<Curve>> {
-    let json: Value = serde_json::from_str(vk)?;
-    to_verifying_key(&json)
-}
-
-// Computes the verification key from a bytes vector containing its JSON serialization
-fn vk_from_vector(vk: &[u8]) -> Result<VerifyingKey<Curve>> {
-    let json = String::from_utf8(vk.to_vec())?;
-    let json: Value = serde_json::from_str(&json)?;
-
-    to_verifying_key(&json)
+// Computes the verification key from a bytes vector containing pre-processed ark-serialized verification key
+// uncompressed, unchecked
+pub fn vk_from_ark_serialized(data: &[u8]) -> Result<VerifyingKey<Curve>> {
+    let vk = VerifyingKey::<Curve>::deserialize_uncompressed_unchecked(data)?;
+    Ok(vk)
 }
 
 // Checks verification key to be correct with respect to proving key

diff --git a/rln/src/public.rs b/rln/src/public.rs
@@ -112,7 +112,7 @@ impl RLN<'_> {
     /// - `tree_height`: the height of the internal Merkle tree
     /// - `circom_vec`: a byte vector containing the ZK circuit (`rln.wasm`) as binary file
     /// - `zkey_vec`: a byte vector containing to the proving key (`rln_final.zkey`)  or (`rln_final.arkzkey`) as binary file
-    /// - `vk_vec`: a byte vector containing to the verification key (`verification_key.json`) as binary file
+    /// - `vk_vec`: a byte vector containing to the verification key (`verification_key.arkvkey`) as binary file
     /// - `tree_config_input`: a reader for a string containing a json with the merkle tree configuration
     ///
     /// Example:
@@ -124,7 +124,7 @@ impl RLN<'_> {
     /// let resources_folder = "./resources/tree_height_20/";
     ///
     /// let mut resources: Vec<Vec<u8>> = Vec::new();
-    /// for filename in ["rln.wasm", "rln_final.zkey", "verification_key.json"] {
+    /// for filename in ["rln.wasm", "rln_final.zkey", "verification_key.arkvkey"] {
     ///     let fullpath = format!("{resources_folder}{filename}");
     ///     let mut file = File::open(&fullpath).expect("no file found");
     ///     let metadata = std::fs::metadata(&fullpath).expect("unable to read metadata");

diff --git a/rln/src/public_api_tests.rs b/rln/src/public_api_tests.rs
@@ -1,11 +1,10 @@
-use crate::circuit::{Curve, Fr, TEST_TREE_HEIGHT};
+use crate::circuit::{Fr, TEST_TREE_HEIGHT};
 use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash};
 use crate::protocol::*;
 use crate::public::RLN;
 use crate::utils::*;
 use ark_groth16::Proof as ArkProof;
 use ark_serialize::{CanonicalDeserialize, Read};
-use num_bigint::BigInt;
 use std::io::Cursor;
 use std::str::FromStr;
 use utils::ZerokitMerkleTree;
@@ -646,24 +645,6 @@ fn test_rln_with_witness() {
 
     let serialized_witness = serialize_witness(&rln_witness).unwrap();
 
-    // Calculate witness outside zerokit (simulating what JS is doing)
-    let inputs = inputs_for_witness_calculation(&rln_witness)
-        .unwrap()
-        .into_iter()
-        .map(|(name, values)| (name.to_string(), values));
-    let calculated_witness = rln
-        .witness_calculator
-        .lock()
-        .expect("witness_calculator mutex should not get poisoned")
-        .calculate_witness_element::<Curve, _>(inputs, false)
-        .map_err(ProofError::WitnessError)
-        .unwrap();
-
-    let calculated_witness_vec: Vec<BigInt> = calculated_witness
-        .into_iter()
-        .map(|v| to_bigint(&v).unwrap())
-        .collect();
-
     // Generating the proof
     let mut input_buffer = Cursor::new(serialized_witness);
     let mut output_buffer = Cursor::new(Vec::<u8>::new());