From 1628a69a2bb2d518254c695cf7ad714ad866951d Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Fri, 4 Aug 2023 11:28:07 -0400
Subject: [PATCH] chore: accept tree_config in new_with_params

---
 rln/src/ffi.rs    |  2 ++
 rln/src/public.rs | 30 +++++++++++++++++++++++++++---
 rln/tests/ffi.rs  |  3 +++
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/rln/src/ffi.rs b/rln/src/ffi.rs
index a8a76d51..55abe32f 100644
--- a/rln/src/ffi.rs
+++ b/rln/src/ffi.rs
@@ -186,6 +186,7 @@ pub extern "C" fn new_with_params(
     circom_buffer: *const Buffer,
     zkey_buffer: *const Buffer,
     vk_buffer: *const Buffer,
+    tree_config: *const Buffer,
     ctx: *mut *mut RLN,
 ) -> bool {
     if let Ok(rln) = RLN::new_with_params(
@@ -193,6 +194,7 @@ pub extern "C" fn new_with_params(
         circom_buffer.process().to_vec(),
         zkey_buffer.process().to_vec(),
         vk_buffer.process().to_vec(),
+        tree_config.process(),
     ) {
         unsafe { *ctx = Box::into_raw(Box::new(rln)) };
         true
diff --git a/rln/src/public.rs b/rln/src/public.rs
index 499711ce..38518b1e 100644
--- a/rln/src/public.rs
+++ b/rln/src/public.rs
@@ -117,6 +117,7 @@ impl RLN<'_> {
     /// - `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`) as binary file
     /// - `vk_vec`: a byte vector containing to the verification key (`verification_key.json`) as binary file
+    /// - `tree_config`: a reader for a string containing a json with the merkle tree configuration
     ///
     /// Example:
     /// ```
@@ -134,6 +135,8 @@ impl RLN<'_> {
     ///     let mut buffer = vec![0; metadata.len() as usize];
     ///     file.read_exact(&mut buffer).expect("buffer overflow");
     ///     resources.push(buffer);
+    ///     let tree_config = "{}".to_string();
+    ///     let tree_config_buffer = &Buffer::from(tree_config.as_bytes());
     /// }
     ///
     /// let mut rln = RLN::new_with_params(
@@ -141,13 +144,15 @@ impl RLN<'_> {
     ///     resources[0].clone(),
     ///     resources[1].clone(),
     ///     resources[2].clone(),
+    ///     tree_config_buffer,
     /// );
     /// ```
-    pub fn new_with_params(
+    pub fn new_with_params<R: Read>(
         tree_height: usize,
         #[cfg(not(target_arch = "wasm32"))] circom_vec: Vec<u8>,
         zkey_vec: Vec<u8>,
         vk_vec: Vec<u8>,
+        #[cfg(not(target_arch = "wasm32"))] mut tree_config_input: R,
     ) -> Result<RLN<'static>> {
         #[cfg(not(target_arch = "wasm32"))]
         let witness_calculator = circom_from_raw(circom_vec)?;
@@ -155,8 +160,27 @@ impl RLN<'_> {
         let proving_key = zkey_from_raw(&zkey_vec)?;
         let verification_key = vk_from_raw(&vk_vec, &zkey_vec)?;
 
-        // We compute a default empty tree
-        let tree = PoseidonTree::default(tree_height)?;
+        let tree = if cfg!(not(target_arch = "wasm32")) {
+            let mut tree_config_vec: Vec<u8> = Vec::new();
+            tree_config_input.read_to_end(&mut tree_config_vec)?;
+            let tree_config_str = String::from_utf8(tree_config_vec)?;
+            let tree_config: <PoseidonTree as ZerokitMerkleTree>::Config =
+                if tree_config_str.is_empty() {
+                    <PoseidonTree as ZerokitMerkleTree>::Config::default()
+                } else {
+                    <PoseidonTree as ZerokitMerkleTree>::Config::from_str(&tree_config_str)?
+                };
+
+            // We compute a default empty tree
+            PoseidonTree::new(
+                tree_height,
+                <PoseidonTree as ZerokitMerkleTree>::Hasher::default_leaf(),
+                tree_config,
+            )?
+        } else {
+            // We compute a default empty tree
+            PoseidonTree::default(tree_height)?
+        };
 
         Ok(RLN {
             #[cfg(not(target_arch = "wasm32"))]
diff --git a/rln/tests/ffi.rs b/rln/tests/ffi.rs
index 0e3dd828..469f9387 100644
--- a/rln/tests/ffi.rs
+++ b/rln/tests/ffi.rs
@@ -617,11 +617,14 @@ mod test {
 
         // Creating a RLN instance passing the raw data
         let mut rln_pointer_raw_bytes = MaybeUninit::<*mut RLN>::uninit();
+        let tree_config = "".to_string();
+        let tree_config_buffer = &Buffer::from(tree_config.as_bytes());
         let success = new_with_params(
             tree_height,
             circom_data,
             zkey_data,
             vk_data,
+            tree_config_buffer,
             rln_pointer_raw_bytes.as_mut_ptr(),
         );
         assert!(success, "RLN object creation failed");