Replace the ark-zkey witness calculator with the one of iden3 (#273)

* Add files via upload * Add files via upload * Add files via upload * Add files via upload * Add files via upload * Add files via upload * Add files via upload * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Update README.md * Delete rln/resources/tree_height_20/rln.wasm * Changes in accordance with the results of the first round of reviewing. * Formatting * Acknowledgements and a readme fix * Minor change: Vec memory allocation
vacp2p · Dec 25, 2024 · e39f156 · e39f156
1 parent 8b04930
commit e39f156
Show file tree

Hide file tree

Showing 19 changed files with 1,981 additions and 189 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/README.md b/README.md
@@ -18,6 +18,8 @@ in Rust.
 
 - [semaphore-rs](https://github.com/worldcoin/semaphore-rs) written in Rust based on ark-circom.
 
+- The circom witness calculation code of the rln crate is based on [circom-witnesscalc](https://github.com/iden3/circom-witnesscalc) by iden3. The execution graph file used by this code has been generated by means of the same iden3 software.
+
 ## Users
 
 Zerokit is used by -

diff --git a/rln-cli/Cargo.toml b/rln-cli/Cargo.toml
@@ -11,3 +11,6 @@ color-eyre = "=0.6.2"
 # serialization
 serde_json = "1.0.48"
 serde = { version = "1.0.130", features = ["derive"] }
+
+[features]
+arkzkey = []
diff --git a/rln-cli/src/main.rs b/rln-cli/src/main.rs
@@ -27,7 +27,7 @@ fn main() -> Result<()> {
             tree_height,
             config,
         }) => {
-            let resources = File::open(&config)?;
+            let resources = File::open(config)?;
             state.rln = Some(RLN::new(*tree_height, resources)?);
             Ok(())
         }
@@ -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.arkvkey"];
+            let filenames = ["rln_final.arkzkey", "verification_key.arkvkey"];
             #[cfg(not(feature = "arkzkey"))]
-            let filenames = ["rln.wasm", "rln_final.zkey", "verification_key.arkvkey"];
+            let filenames = ["rln_final.zkey", "verification_key.arkvkey"];
             for filename in filenames {
                 let fullpath = config.join(Path::new(filename));
                 let mut file = File::open(&fullpath)?;
@@ -49,12 +49,11 @@ fn main() -> Result<()> {
                 file.read_exact(&mut buffer)?;
                 resources.push(buffer);
             }
-            let tree_config_input_file = File::open(&tree_config_input)?;
+            let tree_config_input_file = File::open(tree_config_input)?;
             state.rln = Some(RLN::new_with_params(
                 *tree_height,
                 resources[0].clone(),
                 resources[1].clone(),
-                resources[2].clone(),
                 tree_config_input_file,
             )?);
             Ok(())
@@ -67,31 +66,31 @@ fn main() -> Result<()> {
             Ok(())
         }
         Some(Commands::SetLeaf { index, file }) => {
-            let input_data = File::open(&file)?;
+            let input_data = File::open(file)?;
             state
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .set_leaf(*index, input_data)?;
             Ok(())
         }
         Some(Commands::SetMultipleLeaves { index, file }) => {
-            let input_data = File::open(&file)?;
+            let input_data = File::open(file)?;
             state
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .set_leaves_from(*index, input_data)?;
             Ok(())
         }
         Some(Commands::ResetMultipleLeaves { file }) => {
-            let input_data = File::open(&file)?;
+            let input_data = File::open(file)?;
             state
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .init_tree_with_leaves(input_data)?;
             Ok(())
         }
         Some(Commands::SetNextLeaf { file }) => {
-            let input_data = File::open(&file)?;
+            let input_data = File::open(file)?;
             state
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?
@@ -122,7 +121,7 @@ fn main() -> Result<()> {
             Ok(())
         }
         Some(Commands::Prove { input }) => {
-            let input_data = File::open(&input)?;
+            let input_data = File::open(input)?;
             let writer = std::io::stdout();
             state
                 .rln
@@ -131,15 +130,15 @@ fn main() -> Result<()> {
             Ok(())
         }
         Some(Commands::Verify { file }) => {
-            let input_data = File::open(&file)?;
+            let input_data = File::open(file)?;
             state
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .verify(input_data)?;
             Ok(())
         }
         Some(Commands::GenerateProof { input }) => {
-            let input_data = File::open(&input)?;
+            let input_data = File::open(input)?;
             let writer = std::io::stdout();
             state
                 .rln
@@ -148,8 +147,8 @@ fn main() -> Result<()> {
             Ok(())
         }
         Some(Commands::VerifyWithRoots { input, roots }) => {
-            let input_data = File::open(&input)?;
-            let roots_data = File::open(&roots)?;
+            let input_data = File::open(input)?;
+            let roots_data = File::open(roots)?;
             state
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?

diff --git a/rln/Cargo.toml b/rln/Cargo.toml
@@ -42,6 +42,7 @@ color-eyre = "=0.6.2"
 thiserror = "=1.0.39"
 
 # utilities
+byteorder = "1.4.3"
 cfg-if = "=1.0"
 num-bigint = { version = "=0.4.3", default-features = false, features = [
     "rand",
@@ -51,11 +52,13 @@ once_cell = "=1.17.1"
 lazy_static = "=1.4.0"
 rand = "=0.8.5"
 rand_chacha = "=0.3.1"
+ruint = { version = "1.10.0", features = ["rand", "serde", "ark-ff-04", "num-bigint"] }
 tiny-keccak = { version = "=2.0.2", features = ["keccak"] }
 utils = { package = "zerokit_utils", version = "=0.5.1", path = "../utils/", default-features = false }
 
 
 # serialization
+prost = "0.13.1"
 serde_json = "=1.0.96"
 serde = { version = "=1.0.163", features = ["derive"] }
 

diff --git a/rln/Makefile.toml b/rln/Makefile.toml
@@ -4,7 +4,7 @@ args = ["build", "--release"]
 
 [tasks.test_default]
 command = "cargo"
-args = ["test", "--release"]
+args = ["test", "--release", "--", "--nocapture"]
 
 [tasks.test_stateless]
 command = "cargo"

diff --git a/rln/README.md b/rln/README.md
@@ -12,14 +12,18 @@ git clone https://github.com/vacp2p/zerokit.git
 cd zerokit/rln
 ```
 
-### Build and Test
+ ### Build and Test
 
-To build and test, run the following commands within the module folder
+ To build and test, run the following commands within the module folder
 
-```bash
-cargo make build
-cargo make test
+``` bash
+ cargo make build
+ cargo make test_{mode}
 ```
+The {mode} placeholder should be replaced with
+* **default** for the default tests;
+* **arkzkey** for the tests with the arkzkey feature;
+* **stateless** for the tests with the stateless feature.
 
 ### Compile ZK circuits
 

diff --git a/rln/resources/tree_height_20/graph.bin b/rln/resources/tree_height_20/graph.bin
diff --git a/rln/src/circuit.rs b/rln/src/circuit.rs
@@ -1,5 +1,6 @@
 // This crate provides interfaces for the zero-knowledge circuit and keys
 
+use crate::iden3calc::calc_witness;
 use ark_bn254::{
     Bn254, Fq as ArkFq, Fq2 as ArkFq2, Fr as ArkFr, G1Affine as ArkG1Affine,
     G1Projective as ArkG1Projective, G2Affine as ArkG2Affine, G2Projective as ArkG2Projective,
@@ -9,14 +10,10 @@ use ark_relations::r1cs::ConstraintMatrices;
 use ark_serialize::CanonicalDeserialize;
 use cfg_if::cfg_if;
 use color_eyre::{Report, Result};
+use num_bigint::BigInt;
 
 #[cfg(not(target_arch = "wasm32"))]
-use {
-    ark_circom::WitnessCalculator,
-    lazy_static::lazy_static,
-    std::sync::{Arc, Mutex},
-    wasmer::{Module, Store},
-};
+use ::lazy_static::lazy_static;
 
 #[cfg(feature = "arkzkey")]
 use {
@@ -35,7 +32,7 @@ pub const ARKZKEY_BYTES_UNCOMPR: &[u8] =
 
 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");
+const GRAPH_BYTES: &[u8] = include_bytes!("../resources/tree_height_20/graph.bin");
 
 #[cfg(not(target_arch = "wasm32"))]
 lazy_static! {
@@ -53,11 +50,6 @@ lazy_static! {
 
     #[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 ref WITNESS_CALCULATOR: Arc<Mutex<WitnessCalculator>> = {
-        circom_from_raw(WASM_BYTES).expect("Failed to create witness calculator")
-    };
 }
 
 pub const TEST_TREE_HEIGHT: usize = 20;
@@ -92,6 +84,10 @@ pub fn zkey_from_raw(zkey_data: &[u8]) -> Result<(ProvingKey<Curve>, ConstraintM
     Ok(proving_key_and_matrices)
 }
 
+pub fn calculate_rln_witness<I: IntoIterator<Item = (String, Vec<BigInt>)>>(inputs: I) -> Vec<Fr> {
+    calc_witness(inputs, GRAPH_BYTES)
+}
+
 // Loads the proving key
 #[cfg(not(target_arch = "wasm32"))]
 pub fn zkey_from_folder() -> &'static (ProvingKey<Curve>, ConstraintMatrices<Fr>) {
@@ -118,20 +114,6 @@ pub fn vk_from_folder() -> &'static VerifyingKey<Curve> {
     &VK
 }
 
-// Initializes the witness calculator using a bytes vector
-#[cfg(not(target_arch = "wasm32"))]
-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)?;
-    Ok(Arc::new(Mutex::new(result)))
-}
-
-// Initializes the witness calculator
-#[cfg(not(target_arch = "wasm32"))]
-pub fn circom_from_folder() -> &'static Arc<Mutex<WitnessCalculator>> {
-    &WITNESS_CALCULATOR
-}
-
 // 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>> {

diff --git a/rln/src/ffi.rs b/rln/src/ffi.rs
@@ -228,15 +228,13 @@ pub extern "C" fn new(ctx: *mut *mut RLN) -> bool {
 #[no_mangle]
 pub extern "C" fn new_with_params(
     tree_height: usize,
-    circom_buffer: *const Buffer,
     zkey_buffer: *const Buffer,
     vk_buffer: *const Buffer,
     tree_config: *const Buffer,
     ctx: *mut *mut RLN,
 ) -> bool {
     match RLN::new_with_params(
         tree_height,
-        circom_buffer.process().to_vec(),
         zkey_buffer.process().to_vec(),
         vk_buffer.process().to_vec(),
         tree_config.process(),
@@ -256,16 +254,11 @@ pub extern "C" fn new_with_params(
 #[cfg(feature = "stateless")]
 #[no_mangle]
 pub extern "C" fn new_with_params(
-    circom_buffer: *const Buffer,
     zkey_buffer: *const Buffer,
     vk_buffer: *const Buffer,
     ctx: *mut *mut RLN,
 ) -> bool {
-    match RLN::new_with_params(
-        circom_buffer.process().to_vec(),
-        zkey_buffer.process().to_vec(),
-        vk_buffer.process().to_vec(),
-    ) {
+    match RLN::new_with_params(zkey_buffer.process().to_vec(), vk_buffer.process().to_vec()) {
         Ok(rln) => {
             unsafe { *ctx = Box::into_raw(Box::new(rln)) };
             true

diff --git a/rln/src/iden3calc.rs b/rln/src/iden3calc.rs
@@ -0,0 +1,73 @@
+// This file is based on the code by iden3. Its preimage can be found here:
+// https://github.com/iden3/circom-witnesscalc/blob/5cb365b6e4d9052ecc69d4567fcf5bc061c20e94/src/lib.rs
+
+pub mod graph;
+pub mod proto;
+pub mod storage;
+
+use ark_bn254::Fr;
+use graph::Node;
+use num_bigint::BigInt;
+use ruint::aliases::U256;
+use std::collections::HashMap;
+use storage::deserialize_witnesscalc_graph;
+
+pub type InputSignalsInfo = HashMap<String, (usize, usize)>;
+
+pub fn calc_witness<I: IntoIterator<Item = (String, Vec<BigInt>)>>(
+    inputs: I,
+    graph_data: &[u8],
+) -> Vec<Fr> {
+    let inputs: HashMap<String, Vec<U256>> = inputs
+        .into_iter()
+        .map(|(key, value)| (key, value.iter().map(|v| U256::from(v)).collect()))
+        .collect();
+
+    let (nodes, signals, input_mapping): (Vec<Node>, Vec<usize>, InputSignalsInfo) =
+        deserialize_witnesscalc_graph(std::io::Cursor::new(graph_data)).unwrap();
+
+    let mut inputs_buffer = get_inputs_buffer(get_inputs_size(&nodes));
+    populate_inputs(&inputs, &input_mapping, &mut inputs_buffer);
+
+    graph::evaluate(&nodes, inputs_buffer.as_slice(), &signals)
+}
+
+fn get_inputs_size(nodes: &[Node]) -> usize {
+    let mut start = false;
+    let mut max_index = 0usize;
+    for &node in nodes.iter() {
+        if let Node::Input(i) = node {
+            if i > max_index {
+                max_index = i;
+            }
+            start = true
+        } else if start {
+            break;
+        }
+    }
+    max_index + 1
+}
+
+fn populate_inputs(
+    input_list: &HashMap<String, Vec<U256>>,
+    inputs_info: &InputSignalsInfo,
+    input_buffer: &mut [U256],
+) {
+    for (key, value) in input_list {
+        let (offset, len) = inputs_info[key];
+        if len != value.len() {
+            panic!("Invalid input length for {}", key);
+        }
+
+        for (i, v) in value.iter().enumerate() {
+            input_buffer[offset + i] = *v;
+        }
+    }
+}
+
+/// Allocates inputs vec with position 0 set to 1
+fn get_inputs_buffer(size: usize) -> Vec<U256> {
+    let mut inputs = vec![U256::ZERO; size];
+    inputs[0] = U256::from(1);
+    inputs
+}