feat: expose hash, poseidon_hash and delete_leaf

Co-authored-by: tyshkor <tyshko1@gmail.com>

.github/workflows/nightly-release.yml

@@ -135,6 +135,7 @@ jobs:
           gh release create nightly --prerelease --target master \
             --title 'Nightly build ("master" branch)' \
             --generate-notes \
+            --draft=false \
             --notes-start-tag $start_tag \
             *-archive/*.tar.gz \
         env:

.gitignore

@@ -8,6 +8,7 @@ rln/pmtree_db
 # will have compiled files and executables
 debug/
 target/
+wabt/
 
 # These are backup files generated by rustfmt
 **/*.rs.bk

rln-wasm/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rln-wasm"
-version = "0.0.7"
+version = "0.0.8"
 edition = "2021"
 license = "MIT or Apache2"
 

rln-wasm/src/lib.rs

@@ -3,9 +3,11 @@
 extern crate wasm_bindgen;
 extern crate web_sys;
 
+use std::vec::Vec;
+
 use js_sys::{BigInt as JsBigInt, Object, Uint8Array};
 use num_bigint::BigInt;
-use rln::public::RLN;
+use rln::public::{hash, poseidon_hash, RLN};
 use wasm_bindgen::prelude::*;
 
 #[wasm_bindgen]
@@ -20,6 +22,163 @@ pub struct RLNWrapper {
     instance: RLN<'static>,
 }
 
