refactor(rln): use space-efficient Merkle tree implementation (#32)

* refactor(rln): replace merkle tree implementation

refactor(rln): before switch

refactor(rln): replace merkle tree implementation

fix(rln): cargo fmt, cargo clippy

cargo fmt

* chore(rln): add comments

* chore(rln): cargo fmt; cargo clippy

* refactor(rln): improve Merkle tree crate

- Integrate previous Full storage MT with Optimal (storage) MT implementation
- Align the two implementation to same API/traits implementations
- Align implementations as much as possible to same language and variable names/meaning
- Comment code
- Add descriptions to each module's crate
- address reviewer's comments

* refactor(rln): fmt, clippy, remove pub from next index, update output type for some proc

rln/src/circuit.rs

@@ -1,3 +1,5 @@
+// This crate provides interfaces for the zero-knowledge circuit and keys
+
 use ark_bn254::{Bn254, Fq, Fq2, Fr, G1Affine, G1Projective, G2Affine, G2Projective};
 use ark_circom::{read_zkey, CircomBuilder, CircomConfig, WitnessCalculator};
 use ark_ff::BigInteger256;

rln/src/ffi.rs

@@ -1,3 +1,5 @@
+// This crate implements the public Foreign Function Interface (FFI) for the RLN module
+
 use crate::public::RLN;
 use std::slice;
 
@@ -492,7 +494,7 @@ mod test {
     }
 
     #[test]
-    fn test_100_groth16_proofs_ffi() {
+    fn test_groth16_proofs_performance_ffi() {
         let tree_height = TEST_TREE_HEIGHT;
 
         // We create a RLN instance

rln/src/lib.rs

@@ -189,7 +189,7 @@ mod test {
         // generate merkle tree
         let default_leaf = Field::from(0);
         let mut tree = PoseidonTree::new(tree_height, default_leaf);
-        tree.set(leaf_index, id_commitment.into());
+        tree.set(leaf_index, id_commitment.into()).unwrap();
 
         // We check correct computation of the root
         let root = tree.root();
@@ -297,7 +297,7 @@ mod test {
         assert_eq!(identity_path_index, expected_identity_path_index);
 
         // We check correct verification of the proof
-        assert!(tree.verify(id_commitment.into(), &merkle_proof));
+        assert!(tree.verify(&id_commitment, &merkle_proof).unwrap());
     }
 
     #[test]
@@ -344,7 +344,7 @@ mod test {
         //// generate merkle tree
         let default_leaf = Field::from(0);
         let mut tree = PoseidonTree::new(tree_height, default_leaf);
-        tree.set(leaf_index, id_commitment.into());
+        tree.set(leaf_index, id_commitment.into()).unwrap();
 
         let merkle_proof = tree.proof(leaf_index).expect("proof should exist");
 

rln/src/merkle_tree.rs

@@ -1,132 +1,382 @@
-// Implementation adapted from https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/merkle_tree.rs
+// This crate provides different implementation of Merkle tree
-// In our customization, we expand MerkleTree to have a next_index counter, so that we can add a public API to add leaves to the next available counter with no need to specify the index
+// Currently two interchangeable implementations are supported:
+//    - FullMerkleTree: each tree node is stored
+//    - OptimalMerkleTree: only nodes used to prove accumulation of set leaves are stored
+// Library defaults are set in the poseidon_tree crate
+//
+// Merkle tree implementations are adapted from https://github.com/kilic/rln/blob/master/src/merkle.rs
+// and https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/merkle_tree.rs
 
-//! Implements basic binary Merkle trees
 //!
 //! # To do
 //!
 //! * Disk based storage backend (using mmaped files should be easy)
 
-use semaphore::Field;
+use ark_std::str::FromStr;
 use serde::{Deserialize, Serialize};
+use std::io::{self, Error, ErrorKind};
 use std::{
     cmp::max,
     fmt::Debug,
     iter::{once, repeat, successors},
 };
+use std::{collections::HashMap, hash::Hash};
 
-/// Hash types, values and algorithms for a Merkle tree
+/// 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 node hashes
+    /// Type of the leaf and tree node
-    type Hash: Copy + Clone + Eq + Serialize;
+    type Fr: Copy + Clone + Eq + Serialize;
 
-    /// Compute the hash of an intermediate node
+    /// Returns the default tree leaf
-    fn hash_node(left: &Self::Hash, right: &Self::Hash) -> Self::Hash;
+    fn default_leaf() -> Self::Fr;
+
+    /// Utility to compute the hash of an intermediate node
+    fn hash(input: &[Self::Fr]) -> Self::Fr;
 }
 
-/// Merkle tree with all leaf and intermediate hashes stored
+////////////////////////////////////////////////////////////
-#[derive(Clone, PartialEq, Eq, Debug)]
+/// Optimal Merkle Tree Implementation
-pub struct MerkleTree<H: Hasher> {
+////////////////////////////////////////////////////////////
-    /// total number of levels of the tree, i.e. # of layers including pre-images layer of tree leaves
-    levels: usize,
 
-    /// Hash value of empty subtrees of given levels, starting at leaf level
+/// The Merkle tree structure
-    empty: Vec<H::Hash>,
+pub struct OptimalMerkleTree<H>
+where
+    H: Hasher,
+{
+    /// The depth of the tree, i.e. the number of levels from leaf to root
+    depth: usize,
 
-    /// Hash values of tree nodes and leaves, breadth first order
+    /// The nodes cached from the empty part of the tree (where leaves are set to default).
-    nodes: Vec<H::Hash>,
+    /// 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)
-    pub next_index: usize,
+    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, Serialize)]
+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: 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) -> io::Result<()> {
+        if index >= self.capacity() {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "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 a leaf at the next available index
+    pub fn update_next(&mut self, leaf: H::Fr) -> io::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) -> io::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) -> io::Result<OptimalMerkleProof<H>> {
+        if index >= self.capacity() {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "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;
+            }
+        }
+        assert_eq!(i, 0);
+        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>) -> io::Result<bool> {
+        if witness.length() != self.depth {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "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) {
+        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;
+            }
+        }
+        assert_eq!(depth, 0);
+        assert_eq!(i, 0);
+    }
+}
+
+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>
+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, Serialize, Deserialize)]
-pub enum Branch<H: Hasher> {
+pub enum FullMerkleBranch<H: Hasher> {
     /// Left branch taken, value is the right sibling hash.
-    Left(H::Hash),
+    Left(H::Fr),
 
     /// Right branch taken, value is the left sibling hash.
-    Right(H::Hash),
+    Right(H::Fr),
 }
 
 /// Merkle proof path, bottom to top.
 #[derive(Clone, PartialEq, Eq, Serialize)]
-pub struct Proof<H: Hasher>(pub Vec<Branch<H>>);
+pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
 
-/// For a given node index, return the parent node index
+/// Implementations
-/// Returns None if there is no parent (root node)
+
-const fn parent(index: usize) -> Option<usize> {
+impl<H: Hasher> FullMerkleTree<H> {
-    if index == 0 {
+    pub fn default(depth: usize) -> Self {
-        None
+        FullMerkleTree::<H>::new(depth, H::default_leaf())
-    } else {
-        Some(((index + 1) >> 1) - 1)
     }
-}
 
-/// For a given node index, return index of the first (left) child.
-const fn first_child(index: usize) -> usize {
-    (index << 1) + 1
-}
-
-const fn levels(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
-}
-
-impl<H: Hasher> MerkleTree<H> {
     /// Creates a new `MerkleTree`
-    /// tree_height - the height of the tree made only of hash nodes. 2^tree_height is the maximum number of leaves hash nodes
+    /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
-    pub fn new(tree_height: usize, initial_leaf: H::Hash) -> Self {
+    pub fn new(depth: usize, initial_leaf: H::Fr) -> Self {
-        // total number of levels of the tree, i.e. # of layers including pre-images layer of tree leaves, thus equal to tree_height+1
+        // Compute cache node values, leaf to root
-        let levels = tree_height + 1;
+        let cached_nodes = successors(Some(initial_leaf), |prev| Some(H::hash(&[*prev, *prev])))
-        // Compute empty node values, leaf to root
+            .take(depth + 1)
-        let empty = successors(Some(initial_leaf), |prev| Some(H::hash_node(prev, prev)))
-            .take(levels)
             .collect::<Vec<_>>();
 
         // Compute node values
-        let nodes = empty
+        let nodes = cached_nodes
             .iter()
             .rev()
             .enumerate()
             .flat_map(|(levels, hash)| repeat(hash).take(1 << levels))
             .cloned()
             .collect::<Vec<_>>();
-        debug_assert!(nodes.len() == (1 << levels) - 1);
+        debug_assert!(nodes.len() == (1 << (depth + 1)) - 1);
 
         let next_index = 0;
 
         Self {
-            levels,
+            depth,
-            empty,
+            cached_nodes,
             nodes,
             next_index,
         }
     }
 
-    #[must_use]
+    // Returns the depth of the tree
-    pub fn num_leaves(&self) -> usize {
+    pub fn depth(&self) -> usize {
-        self.levels
+        self.depth
-            .checked_sub(1)
+    }
-            .map(|n| 1 << n)
+
-            .unwrap_or_default()
+    // 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]
-    pub fn root(&self) -> H::Hash {
+    // Returns the root of the tree
+    pub fn root(&self) -> H::Fr {
         self.nodes[0]
     }
 
-    pub fn set(&mut self, leaf: usize, hash: H::Hash) {
+    // Sets a leaf at the specified tree index
-        self.set_range(leaf, once(hash));
+    pub fn set(&mut self, leaf: usize, hash: H::Fr) -> io::Result<()> {
+        self.set_range(leaf, once(hash))?;
         self.next_index = max(self.next_index, leaf + 1);
+        Ok(())
     }
 
-    pub fn set_range<I: IntoIterator<Item = H::Hash>>(&mut self, start: usize, hashes: I) {
+    // Sets tree nodes, starting from start index
-        let index = self.num_leaves() + start - 1;
+    // Function proper of FullMerkleTree implementation
+    fn set_range<I: IntoIterator<Item = H::Fr>>(
+        &mut self,
+        start: usize,
+        hashes: I,
+    ) -> io::Result<()> {
+        let index = self.capacity() + start - 1;
         let mut count = 0;
         // TODO: Error/panic when hashes is longer than available leafs
         for (leaf, hash) in self.nodes[index..].iter_mut().zip(hashes) {
@@ -137,95 +387,140 @@ impl<H: Hasher> MerkleTree<H> {
             self.update_nodes(index, index + (count - 1));
             self.next_index = max(self.next_index, start + count);
         }
+        Ok(())
     }
 
-    fn update_nodes(&mut self, start: usize, end: usize) {
+    // Sets a leaf at the next available index
-        debug_assert_eq!(levels(start), levels(end));
+    pub fn update_next(&mut self, leaf: H::Fr) -> io::Result<()> {
-        if let (Some(start), Some(end)) = (parent(start), parent(end)) {
+        self.set(self.next_index, leaf)?;
-            for parent in start..=end {
+        Ok(())
-                let child = first_child(parent);
-                self.nodes[parent] = H::hash_node(&self.nodes[child], &self.nodes[child + 1]);
-            }
-            self.update_nodes(start, end);
-        }
     }
 
-    #[must_use]
+    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
-    pub fn proof(&self, leaf: usize) -> Option<Proof<H>> {
+    pub fn delete(&mut self, index: usize) -> io::Result<()> {
-        if leaf >= self.num_leaves() {
+        // We reset the leaf only if we previously set a leaf at that index
-            return None;
+        if index < self.next_index {
+            self.set(index, H::default_leaf())?;
         }
-        let mut index = self.num_leaves() + leaf - 1;
+        Ok(())
-        let mut path = Vec::with_capacity(self.levels);
+    }
-        while let Some(parent) = parent(index) {
+
+    // Computes a merkle proof the the leaf at the specified index
+    pub fn proof(&self, leaf: usize) -> io::Result<FullMerkleProof<H>> {
+        if leaf >= self.capacity() {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "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 => Branch::Left(self.nodes[index + 1]),
+                1 => FullMerkleBranch::Left(self.nodes[index + 1]),
-                0 => Branch::Right(self.nodes[index - 1]),
+                0 => FullMerkleBranch::Right(self.nodes[index - 1]),
                 _ => unreachable!(),
             });
             index = parent;
         }
-        Some(Proof(path))
+        Ok(FullMerkleProof(path))
     }
 
-    #[must_use]
+    // Verifies a Merkle proof with respect to the input leaf and the tree root
-    pub fn verify(&self, hash: H::Hash, proof: &Proof<H>) -> bool {
+    pub fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof<H>) -> io::Result<bool> {
-        proof.root(hash) == self.root()
+        Ok(proof.compute_root_from(hash) == self.root())
     }
 
-    #[must_use]
+    // Utilities for updating the tree nodes
-    pub fn leaves(&self) -> &[H::Hash] {
+
-        &self.nodes[(self.num_leaves() - 1)..]
+    /// 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) {
+        debug_assert_eq!(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);
+        }
     }
 }
 
-impl<H: Hasher> Proof<H> {
+impl<H: Hasher> FullMerkleProof<H> {
-    /// Compute the leaf index for this proof
+    #[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 {
-            Branch::Left(_) => index << 1,
+            FullMerkleBranch::Left(_) => index << 1,
-            Branch::Right(_) => (index << 1) + 1,
+            FullMerkleBranch::Right(_) => (index << 1) + 1,
         })
     }
 
     #[must_use]
-    pub fn get_path_elements(&self) -> Vec<H::Hash> {
+    /// Returns the path elements forming a Merkle proof
+    pub fn get_path_elements(&self) -> Vec<H::Fr> {
         self.0
             .iter()
             .map(|x| match x {
-                Branch::Left(value) | Branch::Right(value) => *value,
+                FullMerkleBranch::Left(value) | FullMerkleBranch::Right(value) => *value,
             })
             .collect()
     }
 
-    /// Compute path index (TODO: do we want to keep this here?)
+    /// 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 {
-                Branch::Left(_) => 0,
+                FullMerkleBranch::Left(_) => 0,
-                Branch::Right(_) => 1,
+                FullMerkleBranch::Right(_) => 1,
             })
             .collect()
     }
 
-    /// Compute the Merkle root given a leaf hash
+    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
     #[must_use]
-    pub fn root(&self, hash: H::Hash) -> H::Hash {
+    pub fn compute_root_from(&self, hash: &H::Fr) -> H::Fr {
-        self.0.iter().fold(hash, |hash, branch| match branch {
+        self.0.iter().fold(*hash, |hash, branch| match branch {
-            Branch::Left(sibling) => H::hash_node(&hash, sibling),
+            FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]),
-            Branch::Right(sibling) => H::hash_node(sibling, &hash),
+            FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]),
         })
     }
 }
 
-impl<H> Debug for Branch<H>
+// Debug formatting for printing a (Full) Merkle Proof Branch
+impl<H> Debug for FullMerkleBranch<H>
 where
     H: Hasher,
-    H::Hash: Debug,
+    H::Fr: Debug,
 {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -235,17 +530,22 @@ where
     }
 }
 
-impl<H> Debug for Proof<H>
+// Debug formatting for printing a (Full) Merkle Proof
+impl<H> Debug for FullMerkleProof<H>
 where
     H: Hasher,
-    H::Hash: Debug,
+    H::Fr: Debug,
 {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_tuple("Proof").field(&self.0).finish()
     }
 }
 
-#[cfg(test)]
+////////////////////////////////////////////////////////////
+/// Tests
+////////////////////////////////////////////////////////////
+
+// Tests adapted from https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/merkle_tree.rs
 pub mod test {
     use super::*;
     use hex_literal::hex;
@@ -254,127 +554,112 @@ pub mod test {
     struct Keccak256;
 
     impl Hasher for Keccak256 {
-        type Hash = [u8; 32];
+        type Fr = [u8; 32];
 
-        fn hash_node(left: &Self::Hash, right: &Self::Hash) -> Self::Hash {
+        fn default_leaf() -> Self::Fr {
+            [0; 32]
+        }
+
+        fn hash(inputs: &[Self::Fr]) -> Self::Fr {
             let mut output = [0; 32];
             let mut hasher = Keccak::v256();
-            hasher.update(left);
+            for element in inputs {
-            hasher.update(right);
+                hasher.update(element);
+            }
             hasher.finalize(&mut output);
             output
         }
     }
 
-    #[test]
-    fn test_index_calculus() {
-        assert_eq!(parent(0), None);
-        assert_eq!(parent(1), Some(0));
-        assert_eq!(parent(2), Some(0));
-        assert_eq!(parent(3), Some(1));
-        assert_eq!(parent(4), Some(1));
-        assert_eq!(parent(5), Some(2));
-        assert_eq!(parent(6), Some(2));
-        assert_eq!(first_child(0), 1);
-        assert_eq!(first_child(2), 5);
-        assert_eq!(levels(0), 0);
-        assert_eq!(levels(1), 1);
-        assert_eq!(levels(2), 1);
-        assert_eq!(levels(3), 2);
-        assert_eq!(levels(6), 2);
-    }
-
     #[test]
     fn test_root() {
-        let mut tree = MerkleTree::<Keccak256>::new(2, [0; 32]);
+        let leaves = [
-        assert_eq!(
-            tree.root(),
-            hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30")
-        );
-        tree.set(
-            0,
             hex!("0000000000000000000000000000000000000000000000000000000000000001"),
-        );
-        assert_eq!(
-            tree.root(),
-            hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95")
-        );
-        tree.set(
-            1,
             hex!("0000000000000000000000000000000000000000000000000000000000000002"),
-        );
-        assert_eq!(
-            tree.root(),
-            hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b")
-        );
-        tree.set(
-            2,
             hex!("0000000000000000000000000000000000000000000000000000000000000003"),
-        );
-        assert_eq!(
-            tree.root(),
-            hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c")
-        );
-        tree.set(
-            3,
             hex!("0000000000000000000000000000000000000000000000000000000000000004"),
-        );
+        ];
-        assert_eq!(
+
-            tree.root(),
+        let default_tree_root =
-            hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36")
+            hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30");
-        );
+
+        let roots = [
+            hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95"),
+            hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b"),
+            hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c"),
+            hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"),
+        ];
+
+        let mut tree = FullMerkleTree::<Keccak256>::new(2, [0; 32]);
+        assert_eq!(tree.root(), default_tree_root);
+        for i in 0..leaves.len() {
+            tree.set(i, leaves[i]).unwrap();
+            assert_eq!(tree.root(), roots[i]);
+        }
+
+        let mut tree = OptimalMerkleTree::<Keccak256>::new(2, [0; 32]);
+        assert_eq!(tree.root(), default_tree_root);
+        for i in 0..leaves.len() {
+            tree.set(i, leaves[i]).unwrap();
+            assert_eq!(tree.root(), roots[i]);
+        }
     }
 
     #[test]
     fn test_proof() {
-        let mut tree = MerkleTree::<Keccak256>::new(2, [0; 32]);
+        let leaves = [
-        tree.set(
-            0,
             hex!("0000000000000000000000000000000000000000000000000000000000000001"),
-        );
-        tree.set(
-            1,
             hex!("0000000000000000000000000000000000000000000000000000000000000002"),
-        );
-        tree.set(
-            2,
             hex!("0000000000000000000000000000000000000000000000000000000000000003"),
-        );
-        tree.set(
-            3,
             hex!("0000000000000000000000000000000000000000000000000000000000000004"),
-        );
+        ];
 
-        let proof = tree.proof(2).expect("proof should exist");
+        // We thest the FullMerkleTree implementation
-        assert_eq!(proof.leaf_index(), 2);
+        let mut tree = FullMerkleTree::<Keccak256>::new(2, [0; 32]);
-        assert!(tree.verify(
+        for i in 0..leaves.len() {
-            hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+            // We set the leaves
-            &proof
+            tree.set(i, leaves[i]).unwrap();
-        ));
-        assert!(!tree.verify(
-            hex!("0000000000000000000000000000000000000000000000000000000000000001"),
-            &proof
-        ));
-    }
 
-    #[test]
+            // We compute a merkle proof
-    fn test_position() {
+            let proof = tree.proof(i).expect("index should be set");
-        let mut tree = MerkleTree::<Keccak256>::new(2, [0; 32]);
+
-        tree.set(
+            // We verify if the merkle proof corresponds to the right leaf index
-            0,
+            assert_eq!(proof.leaf_index(), i);
-            hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+
-        );
+            // We verify the proof
-        tree.set(
+            assert!(tree.verify(&leaves[i], &proof).unwrap());
-            1,
+
-            hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+            // We ensure that the Merkle proof and the leaf generate the same root as the tree
-        );
+            assert_eq!(proof.compute_root_from(&leaves[i]), tree.root());
-        tree.set(
+
-            2,
+            // We check that the proof is not valid for another leaf
-            hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+            assert!(!tree
-        );
+                .verify(&leaves[(i + 1) % leaves.len()], &proof)
-        tree.set(
+                .unwrap());
-            3,
+        }
-            hex!("0000000000000000000000000000000000000000000000000000000000000004"),
+
-        );
+        // We test the OptimalMerkleTree implementation
+        let mut tree = OptimalMerkleTree::<Keccak256>::new(2, [0; 32]);
+        for i in 0..leaves.len() {
+            // We set the leaves
+            tree.set(i, leaves[i]).unwrap();
+
+            // We compute a merkle proof
+            let proof = tree.proof(i).expect("index should be set");
+
+            // We verify if the merkle proof corresponds to the right leaf index
+            assert_eq!(proof.leaf_index(), i);
+
+            // We verify the proof
+            assert!(tree.verify(&leaves[i], &proof).unwrap());
+
+            // We ensure that the Merkle proof and the leaf generate the same root as the tree
+            assert_eq!(proof.compute_root_from(&leaves[i]), tree.root());
+
+            // We check that the proof is not valid for another leaf
+            assert!(!tree
+                .verify(&leaves[(i + 1) % leaves.len()], &proof)
+                .unwrap());
+        }
     }
 }

rln/src/poseidon_tree.rs

@@ -1,25 +1,107 @@
-// Implementation taken from https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs (no differences)
+// This crate defines RLN module default Merkle tree implementation and Hasher
-// Implements Merkle trees with Poseidon hash for the customized semaphore-rs merkle_tree implementation
+// Implementation inspired by https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs (no differences)
 
-use crate::merkle_tree::{self, Hasher, MerkleTree};
+use crate::merkle_tree::{
+    FullMerkleProof, FullMerkleTree, Hasher, OptimalMerkleProof, OptimalMerkleTree,
+};
 use semaphore::{poseidon_hash, Field};
-
 use serde::{Deserialize, Serialize};
 
+// The zerokit RLN default Merkle tree implementation.
+// To switch to FullMerkleTree implementation it is enough to redefine the following two types
 #[allow(dead_code)]
-pub type PoseidonTree = MerkleTree<PoseidonHash>;
+pub type PoseidonTree = OptimalMerkleTree<PoseidonHash>;
-#[allow(dead_code)]
+pub type MerkleProof = OptimalMerkleProof<PoseidonHash>;
-pub type Branch = merkle_tree::Branch<PoseidonHash>;
+//pub type PoseidonTree = FullMerkleTree<PoseidonHash>;
-#[allow(dead_code)]
+//pub type MerkleProof = FullMerkleProof<PoseidonHash>;
-pub type Proof = merkle_tree::Proof<PoseidonHash>;
 
+// The zerokit RLN default Hasher
 #[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub struct PoseidonHash;
 
 impl Hasher for PoseidonHash {
-    type Hash = Field;
+    type Fr = Field;
 
-    fn hash_node(left: &Self::Hash, right: &Self::Hash) -> Self::Hash {
+    fn default_leaf() -> Self::Fr {
-        poseidon_hash(&[*left, *right])
+        Self::Fr::from(0)
+    }
+
+    fn hash(inputs: &[Self::Fr]) -> Self::Fr {
+        poseidon_hash(inputs)
     }
 }
+
+////////////////////////////////////////////////////////////
+/// Tests
+////////////////////////////////////////////////////////////
+
+#[test]
+/// A basic performance comparison between the two supported Merkle Tree implementations
+fn test_merkle_implementations_performances() {
+    use std::time::{Duration, Instant};
+
+    let tree_height = 20;
+    let sample_size = 100;
+
+    let leaves: Vec<Field> = (0..sample_size).map(|s| Field::from(s)).collect();
+
+    let mut gen_time_full: u128 = 0;
+    let mut upd_time_full: u128 = 0;
+    let mut gen_time_opt: u128 = 0;
+    let mut upd_time_opt: u128 = 0;
+
+    for _ in 0..sample_size.try_into().unwrap() {
+        let now = Instant::now();
+        FullMerkleTree::<PoseidonHash>::default(tree_height);
+        gen_time_full += now.elapsed().as_nanos();
+
+        let now = Instant::now();
+        OptimalMerkleTree::<PoseidonHash>::default(tree_height);
+        gen_time_opt += now.elapsed().as_nanos();
+    }
+
+    let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height);
+    let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height);
+    for i in 0..sample_size.try_into().unwrap() {
+        let now = Instant::now();
+        tree_full.set(i, leaves[i]).unwrap();
+        upd_time_full += now.elapsed().as_nanos();
+        let proof = tree_full.proof(i).expect("index should be set");
+        assert_eq!(proof.leaf_index(), i);
+
+        let now = Instant::now();
+        tree_opt.set(i, leaves[i]).unwrap();
+        upd_time_opt += now.elapsed().as_nanos();
+        let proof = tree_opt.proof(i).expect("index should be set");
+        assert_eq!(proof.leaf_index(), i);
+    }
+
+    println!("Average tree generation time:");
+    println!(
+        "   - Full Merkle Tree:  {:?}",
+        Duration::from_nanos(
+            (gen_time_full / u128::from(sample_size))
+                .try_into()
+                .unwrap()
+        )
+    );
+    println!(
+        "   - Optimal Merkle Tree: {:?}",
+        Duration::from_nanos((gen_time_opt / u128::from(sample_size)).try_into().unwrap())
+    );
+
+    println!("Average update_next execution time:");
+    println!(
+        "   - Full Merkle Tree: {:?}",
+        Duration::from_nanos(
+            (upd_time_full / u128::from(sample_size))
+                .try_into()
+                .unwrap()
+        )
+    );
+
+    println!(
+        "   - Optimal Merkle Tree: {:?}",
+        Duration::from_nanos((upd_time_opt / u128::from(sample_size)).try_into().unwrap())
+    );
+}

rln/src/protocol.rs

@@ -1,6 +1,6 @@
+// This crate collects all the underlying primitives used to implement RLN
+
 use crate::circuit::{CIRCOM, VK, ZKEY};
-use crate::merkle_tree::{self, Branch};
-use crate::poseidon_tree::PoseidonHash;
 use ark_bn254::{Bn254, Fr, Parameters};
 use ark_circom::{read_zkey, CircomBuilder, CircomConfig, CircomReduction, WitnessCalculator};
 use ark_ec::bn::Bn;
@@ -237,7 +237,7 @@ pub fn rln_witness_from_json(input_json_str: &str) -> RLNWitnessInput {
 
 pub fn rln_witness_from_values(
     identity_secret: Field,
-    merkle_proof: &merkle_tree::Proof<PoseidonHash>,
+    merkle_proof: &MerkleProof,
     x: Field,
     epoch: Field,
     //rln_identifier: Field,

rln/src/public.rs

@@ -1,5 +1,4 @@
-use crate::poseidon_tree::PoseidonTree;
+/// This is the main public API for RLN module. It is used by the FFI, and should be
-/// This is the main public API for RLN. It is used by the FFI, and should be
 /// used by tests etc as well
 ///
 use ark_bn254::{Bn254, Fr};
@@ -25,6 +24,7 @@ use ark_ff::bytes::ToBytes;
 use ark_serialize::{Read, Write};
 
 use crate::circuit::{CIRCOM, TEST_RESOURCES_FOLDER, TEST_TREE_HEIGHT, VK, ZKEY};
+use crate::poseidon_tree::PoseidonTree;
 use crate::protocol::{self, *};
 use crate::utils::*;
 
@@ -55,8 +55,7 @@ impl RLN<'_> {
         let verification_key = VK(&resources_folder);
 
         // We compute a default empty tree
-        let leaf = Field::from(0);
+        let tree = PoseidonTree::default(tree_height);
-        let tree = PoseidonTree::new(tree_height, leaf);
 
         RLN {
             witness_calculator,
@@ -72,8 +71,7 @@ impl RLN<'_> {
     ////////////////////////////////////////////////////////
     pub fn set_tree(&mut self, tree_height: usize) -> io::Result<()> {
         // We compute a default empty tree of desired height
-        let leaf = Field::from(0);
+        self.tree = PoseidonTree::default(tree_height);
-        self.tree = PoseidonTree::new(tree_height, leaf);
 
         Ok(())
     }
@@ -85,7 +83,7 @@ impl RLN<'_> {
 
         // We set the leaf at input index
         let (leaf, _) = bytes_le_to_field(&leaf_byte);
-        self.tree.set(index, leaf);
+        self.tree.set(index, leaf)?;
 
         Ok(())
     }
@@ -100,7 +98,7 @@ impl RLN<'_> {
 
         // We set the leaves
         for (i, leaf) in leaves.iter().enumerate() {
-            self.tree.set(i, *leaf);
+            self.tree.set(i, *leaf)?;
         }
 
         Ok(())
@@ -114,19 +112,14 @@ impl RLN<'_> {
 
         // We set the leaf at input index
         let (leaf, _) = bytes_le_to_field(&leaf_byte);
-        self.tree.set(self.tree.next_index, leaf);
+        self.tree.update_next(leaf)?;
 
         Ok(())
     }
 
     // Deleting a leaf corresponds to set its value to the default 0 leaf
     pub fn delete_leaf(&mut self, index: usize) -> io::Result<()> {
-        // We reset the leaf only if we previously set a leaf at that index
+        self.tree.delete(index)?;
-        if index < self.tree.next_index {
-            let leaf = Field::from(0);
-            self.tree.set(index, leaf);
-        }
-
         Ok(())
     }
 
@@ -328,8 +321,8 @@ mod test {
 
         // We first add leaves one by one specifying the index
         for (i, leaf) in leaves.iter().enumerate() {
-            // We check if internal index is properly set
+            // We check if the number of leaves set is consistent
-            assert_eq!(rln.tree.next_index, i);
+            assert_eq!(rln.tree.leaves_set(), i);
 
             let mut buffer = Cursor::new(field_to_bytes_le(&leaf));
             rln.set_leaf(i, &mut buffer).unwrap();
@@ -349,8 +342,8 @@ mod test {
             rln.set_next_leaf(&mut buffer).unwrap();
         }
 
-        // We check if internal index is properly set
+        // We check if numbers of leaves set is consistent
-        assert_eq!(rln.tree.next_index, no_of_leaves);
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
 
         // We get the root of the tree obtained adding leaves using the internal index
         let mut buffer = Cursor::new(Vec::<u8>::new());
@@ -366,8 +359,8 @@ mod test {
         let mut buffer = Cursor::new(vec_field_to_bytes_le(&leaves));
         rln.set_leaves(&mut buffer).unwrap();
 
-        // We check if internal index is properly set
+        // We check if number of leaves set is consistent
-        assert_eq!(rln.tree.next_index, no_of_leaves);
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
 
         // We get the root of the tree obtained adding leaves in batch
         let mut buffer = Cursor::new(Vec::<u8>::new());
@@ -382,8 +375,8 @@ mod test {
             rln.delete_leaf(i).unwrap();
         }
 
-        // We check if internal index is properly set
+        // We check if number of leaves set is consistent
-        assert_eq!(rln.tree.next_index, no_of_leaves);
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
 
         let mut buffer = Cursor::new(Vec::<u8>::new());
         rln.get_root(&mut buffer).unwrap();
@@ -601,7 +594,7 @@ mod test {
         let (identity_secret, id_commitment) = keygen();
 
         // We set as leaf id_commitment after storing its index
-        let identity_index = u64::try_from(rln.tree.next_index).unwrap();
+        let identity_index = u64::try_from(rln.tree.leaves_set()).unwrap();
         let mut buffer = Cursor::new(field_to_bytes_le(&id_commitment));
         rln.set_next_leaf(&mut buffer).unwrap();
 

rln/src/utils.rs

@@ -1,3 +1,5 @@
+// This crate provides cross-module useful utilities (mainly type conversions) not necessarily specific to RLN
+
 use ark_bn254::{Bn254, Fr, Parameters};
 use ark_ff::{BigInteger, Field as ArkField, FpParameters, PrimeField};
 use ark_std::str::FromStr;