+// Macro to call methods with arbitrary amount of arguments,
+// which have the last argument is output buffer pointer
+// First argument to the macro is context,
+// second is the actual method on `RLN`
+// third is the aforementioned output buffer argument
+// rest are all other arguments to the method
+macro_rules! call_with_output_and_error_msg {
+    // this variant is needed for the case when
+    // there are zero other arguments
+    ($instance:expr, $method:ident, $error_msg:expr) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            let new_instance = $instance.process();
+            if let Err(err) = new_instance.instance.$method(&mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+    ($instance:expr, $method:ident, $error_msg:expr, $( $arg:expr ),* ) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            let new_instance = $instance.process();
+            if let Err(err) = new_instance.instance.$method($($arg.process()),*, &mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+}
+
+// Macro to call_with_error_msg methods with arbitrary amount of arguments,
+// First argument to the macro is context,
+// second is the actual method on `RLNWrapper`
+// rest are all other arguments to the method
+macro_rules! call_with_error_msg {
+    ($instance:expr, $method:ident, $error_msg:expr $(, $arg:expr)*) => {
+        {
+            let new_instance: &mut RLNWrapper = $instance.process();
+            if let Err(err) = new_instance.instance.$method($($arg.process()),*) {
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                Ok(())
+            }
+
+        }
+    }
+}
+
+macro_rules! call {
+    ($instance:expr, $method:ident $(, $arg:expr)*) => {
+        {
+            let new_instance: &mut RLNWrapper = $instance.process();
+            new_instance.instance.$method($($arg.process()),*)
+        }
+    }
+}
+
+macro_rules! call_bool_method_with_error_msg {
+    ($instance:expr, $method:ident, $error_msg:expr $(, $arg:expr)*) => {
+        {
+            let new_instance: &RLNWrapper = $instance.process();
+            new_instance.instance.$method($($arg.process()),*).map_err(|err| format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+        }
+    }
+}
+
+// Macro to execute a function with arbitrary amount of arguments,
+// First argument is the function to execute
+// Rest are all other arguments to the method
+macro_rules! fn_call_with_output_and_error_msg {
+    // this variant is needed for the case when
+    // there are zero other arguments
+    ($func:ident, $error_msg:expr) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            if let Err(err) = $func(&mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+    ($func:ident, $error_msg:expr, $( $arg:expr ),* ) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            if let Err(err) = $func($($arg.process()),*, &mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+}
+
+trait ProcessArg {
+    type ReturnType;
+    fn process(self) -> Self::ReturnType;
+}
+
+impl ProcessArg for usize {
+    type ReturnType = usize;
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
+impl<T> ProcessArg for Vec<T> {
+    type ReturnType = Vec<T>;
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
+impl<'a> ProcessArg for *const RLN<'a> {
+    type ReturnType = &'a RLN<'a>;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &*self }
+    }
+}
+
+impl ProcessArg for *const RLNWrapper {
+    type ReturnType = &'static RLNWrapper;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &*self }
+    }
+}
+
+impl ProcessArg for *mut RLNWrapper {
+    type ReturnType = &'static mut RLNWrapper;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &mut *self }
+    }
+}
+
+impl<'a> ProcessArg for &'a [u8] {
+    type ReturnType = &'a [u8];
+
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = newRLN)]
 pub fn wasm_new(
@@ -39,24 +198,20 @@ pub fn wasm_get_serialized_rln_witness(
     ctx: *mut RLNWrapper,
     input: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &mut *ctx };
-    let rln_witness = wrapper
-        .instance
-        .get_serialized_rln_witness(&input.to_vec()[..])
+    let rln_witness = call!(ctx, get_serialized_rln_witness, &input.to_vec()[..])
         .map_err(|err| format!("{:#?}", err))?;
-
     Ok(Uint8Array::from(&rln_witness[..]))
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = insertMember)]
 pub fn wasm_set_next_leaf(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper.instance.set_next_leaf(&input.to_vec()[..]).is_ok() {
-        Ok(())
-    } else {
-        Err("could not insert member into merkle tree".into())
-    }
+    call_with_error_msg!(
+        ctx,
+        set_next_leaf,
+        "could not insert member into merkle tree".to_string(),
+        &input.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -66,31 +221,30 @@ pub fn wasm_set_leaves_from(
     index: usize,
     input: Uint8Array,
 ) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper
-        .instance
-        .set_leaves_from(index as usize, &input.to_vec()[..])
-        .is_ok()
-    {
-        Ok(())
-    } else {
-        Err("could not set multiple leaves".into())
-    }
+    call_with_error_msg!(
+        ctx,
+        set_leaves_from,
+        "could not set multiple leaves".to_string(),
+        index,
+        &*input.to_vec()
+    )
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[wasm_bindgen(js_name = deleteLeaf)]
+pub fn wasm_delete_leaf(ctx: *mut RLNWrapper, index: usize) -> Result<(), String> {
+    call_with_error_msg!(ctx, delete_leaf, "could not delete leaf".to_string(), index)
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = initTreeWithLeaves)]
 pub fn wasm_init_tree_with_leaves(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper
-        .instance
-        .init_tree_with_leaves(&input.to_vec()[..])
-        .is_ok()
-    {
-        Ok(())
-    } else {
-        Err("could not init merkle tree".into())
-    }
+    call_with_error_msg!(
+        ctx,
+        init_tree_with_leaves,
+        "could not init merkle tree".to_string(),
+        &*input.to_vec()
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -99,12 +253,8 @@ pub fn rln_witness_to_json(
     ctx: *mut RLNWrapper,
     serialized_witness: Uint8Array,
 ) -> Result<Object, String> {
-    let wrapper = unsafe { &mut *ctx };
-    let inputs = wrapper
-        .instance
-        .get_rln_witness_json(&serialized_witness.to_vec()[..])
+    let inputs = call!(ctx, get_rln_witness_json, &serialized_witness.to_vec()[..])
         .map_err(|err| err.to_string())?;
-
     let js_value = serde_wasm_bindgen::to_value(&inputs).map_err(|err| err.to_string())?;
     Object::from_entries(&js_value).map_err(|err| format!("{:#?}", err))
 }
@@ -116,8 +266,6 @@ pub fn generate_rln_proof_with_witness(
     calculated_witness: Vec<JsBigInt>,
     serialized_witness: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &mut *ctx };
-
     let mut witness_vec: Vec<BigInt> = vec![];
 
     for v in calculated_witness {
@@ -131,69 +279,36 @@ pub fn generate_rln_proof_with_witness(
         );
     }
 
-    let mut output_data: Vec<u8> = Vec::new();
-
-    if wrapper
-        .instance
-        .generate_rln_proof_with_witness(witness_vec, serialized_witness.to_vec(), &mut output_data)
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate proof".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        generate_rln_proof_with_witness,
+        "could not generate proof",
+        witness_vec,
+        serialized_witness.to_vec()
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = generateMembershipKey)]
 pub fn wasm_key_gen(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper.instance.key_gen(&mut output_data).is_ok() {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership keys".into())
-    }
+    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> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper.instance.extended_key_gen(&mut output_data).is_ok() {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership keys".into())
-    }
+    call_with_output_and_error_msg!(ctx, extended_key_gen, "could not generate membership keys")
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = generateSeededMembershipKey)]
 pub fn wasm_seeded_key_gen(ctx: *const RLNWrapper, seed: Uint8Array) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper
-        .instance
-        .seeded_key_gen(&seed.to_vec()[..], &mut output_data)
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership key".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        seeded_key_gen,
+        "could not generate membership key",
+        &seed.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -202,20 +317,12 @@ pub fn wasm_seeded_extended_key_gen(
     ctx: *const RLNWrapper,
     seed: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper
-        .instance
-        .seeded_extended_key_gen(&seed.to_vec()[..], &mut output_data)
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership key".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        seeded_extended_key_gen,
+        "could not generate membership key",
+        &seed.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -225,38 +332,24 @@ pub fn wasm_recover_id_secret(
     input_proof_data_1: Uint8Array,
     input_proof_data_2: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper
-        .instance
-        .recover_id_secret(
-            &input_proof_data_1.to_vec()[..],
-            &input_proof_data_2.to_vec()[..],
-            &mut output_data,
-        )
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not recover id secret".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        recover_id_secret,
+        "could not recover id secret",
+        &input_proof_data_1.to_vec()[..],
+        &input_proof_data_2.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = verifyRLNProof)]
 pub fn wasm_verify_rln_proof(ctx: *const RLNWrapper, proof: Uint8Array) -> Result<bool, String> {
-    let wrapper = unsafe { &*ctx };
-    if match wrapper.instance.verify_rln_proof(&proof.to_vec()[..]) {
-        Ok(verified) => verified,
-        Err(_) => return Err("error while verifying rln proof".into()),
-    } {
-        return Ok(true);
-    }
-
-    Ok(false)
+    call_bool_method_with_error_msg!(
+        ctx,
+        verify_rln_proof,
+        "error while verifying rln proof".to_string(),
+        &proof.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -266,31 +359,31 @@ pub fn wasm_verify_with_roots(
     proof: Uint8Array,
     roots: Uint8Array,
 ) -> Result<bool, String> {
-    let wrapper = unsafe { &*ctx };
-    if match wrapper
-        .instance
-        .verify_with_roots(&proof.to_vec()[..], &roots.to_vec()[..])
-    {
-        Ok(verified) => verified,
-        Err(_) => return Err("error while verifying proof with roots".into()),
-    } {
-        return Ok(true);
-    }
-
-    Ok(false)
+    call_bool_method_with_error_msg!(
+        ctx,
+        verify_with_roots,
+        "error while verifying proof with roots".to_string(),
+        &proof.to_vec()[..],
+        &roots.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = getRoot)]
 pub fn wasm_get_root(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper.instance.get_root(&mut output_data).is_ok() {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not obtain root".into())
-    }
+    call_with_output_and_error_msg!(ctx, get_root, "could not obtain root")
+}
+
+#[wasm_bindgen(js_name = hash)]
+pub fn wasm_hash(input: Uint8Array) -> Result<Uint8Array, String> {
+    fn_call_with_output_and_error_msg!(hash, "could not generate hash", &input.to_vec()[..])
+}
+
+#[wasm_bindgen(js_name = poseidonHash)]
+pub fn wasm_poseidon_hash(input: Uint8Array) -> Result<Uint8Array, String> {
+    fn_call_with_output_and_error_msg!(
+        poseidon_hash,
+        "could not generate poseidon hash",
+        &input.to_vec()[..]
+    )
 }

rln/src/hashers.rs

@@ -1,7 +1,8 @@
 // This crate instantiate the Poseidon hash algorithm
 
-use crate::circuit::Fr;
+use crate::{circuit::Fr, utils::bytes_le_to_fr};
 use once_cell::sync::Lazy;
+use tiny_keccak::{Hasher, Keccak};
 use utils::poseidon::Poseidon;
 
 // These indexed constants hardcodes the supported round parameters tuples (t, RF, RN, SKIP_MATRICES) for the Bn254 scalar field
@@ -26,3 +27,17 @@ pub fn poseidon_hash(input: &[Fr]) -> Fr {
         .hash(input.to_vec())
         .expect("hash with fixed input size can't fail")
 }
+
+// Hashes arbitrary signal to the underlying prime field
+pub fn hash_to_field(signal: &[u8]) -> Fr {
+    // We hash the input signal using Keccak256
+    // (note that a bigger curve order might require a bigger hash blocksize)
+    let mut hash = [0; 32];
+    let mut hasher = Keccak::v256();
+    hasher.update(signal);
+    hasher.finalize(&mut hash);
+
+    // We export the hash as a field element
+    let (el, _) = bytes_le_to_fr(hash.as_ref());
+    el
+}

rln/src/lib.rs

@@ -1,7 +1,7 @@
 #![allow(dead_code)]
 
 pub mod circuit;
-pub mod poseidon_hash;
+pub mod hashers;
 pub mod poseidon_tree;
 pub mod protocol;
 pub mod public;

rln/src/poseidon_tree.rs

@@ -3,7 +3,7 @@
 // Implementation inspired by https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs (no differences)
 
 use crate::circuit::Fr;
-use crate::poseidon_hash::poseidon_hash;
+use crate::hashers::poseidon_hash;
 use cfg_if::cfg_if;
 use utils::merkle_tree::*;
 

rln/src/protocol.rs

@@ -17,11 +17,13 @@ use thiserror::Error;
 use tiny_keccak::{Hasher as _, Keccak};
 
 use crate::circuit::{Curve, Fr};
-use crate::poseidon_hash::poseidon_hash;
+use crate::hashers::hash_to_field;
+use crate::hashers::poseidon_hash;
 use crate::poseidon_tree::*;
 use crate::public::RLN_IDENTIFIER;
 use crate::utils::*;
 use cfg_if::cfg_if;
+use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 
 ///////////////////////////////////////////////////////
 // RLN Witness data structure and utility functions
@@ -482,20 +484,6 @@ pub fn extended_seeded_keygen(signal: &[u8]) -> (Fr, Fr, Fr, Fr) {
     )
 }
 
-// Hashes arbitrary signal to the underlying prime field
-pub fn hash_to_field(signal: &[u8]) -> Fr {
-    // We hash the input signal using Keccak256
-    // (note that a bigger curve order might require a bigger hash blocksize)
-    let mut hash = [0; 32];
-    let mut hasher = Keccak::v256();
-    hasher.update(signal);
-    hasher.finalize(&mut hash);
-
-    // We export the hash as a field element
-    let (el, _) = bytes_le_to_fr(hash.as_ref());
-    el
-}
-
 pub fn compute_id_secret(
     share1: (Fr, Fr),
     share2: (Fr, Fr),

rln/src/public.rs

@@ -1,5 +1,5 @@
 use crate::circuit::{vk_from_raw, zkey_from_raw, Curve, Fr};
-use crate::poseidon_hash::poseidon_hash as utils_poseidon_hash;
+use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash};
 use crate::poseidon_tree::PoseidonTree;
 use crate::protocol::*;
 use crate::utils::*;
@@ -13,6 +13,7 @@ use cfg_if::cfg_if;
 use color_eyre::Result;
 use num_bigint::BigInt;
 use std::io::Cursor;
+use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 // use rkyv::Deserialize;
 
 cfg_if! {
@@ -1059,6 +1060,7 @@ mod test {
     use super::*;
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
+    use utils::ZerokitMerkleTree;
     // use rkyv::Deserialize;
 
     #[test]

rln/tests/ffi.rs

@@ -4,7 +4,7 @@ mod test {
     use rand::Rng;
     use rln::circuit::*;
     use rln::ffi::{hash as ffi_hash, poseidon_hash as ffi_poseidon_hash, *};
-    use rln::poseidon_hash::{poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
+    use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
     use rln::protocol::*;
     use rln::public::RLN;
     use rln::utils::*;

rln/tests/poseidon_tree.rs

@@ -6,7 +6,7 @@
 mod test {
     use rln::circuit::*;
     use rln::poseidon_tree::*;
-    use utils::{FullMerkleTree, OptimalMerkleTree};
+    use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     #[test]
     /// A basic performance comparison between the two supported Merkle Tree implementations
@@ -88,7 +88,7 @@ mod pmtree_test {
 
     use pmtree::*;
     use rln::circuit::Fr;
-    use rln::poseidon_hash::poseidon_hash;
+    use rln::hashers::{hash_to_field, poseidon_hash};
     use rln::poseidon_tree::PoseidonHash;
     use rln::protocol::hash_to_field;
     use rln::utils::str_to_fr;

rln/tests/protocol.rs

@@ -4,10 +4,11 @@ mod test {
         circom_from_folder, vk_from_folder, zkey_from_folder, Fr, TEST_RESOURCES_FOLDER,
         TEST_TREE_HEIGHT,
     };
-    use rln::poseidon_hash::poseidon_hash;
+    use rln::hashers::{hash_to_field, poseidon_hash};
     use rln::poseidon_tree::PoseidonTree;
     use rln::protocol::*;
     use rln::utils::str_to_fr;
+    use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 
     // Input generated with https://github.com/oskarth/zk-kit/commit/b6a872f7160c7c14e10a0ea40acab99cbb23c9a8
     const WITNESS_JSON_15: &str = r#"

rln/tests/public.rs

@@ -3,8 +3,8 @@ mod test {
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
     use rln::circuit::{Fr, TEST_RESOURCES_FOLDER, TEST_TREE_HEIGHT};
-    use rln::poseidon_hash::{poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
-    use rln::protocol::{compute_tree_root, deserialize_identity_tuple, hash_to_field};
+    use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
+    use rln::protocol::{compute_tree_root, deserialize_identity_tuple};
     use rln::public::{hash as public_hash, poseidon_hash as public_poseidon_hash, RLN};
     use rln::utils::*;
     use std::io::Cursor;

utils/src/merkle_tree/full_merkle_tree.rs

@@ -0,0 +1,289 @@
+use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
+use color_eyre::{Report, Result};
+use std::{
+    cmp::max,
+    fmt::Debug,
+    iter::{once, repeat, successors},
+};
+
+////////////////////////////////////////////////////////////
+/// Full Merkle Tree Implementation
+////////////////////////////////////////////////////////////
+
+/// Merkle tree with all leaf and intermediate hashes stored
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct FullMerkleTree<H: Hasher> {
+    /// The depth of the tree, i.e. the number of levels from leaf to root
+    depth: usize,
+
+    /// The nodes cached from the empty part of the tree (where leaves are set to default).
+    /// Since the rightmost part of the tree is usually changed much later than its creation,
+    /// we can prove accumulation of elements in the leftmost part, with no need to initialize the full tree
+    /// and by caching few intermediate nodes to the root computed from default leaves
+    cached_nodes: Vec<H::Fr>,
+
+    /// The tree nodes
+    nodes: Vec<H::Fr>,
+
+    // The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
+    // (deletions leave next_index unchanged)
+    next_index: usize,
+}
+
+/// Element of a Merkle proof
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum FullMerkleBranch<H: Hasher> {
+    /// Left branch taken, value is the right sibling hash.
+    Left(H::Fr),
+
+    /// Right branch taken, value is the left sibling hash.
+    Right(H::Fr),
+}
+
+/// Merkle proof path, bottom to top.
+#[derive(Clone, PartialEq, Eq)]
+pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
+
+/// Implementations
+
+impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
+    type Proof = FullMerkleProof<H>;
+    fn default(depth: usize) -> Self {
+        FullMerkleTree::<H>::new(depth, H::default_leaf())
+    }
+
+    /// Creates a new `MerkleTree`
+    /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
+    fn new(depth: usize, initial_leaf: H::Fr) -> Self {
+        // Compute cache node values, leaf to root
+        let cached_nodes = successors(Some(initial_leaf), |prev| Some(H::hash(&[*prev, *prev])))
+            .take(depth + 1)
+            .collect::<Vec<_>>();
+
+        // Compute node values
+        let nodes = cached_nodes
+            .iter()
+            .rev()
+            .enumerate()
+            .flat_map(|(levels, hash)| repeat(hash).take(1 << levels))
+            .cloned()
+            .collect::<Vec<_>>();
+        debug_assert!(nodes.len() == (1 << (depth + 1)) - 1);
+
+        let next_index = 0;
+
+        Self {
+            depth,
+            cached_nodes,
+            nodes,
+            next_index,
+        }
+    }
+
+    // Returns the depth of the tree
+    fn depth(&self) -> usize {
+        self.depth
+    }
+
+    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
+    fn capacity(&self) -> usize {
+        1 << self.depth
+    }
+
+    // Returns the total number of leaves set
+    fn leaves_set(&mut self) -> usize {
+        self.next_index
+    }
+
+    #[must_use]
+    // Returns the root of the tree
+    fn root(&self) -> H::Fr {
+        self.nodes[0]
+    }
+
+    // Sets a leaf at the specified tree index
+    fn set(&mut self, leaf: usize, hash: H::Fr) -> Result<()> {
+        self.set_range(leaf, once(hash))?;
+        self.next_index = max(self.next_index, leaf + 1);
+        Ok(())
+    }
+
+    // Sets tree nodes, starting from start index
+    // Function proper of FullMerkleTree implementation
+    fn set_range<I: IntoIterator<Item = H::Fr>>(&mut self, start: usize, hashes: I) -> Result<()> {
+        let index = self.capacity() + start - 1;
+        let mut count = 0;
+        // first count number of hashes, and check that they fit in the tree
+        // then insert into the tree
+        let hashes = hashes.into_iter().collect::<Vec<_>>();
+        if hashes.len() + start > self.capacity() {
+            return Err(Report::msg("provided hashes do not fit in the tree"));
+        }
+        hashes.into_iter().for_each(|hash| {
+            self.nodes[index + count] = hash;
+            count += 1;
+        });
+        if count != 0 {
+            self.update_nodes(index, index + (count - 1))?;
+            self.next_index = max(self.next_index, start + count);
+        }
+        Ok(())
+    }
+
+    // Sets a leaf at the next available index
+    fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
+        self.set(self.next_index, leaf)?;
+        Ok(())
+    }
+
+    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
+    fn delete(&mut self, index: usize) -> Result<()> {
+        // We reset the leaf only if we previously set a leaf at that index
+        if index < self.next_index {
+            self.set(index, H::default_leaf())?;
+        }
+        Ok(())
+    }
+
+    // Computes a merkle proof the the leaf at the specified index
+    fn proof(&self, leaf: usize) -> Result<FullMerkleProof<H>> {
+        if leaf >= self.capacity() {
+            return Err(Report::msg("index exceeds set size"));
+        }
+        let mut index = self.capacity() + leaf - 1;
+        let mut path = Vec::with_capacity(self.depth + 1);
+        while let Some(parent) = self.parent(index) {
+            // Add proof for node at index to parent
+            path.push(match index & 1 {
+                1 => FullMerkleBranch::Left(self.nodes[index + 1]),
+                0 => FullMerkleBranch::Right(self.nodes[index - 1]),
+                _ => unreachable!(),
+            });
+            index = parent;
+        }
+        Ok(FullMerkleProof(path))
+    }
+
+    // Verifies a Merkle proof with respect to the input leaf and the tree root
+    fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof<H>) -> Result<bool> {
+        Ok(proof.compute_root_from(hash) == self.root())
+    }
+}
+
+impl<H: Hasher> FullMerkleTree<H>
+where
+    H: Hasher,
+{
+    // Utilities for updating the tree nodes
+
+    /// For a given node index, return the parent node index
+    /// Returns None if there is no parent (root node)
+    fn parent(&self, index: usize) -> Option<usize> {
+        if index == 0 {
+            None
+        } else {
+            Some(((index + 1) >> 1) - 1)
+        }
+    }
+
+    /// For a given node index, return index of the first (left) child.
+    fn first_child(&self, index: usize) -> usize {
+        (index << 1) + 1
+    }
+
+    fn levels(&self, index: usize) -> usize {
+        // `n.next_power_of_two()` will return `n` iff `n` is a power of two.
+        // The extra offset corrects this.
+        (index + 2).next_power_of_two().trailing_zeros() as usize - 1
+    }
+
+    fn update_nodes(&mut self, start: usize, end: usize) -> Result<()> {
+        if self.levels(start) != self.levels(end) {
+            return Err(Report::msg("self.levels(start) != self.levels(end)"));
+        }
+        if let (Some(start), Some(end)) = (self.parent(start), self.parent(end)) {
+            for parent in start..=end {
+                let child = self.first_child(parent);
+                self.nodes[parent] = H::hash(&[self.nodes[child], self.nodes[child + 1]]);
+            }
+            self.update_nodes(start, end)?;
+        }
+        Ok(())
+    }
+}
+
+impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
+    type Index = u8;
+
+    #[must_use]
+    // Returns the length of a Merkle proof
+    fn length(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Computes the leaf index corresponding to a Merkle proof
+    #[must_use]
+    fn leaf_index(&self) -> usize {
+        self.0.iter().rev().fold(0, |index, branch| match branch {
+            FullMerkleBranch::Left(_) => index << 1,
+            FullMerkleBranch::Right(_) => (index << 1) + 1,
+        })
+    }
+
+    #[must_use]
+    /// Returns the path elements forming a Merkle proof
+    fn get_path_elements(&self) -> Vec<H::Fr> {
+        self.0
+            .iter()
+            .map(|x| match x {
+                FullMerkleBranch::Left(value) | FullMerkleBranch::Right(value) => *value,
+            })
+            .collect()
+    }
+
+    /// Returns the path indexes forming a Merkle proof
+    #[must_use]
+    fn get_path_index(&self) -> Vec<Self::Index> {
+        self.0
+            .iter()
+            .map(|branch| match branch {
+                FullMerkleBranch::Left(_) => 0,
+                FullMerkleBranch::Right(_) => 1,
+            })
+            .collect()
+    }
+
+    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
+    #[must_use]
+    fn compute_root_from(&self, hash: &H::Fr) -> H::Fr {
+        self.0.iter().fold(*hash, |hash, branch| match branch {
+            FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]),
+            FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]),
+        })
+    }
+}
+
+// Debug formatting for printing a (Full) Merkle Proof Branch
+impl<H> Debug for FullMerkleBranch<H>
+where
+    H: Hasher,
+    H::Fr: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Left(arg0) => f.debug_tuple("Left").field(arg0).finish(),
+            Self::Right(arg0) => f.debug_tuple("Right").field(arg0).finish(),
+        }
+    }
+}
+
+// Debug formatting for printing a (Full) Merkle Proof
+impl<H> Debug for FullMerkleProof<H>
+where
+    H: Hasher,
+    H::Fr: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Proof").field(&self.0).finish()
+    }
+}

utils/src/merkle_tree/merkle_tree.rs

@@ -13,22 +13,13 @@
 //! * Disk based storage backend (using mmaped files should be easy)
 //! * Implement serialization for tree and Merkle proof
 
-#![allow(dead_code)]
-
-use std::collections::HashMap;
-use std::{
-    cmp::max,
-    fmt::Debug,
-    iter::{once, repeat, successors},
-};
-
-use color_eyre::{Report, Result};
+use color_eyre::Result;
 
 /// In the Hasher trait we define the node type, the default leaf
 /// and the hash function used to initialize a Merkle Tree implementation
 pub trait Hasher {
     /// Type of the leaf and tree node
-    type Fr: Copy + Clone + Eq;
+    type Fr: Clone + Copy + Eq;
 
     /// Returns the default tree leaf
     fn default_leaf() -> Self::Fr;
@@ -37,528 +28,39 @@ pub trait Hasher {
     fn hash(input: &[Self::Fr]) -> Self::Fr;
 }
 
-////////////////////////////////////////////////////////////
-/// Optimal Merkle Tree Implementation
-////////////////////////////////////////////////////////////
-
-/// The Merkle tree structure
-#[derive(Clone, PartialEq, Eq, Debug)]
-pub struct OptimalMerkleTree<H>
+/// In the ZerokitMerkleTree trait we define the methods that are required to be implemented by a Merkle tree
+/// Including, OptimalMerkleTree, FullMerkleTree, Pmtree
+pub trait ZerokitMerkleTree<H: Hasher>
 where
     H: Hasher,
 {
-    /// The depth of the tree, i.e. the number of levels from leaf to root
-    depth: usize,
+    type Proof;
 
-    /// The nodes cached from the empty part of the tree (where leaves are set to default).
-    /// Since the rightmost part of the tree is usually changed much later than its creation,
-    /// we can prove accumulation of elements in the leftmost part, with no need to initialize the full tree
-    /// and by caching few intermediate nodes to the root computed from default leaves
-    cached_nodes: Vec<H::Fr>,
-
-    /// The tree nodes
-    nodes: HashMap<(usize, usize), H::Fr>,
-
-    // The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
-    // (deletions leave next_index unchanged)
-    next_index: usize,
+    fn default(depth: usize) -> Self;
+    fn new(depth: usize, default_leaf: H::Fr) -> Self;
+    fn depth(&self) -> usize;
+    fn capacity(&self) -> usize;
+    fn leaves_set(&mut self) -> usize;
+    fn root(&self) -> H::Fr;
+    fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()>;
+    fn set_range<I>(&mut self, start: usize, leaves: I) -> Result<()>
+    where
+        I: IntoIterator<Item = H::Fr>;
+    fn update_next(&mut self, leaf: H::Fr) -> Result<()>;
+    fn delete(&mut self, index: usize) -> Result<()>;
+    fn proof(&self, index: usize) -> Result<Self::Proof>;
+    fn verify(&self, leaf: &H::Fr, witness: &Self::Proof) -> Result<bool>;
 }
 
-/// The Merkle proof
-/// Contains a vector of (node, branch_index) that defines the proof path elements and branch direction (1 or 0)
-#[derive(Clone, PartialEq, Eq)]
-pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
-
-/// Implementations
-
-impl<H: Hasher> OptimalMerkleTree<H> {
-    pub fn default(depth: usize) -> Self {
-        OptimalMerkleTree::<H>::new(depth, H::default_leaf())
-    }
-
-    /// Creates a new `MerkleTree`
-    /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
-    pub fn new(depth: usize, default_leaf: H::Fr) -> Self {
-        let mut cached_nodes: Vec<H::Fr> = Vec::with_capacity(depth + 1);
-        cached_nodes.push(default_leaf);
-        for i in 0..depth {
-            cached_nodes.push(H::hash(&[cached_nodes[i]; 2]));
-        }
-        cached_nodes.reverse();
-        OptimalMerkleTree {
-            cached_nodes: cached_nodes.clone(),
-            depth,
-            nodes: HashMap::new(),
-            next_index: 0,
-        }
-    }
-
-    // Returns the depth of the tree
-    pub fn depth(&self) -> usize {
-        self.depth
-    }
-
-    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
-    pub fn capacity(&self) -> usize {
-        1 << self.depth
-    }
-
-    // Returns the total number of leaves set
-    pub fn leaves_set(&mut self) -> usize {
-        self.next_index
-    }
-
-    #[must_use]
-    // Returns the root of the tree
-    pub fn root(&self) -> H::Fr {
-        self.get_node(0, 0)
-    }
-
-    // Sets a leaf at the specified tree index
-    pub fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()> {
-        if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        self.nodes.insert((self.depth, index), leaf);
-        self.recalculate_from(index)?;
-        self.next_index = max(self.next_index, index + 1);
-        Ok(())
-    }
-
-    // Sets multiple leaves from the specified tree index
-    pub fn set_range<I: IntoIterator<Item = H::Fr>>(
-        &mut self,
-        start: usize,
-        leaves: I,
-    ) -> Result<()> {
-        let leaves = leaves.into_iter().collect::<Vec<_>>();
-        // check if the range is valid
-        if start + leaves.len() > self.capacity() {
-            return Err(Report::msg("provided range exceeds set size"));
-        }
-        for (i, leaf) in leaves.iter().enumerate() {
-            self.nodes.insert((self.depth, start + i), *leaf);
-            self.recalculate_from(start + i)?;
-        }
-        self.next_index = max(self.next_index, start + leaves.len());
-        Ok(())
-    }
-
-    // Sets a leaf at the next available index
-    pub fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
-        self.set(self.next_index, leaf)?;
-        Ok(())
-    }
-
-    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
-    pub fn delete(&mut self, index: usize) -> Result<()> {
-        // We reset the leaf only if we previously set a leaf at that index
-        if index < self.next_index {
-            self.set(index, H::default_leaf())?;
-        }
-        Ok(())
-    }
-
-    // Computes a merkle proof the the leaf at the specified index
-    pub fn proof(&self, index: usize) -> Result<OptimalMerkleProof<H>> {
-        if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        let mut witness = Vec::<(H::Fr, u8)>::with_capacity(self.depth);
-        let mut i = index;
-        let mut depth = self.depth;
-        loop {
-            i ^= 1;
-            witness.push((self.get_node(depth, i), (1 - (i & 1)).try_into().unwrap()));
-            i >>= 1;
-            depth -= 1;
-            if depth == 0 {
-                break;
-            }
-        }
-        if i != 0 {
-            Err(Report::msg("i != 0"))
-        } else {
-            Ok(OptimalMerkleProof(witness))
-        }
-    }
-
-    // Verifies a Merkle proof with respect to the input leaf and the tree root
-    pub fn verify(&self, leaf: &H::Fr, witness: &OptimalMerkleProof<H>) -> Result<bool> {
-        if witness.length() != self.depth {
-            return Err(Report::msg("witness length doesn't match tree depth"));
-        }
-        let expected_root = witness.compute_root_from(leaf);
-        Ok(expected_root.eq(&self.root()))
-    }
-
-    // Utilities for updating the tree nodes
-
-    fn get_node(&self, depth: usize, index: usize) -> H::Fr {
-        let node = *self
-            .nodes
-            .get(&(depth, index))
-            .unwrap_or_else(|| &self.cached_nodes[depth]);
-        node
-    }
-
-    fn get_leaf(&self, index: usize) -> H::Fr {
-        self.get_node(self.depth, index)
-    }
-
-    fn hash_couple(&mut self, depth: usize, index: usize) -> H::Fr {
-        let b = index & !1;
-        H::hash(&[self.get_node(depth, b), self.get_node(depth, b + 1)])
-    }
-
-    fn recalculate_from(&mut self, index: usize) -> Result<()> {
-        let mut i = index;
-        let mut depth = self.depth;
-        loop {
-            let h = self.hash_couple(depth, i);
-            i >>= 1;
-            depth -= 1;
-            self.nodes.insert((depth, i), h);
-            if depth == 0 {
-                break;
-            }
-        }
-        if depth != 0 {
-            return Err(Report::msg("did not reach the depth"));
-        }
-        if i != 0 {
-            return Err(Report::msg("did not go through all indexes"));
-        }
-        Ok(())
-    }
-}
-
-impl<H: Hasher> OptimalMerkleProof<H> {
-    #[must_use]
-    // Returns the length of a Merkle proof
-    pub fn length(&self) -> usize {
-        self.0.len()
-    }
-
-    /// Computes the leaf index corresponding to a Merkle proof
-    #[must_use]
-    pub fn leaf_index(&self) -> usize {
-        // In current implementation the path indexes in a proof correspond to the binary representation of the leaf index
-        let mut binary_repr = self.get_path_index();
-        binary_repr.reverse();
-        binary_repr
-            .into_iter()
-            .fold(0, |acc, digit| (acc << 1) + usize::from(digit))
-    }
-
-    #[must_use]
-    /// Returns the path elements forming a Merkle proof
-    pub fn get_path_elements(&self) -> Vec<H::Fr> {
-        self.0.iter().map(|x| x.0).collect()
-    }
-
-    /// Returns the path indexes forming a Merkle proof
-    #[must_use]
-    pub fn get_path_index(&self) -> Vec<u8> {
-        self.0.iter().map(|x| x.1).collect()
-    }
-
-    #[must_use]
-    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
-    pub fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr {
-        let mut acc: H::Fr = *leaf;
-        for w in self.0.iter() {
-            if w.1 == 0 {
-                acc = H::hash(&[acc, w.0]);
-            } else {
-                acc = H::hash(&[w.0, acc]);
-            }
-        }
-        acc
-    }
-}
-
-// Debug formatting for printing a (Optimal) Merkle Proof
-impl<H> Debug for OptimalMerkleProof<H>
+pub trait ZerokitMerkleProof<H: Hasher>
 where
     H: Hasher,
-    H::Fr: Debug,
 {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_tuple("Proof").field(&self.0).finish()
-    }
-}
-
-////////////////////////////////////////////////////////////
-/// Full Merkle Tree Implementation
-////////////////////////////////////////////////////////////
-
-/// Merkle tree with all leaf and intermediate hashes stored
-#[derive(Clone, PartialEq, Eq, Debug)]
-pub struct FullMerkleTree<H: Hasher> {
-    /// The depth of the tree, i.e. the number of levels from leaf to root
-    depth: usize,
-
-    /// The nodes cached from the empty part of the tree (where leaves are set to default).
-    /// Since the rightmost part of the tree is usually changed much later than its creation,
-    /// we can prove accumulation of elements in the leftmost part, with no need to initialize the full tree
-    /// and by caching few intermediate nodes to the root computed from default leaves
-    cached_nodes: Vec<H::Fr>,
-
-    /// The tree nodes
-    nodes: Vec<H::Fr>,
-
-    // The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
-    // (deletions leave next_index unchanged)
-    next_index: usize,
-}
-
-/// Element of a Merkle proof
-#[derive(Clone, Copy, PartialEq, Eq)]
-pub enum FullMerkleBranch<H: Hasher> {
-    /// Left branch taken, value is the right sibling hash.
-    Left(H::Fr),
-
-    /// Right branch taken, value is the left sibling hash.
-    Right(H::Fr),
-}
-
-/// Merkle proof path, bottom to top.
-#[derive(Clone, PartialEq, Eq)]
-pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
-
-/// Implementations
-
-impl<H: Hasher> FullMerkleTree<H> {
-    pub fn default(depth: usize) -> Self {
-        FullMerkleTree::<H>::new(depth, H::default_leaf())
-    }
-
-    /// Creates a new `MerkleTree`
-    /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
-    pub fn new(depth: usize, initial_leaf: H::Fr) -> Self {
-        // Compute cache node values, leaf to root
-        let cached_nodes = successors(Some(initial_leaf), |prev| Some(H::hash(&[*prev, *prev])))
-            .take(depth + 1)
-            .collect::<Vec<_>>();
-
-        // Compute node values
-        let nodes = cached_nodes
-            .iter()
-            .rev()
-            .enumerate()
-            .flat_map(|(levels, hash)| repeat(hash).take(1 << levels))
-            .cloned()
-            .collect::<Vec<_>>();
-        debug_assert!(nodes.len() == (1 << (depth + 1)) - 1);
-
-        let next_index = 0;
-
-        Self {
-            depth,
-            cached_nodes,
-            nodes,
-            next_index,
-        }
-    }
-
-    // Returns the depth of the tree
-    pub fn depth(&self) -> usize {
-        self.depth
-    }
-
-    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
-    pub fn capacity(&self) -> usize {
-        1 << self.depth
-    }
-
-    // Returns the total number of leaves set
-    pub fn leaves_set(&mut self) -> usize {
-        self.next_index
-    }
-
-    #[must_use]
-    // Returns the root of the tree
-    pub fn root(&self) -> H::Fr {
-        self.nodes[0]
-    }
-
-    // Sets a leaf at the specified tree index
-    pub fn set(&mut self, leaf: usize, hash: H::Fr) -> Result<()> {
-        self.set_range(leaf, once(hash))?;
-        self.next_index = max(self.next_index, leaf + 1);
-        Ok(())
-    }
-
-    // Sets tree nodes, starting from start index
-    // Function proper of FullMerkleTree implementation
-    fn set_range<I: IntoIterator<Item = H::Fr>>(&mut self, start: usize, hashes: I) -> Result<()> {
-        let index = self.capacity() + start - 1;
-        let mut count = 0;
-        // first count number of hashes, and check that they fit in the tree
-        // then insert into the tree
-        let hashes = hashes.into_iter().collect::<Vec<_>>();
-        if hashes.len() + start > self.capacity() {
-            return Err(Report::msg("provided hashes do not fit in the tree"));
-        }
-        hashes.into_iter().for_each(|hash| {
-            self.nodes[index + count] = hash;
-            count += 1;
-        });
-        if count != 0 {
-            self.update_nodes(index, index + (count - 1))?;
-            self.next_index = max(self.next_index, start + count);
-        }
-        Ok(())
-    }
-
-    // Sets a leaf at the next available index
-    pub fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
-        self.set(self.next_index, leaf)?;
-        Ok(())
-    }
-
-    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
-    pub fn delete(&mut self, index: usize) -> Result<()> {
-        // We reset the leaf only if we previously set a leaf at that index
-        if index < self.next_index {
-            self.set(index, H::default_leaf())?;
-        }
-        Ok(())
-    }
-
-    // Computes a merkle proof the the leaf at the specified index
-    pub fn proof(&self, leaf: usize) -> Result<FullMerkleProof<H>> {
-        if leaf >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        let mut index = self.capacity() + leaf - 1;
-        let mut path = Vec::with_capacity(self.depth + 1);
-        while let Some(parent) = self.parent(index) {
-            // Add proof for node at index to parent
-            path.push(match index & 1 {
-                1 => FullMerkleBranch::Left(self.nodes[index + 1]),
-                0 => FullMerkleBranch::Right(self.nodes[index - 1]),
-                _ => unreachable!(),
-            });
-            index = parent;
-        }
-        Ok(FullMerkleProof(path))
-    }
-
-    // Verifies a Merkle proof with respect to the input leaf and the tree root
-    pub fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof<H>) -> Result<bool> {
-        Ok(proof.compute_root_from(hash) == self.root())
-    }
-
-    // Utilities for updating the tree nodes
-
-    /// For a given node index, return the parent node index
-    /// Returns None if there is no parent (root node)
-    fn parent(&self, index: usize) -> Option<usize> {
-        if index == 0 {
-            None
-        } else {
-            Some(((index + 1) >> 1) - 1)
-        }
-    }
-
-    /// For a given node index, return index of the first (left) child.
-    fn first_child(&self, index: usize) -> usize {
-        (index << 1) + 1
-    }
-
-    fn levels(&self, index: usize) -> usize {
-        // `n.next_power_of_two()` will return `n` iff `n` is a power of two.
-        // The extra offset corrects this.
-        (index + 2).next_power_of_two().trailing_zeros() as usize - 1
-    }
-
-    fn update_nodes(&mut self, start: usize, end: usize) -> Result<()> {
-        if self.levels(start) != self.levels(end) {
-            return Err(Report::msg("self.levels(start) != self.levels(end)"));
-        }
-        if let (Some(start), Some(end)) = (self.parent(start), self.parent(end)) {
-            for parent in start..=end {
-                let child = self.first_child(parent);
-                self.nodes[parent] = H::hash(&[self.nodes[child], self.nodes[child + 1]]);
-            }
-            self.update_nodes(start, end)?;
-        }
-        Ok(())
-    }
-}
-
-impl<H: Hasher> FullMerkleProof<H> {
-    #[must_use]
-    // Returns the length of a Merkle proof
-    pub fn length(&self) -> usize {
-        self.0.len()
-    }
-
-    /// Computes the leaf index corresponding to a Merkle proof
-    #[must_use]
-    pub fn leaf_index(&self) -> usize {
-        self.0.iter().rev().fold(0, |index, branch| match branch {
-            FullMerkleBranch::Left(_) => index << 1,
-            FullMerkleBranch::Right(_) => (index << 1) + 1,
-        })
-    }
-
-    #[must_use]
-    /// Returns the path elements forming a Merkle proof
-    pub fn get_path_elements(&self) -> Vec<H::Fr> {
-        self.0
-            .iter()
-            .map(|x| match x {
-                FullMerkleBranch::Left(value) | FullMerkleBranch::Right(value) => *value,
-            })
-            .collect()
-    }
-
-    /// Returns the path indexes forming a Merkle proof
-    #[must_use]
-    pub fn get_path_index(&self) -> Vec<u8> {
-        self.0
-            .iter()
-            .map(|branch| match branch {
-                FullMerkleBranch::Left(_) => 0,
-                FullMerkleBranch::Right(_) => 1,
-            })
-            .collect()
-    }
-
-    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
-    #[must_use]
-    pub fn compute_root_from(&self, hash: &H::Fr) -> H::Fr {
-        self.0.iter().fold(*hash, |hash, branch| match branch {
-            FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]),
-            FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]),
-        })
-    }
-}
-
-// Debug formatting for printing a (Full) Merkle Proof Branch
-impl<H> Debug for FullMerkleBranch<H>
-where
-    H: Hasher,
-    H::Fr: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Left(arg0) => f.debug_tuple("Left").field(arg0).finish(),
-            Self::Right(arg0) => f.debug_tuple("Right").field(arg0).finish(),
-        }
-    }
-}
-
-// Debug formatting for printing a (Full) Merkle Proof
-impl<H> Debug for FullMerkleProof<H>
-where
-    H: Hasher,
-    H::Fr: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_tuple("Proof").field(&self.0).finish()
-    }
+    type Index;
+
+    fn length(&self) -> usize;
+    fn leaf_index(&self) -> usize;
+    fn get_path_elements(&self) -> Vec<H::Fr>;
+    fn get_path_index(&self) -> Vec<Self::Index>;
+    fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr;
 }

utils/src/merkle_tree/mod.rs

@@ -1,3 +1,7 @@
+pub mod full_merkle_tree;
 #[allow(clippy::module_inception)]
 pub mod merkle_tree;
+pub mod optimal_merkle_tree;
+pub use self::full_merkle_tree::*;
 pub use self::merkle_tree::*;
+pub use self::optimal_merkle_tree::*;

utils/src/merkle_tree/optimal_merkle_tree.rs

@@ -0,0 +1,266 @@
+use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
+use color_eyre::{Report, Result};
+use std::collections::HashMap;
+use std::{cmp::max, fmt::Debug};
+
+////////////////////////////////////////////////////////////
+/// Optimal Merkle Tree Implementation
+////////////////////////////////////////////////////////////
+
+/// The Merkle tree structure
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct OptimalMerkleTree<H>
+where
+    H: Hasher,
+{
+    /// The depth of the tree, i.e. the number of levels from leaf to root
+    depth: usize,
+
+    /// The nodes cached from the empty part of the tree (where leaves are set to default).
+    /// Since the rightmost part of the tree is usually changed much later than its creation,
+    /// we can prove accumulation of elements in the leftmost part, with no need to initialize the full tree
+    /// and by caching few intermediate nodes to the root computed from default leaves
+    cached_nodes: Vec<H::Fr>,
+
+    /// The tree nodes
+    nodes: HashMap<(usize, usize), H::Fr>,
+
+    // The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
+    // (deletions leave next_index unchanged)
+    next_index: usize,
+}
+
+/// The Merkle proof
+/// Contains a vector of (node, branch_index) that defines the proof path elements and branch direction (1 or 0)
+#[derive(Clone, PartialEq, Eq)]
+pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
+
+/// Implementations
+
+impl<H: Hasher> ZerokitMerkleTree<H> for OptimalMerkleTree<H>
+where
+    H: Hasher,
+{
+    type Proof = OptimalMerkleProof<H>;
+
+    fn default(depth: usize) -> Self {
+        OptimalMerkleTree::<H>::new(depth, H::default_leaf())
+    }
+
+    /// Creates a new `MerkleTree`
+    /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
+    fn new(depth: usize, default_leaf: H::Fr) -> Self {
+        let mut cached_nodes: Vec<H::Fr> = Vec::with_capacity(depth + 1);
+        cached_nodes.push(default_leaf);
+        for i in 0..depth {
+            cached_nodes.push(H::hash(&[cached_nodes[i]; 2]));
+        }
+        cached_nodes.reverse();
+        OptimalMerkleTree {
+            cached_nodes: cached_nodes.clone(),
+            depth,
+            nodes: HashMap::new(),
+            next_index: 0,
+        }
+    }
+
+    // Returns the depth of the tree
+    fn depth(&self) -> usize {
+        self.depth
+    }
+
+    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
+    fn capacity(&self) -> usize {
+        1 << self.depth
+    }
+
+    // Returns the total number of leaves set
+    fn leaves_set(&mut self) -> usize {
+        self.next_index
+    }
+
+    #[must_use]
+    // Returns the root of the tree
+    fn root(&self) -> H::Fr {
+        self.get_node(0, 0)
+    }
+
+    // Sets a leaf at the specified tree index
+    fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()> {
+        if index >= self.capacity() {
+            return Err(Report::msg("index exceeds set size"));
+        }
+        self.nodes.insert((self.depth, index), leaf);
+        self.recalculate_from(index)?;
+        self.next_index = max(self.next_index, index + 1);
+        Ok(())
+    }
+
+    // Sets multiple leaves from the specified tree index
+    fn set_range<I: IntoIterator<Item = H::Fr>>(&mut self, start: usize, leaves: I) -> Result<()> {
+        let leaves = leaves.into_iter().collect::<Vec<_>>();
+        // check if the range is valid
+        if start + leaves.len() > self.capacity() {
+            return Err(Report::msg("provided range exceeds set size"));
+        }
+        for (i, leaf) in leaves.iter().enumerate() {
+            self.nodes.insert((self.depth, start + i), *leaf);
+            self.recalculate_from(start + i)?;
+        }
+        self.next_index = max(self.next_index, start + leaves.len());
+        Ok(())
+    }
+
+    // Sets a leaf at the next available index
+    fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
+        self.set(self.next_index, leaf)?;
+        Ok(())
+    }
+
+    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
+    fn delete(&mut self, index: usize) -> Result<()> {
+        // We reset the leaf only if we previously set a leaf at that index
+        if index < self.next_index {
+            self.set(index, H::default_leaf())?;
+        }
+        Ok(())
+    }
+
+    // Computes a merkle proof the the leaf at the specified index
+    fn proof(&self, index: usize) -> Result<Self::Proof> {
+        if index >= self.capacity() {
+            return Err(Report::msg("index exceeds set size"));
+        }
+        let mut witness = Vec::<(H::Fr, u8)>::with_capacity(self.depth);
+        let mut i = index;
+        let mut depth = self.depth;
+        loop {
+            i ^= 1;
+            witness.push((self.get_node(depth, i), (1 - (i & 1)).try_into().unwrap()));
+            i >>= 1;
+            depth -= 1;
+            if depth == 0 {
+                break;
+            }
+        }
+        if i != 0 {
+            Err(Report::msg("i != 0"))
+        } else {
+            Ok(OptimalMerkleProof(witness))
+        }
+    }
+
+    // Verifies a Merkle proof with respect to the input leaf and the tree root
+    fn verify(&self, leaf: &H::Fr, witness: &Self::Proof) -> Result<bool> {
+        if witness.length() != self.depth {
+            return Err(Report::msg("witness length doesn't match tree depth"));
+        }
+        let expected_root = witness.compute_root_from(leaf);
+        Ok(expected_root.eq(&self.root()))
+    }
+}
+
+impl<H: Hasher> OptimalMerkleTree<H>
+where
+    H: Hasher,
+{
+    // Utilities for updating the tree nodes
+
+    fn get_node(&self, depth: usize, index: usize) -> H::Fr {
+        let node = *self
+            .nodes
+            .get(&(depth, index))
+            .unwrap_or_else(|| &self.cached_nodes[depth]);
+        node
+    }
+
+    pub fn get_leaf(&self, index: usize) -> H::Fr {
+        self.get_node(self.depth, index)
+    }
+
+    fn hash_couple(&mut self, depth: usize, index: usize) -> H::Fr {
+        let b = index & !1;
+        H::hash(&[self.get_node(depth, b), self.get_node(depth, b + 1)])
+    }
+
+    fn recalculate_from(&mut self, index: usize) -> Result<()> {
+        let mut i = index;
+        let mut depth = self.depth;
+        loop {
+            let h = self.hash_couple(depth, i);
+            i >>= 1;
+            depth -= 1;
+            self.nodes.insert((depth, i), h);
+            if depth == 0 {
+                break;
+            }
+        }
+        if depth != 0 {
+            return Err(Report::msg("did not reach the depth"));
+        }
+        if i != 0 {
+            return Err(Report::msg("did not go through all indexes"));
+        }
+        Ok(())
+    }
+}
+
+impl<H: Hasher> ZerokitMerkleProof<H> for OptimalMerkleProof<H>
+where
+    H: Hasher,
+{
+    type Index = u8;
+    #[must_use]
+    // Returns the length of a Merkle proof
+    fn length(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Computes the leaf index corresponding to a Merkle proof
+    #[must_use]
+    fn leaf_index(&self) -> usize {
+        // In current implementation the path indexes in a proof correspond to the binary representation of the leaf index
+        let mut binary_repr = self.get_path_index();
+        binary_repr.reverse();
+        binary_repr
+            .into_iter()
+            .fold(0, |acc, digit| (acc << 1) + usize::from(digit))
+    }
+
+    #[must_use]
+    /// Returns the path elements forming a Merkle proof
+    fn get_path_elements(&self) -> Vec<H::Fr> {
+        self.0.iter().map(|x| x.0).collect()
+    }
+
+    /// Returns the path indexes forming a Merkle proof
+    #[must_use]
+    fn get_path_index(&self) -> Vec<u8> {
+        self.0.iter().map(|x| x.1).collect()
+    }
+
+    #[must_use]
+    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
+    fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr {
+        let mut acc: H::Fr = *leaf;
+        for w in self.0.iter() {
+            if w.1 == 0 {
+                acc = H::hash(&[acc, w.0]);
+            } else {
+                acc = H::hash(&[w.0, acc]);
+            }
+        }
+        acc
+    }
+}
+
+// Debug formatting for printing a (Optimal) Merkle Proof
+impl<H> Debug for OptimalMerkleProof<H>
+where
+    H: Hasher,
+    H::Fr: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Proof").field(&self.0).finish()
+    }
+}

utils/tests/merkle_tree.rs

@@ -3,7 +3,7 @@
 mod test {
     use hex_literal::hex;
     use tiny_keccak::{Hasher as _, Keccak};
-    use utils::{FullMerkleTree, Hasher, OptimalMerkleTree};
+    use utils::{FullMerkleTree, Hasher, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     struct Keccak256;