feat(rln): function for checking indices of leaves which are set to z…

…ero (#249)

* add function for empty leaves

* fix from linter

* fix rebase

* update test in utils

* fix

* fix(trees): inconsistencies in override_range (#250)

* fix tests

---------

Co-authored-by: Aaryamann Challani <43716372+rymnc@users.noreply.github.com>

rln/benches/pmtree_benchmark.rs

@@ -44,6 +44,12 @@ pub fn pmtree_benchmark(c: &mut Criterion) {
             tree.get_subtree_root(1, 0).unwrap();
         })
     });
+
+    c.bench_function("Pmtree::get_empty_leaves_indices", |b| {
+        b.iter(|| {
+            tree.get_empty_leaves_indices();
+        })
+    });
 }
 
 criterion_group!(benches, pmtree_benchmark);

rln/src/pm_tree_adapter.rs

@@ -17,6 +17,9 @@ const METADATA_KEY: [u8; 8] = *b"metadata";
 
 pub struct PmTree {
     tree: pmtree::MerkleTree<SledDB, PoseidonHash>,
+    /// The indices of leaves which are set into zero upto next_index.
+    /// Set to 0 if the leaf is empty and set to 1 in otherwise.
+    cached_leaves_indices: Vec<u8>,
     // metadata that an application may use to store additional information
     metadata: Vec<u8>,
 }
@@ -144,6 +147,7 @@ impl ZerokitMerkleTree for PmTree {
 
         Ok(PmTree {
             tree,
+            cached_leaves_indices: vec![0; 1 << depth],
             metadata: Vec::new(),
         })
     }
@@ -156,7 +160,7 @@ impl ZerokitMerkleTree for PmTree {
         self.tree.capacity()
     }
 
-    fn leaves_set(&mut self) -> usize {
+    fn leaves_set(&self) -> usize {
         self.tree.leaves_set()
     }
 
@@ -171,17 +175,24 @@ impl ZerokitMerkleTree for PmTree {
     fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>) -> Result<()> {
         self.tree
             .set(index, leaf)
-            .map_err(|e| Report::msg(e.to_string()))
+            .map_err(|e| Report::msg(e.to_string()))?;
+        self.cached_leaves_indices[index] = 1;
+        Ok(())
     }
 
     fn set_range<I: IntoIterator<Item = FrOf<Self::Hasher>>>(
         &mut self,
         start: usize,
         values: I,
     ) -> Result<()> {
+        let v = values.into_iter().collect::<Vec<_>>();
         self.tree
-            .set_range(start, values)
-            .map_err(|e| Report::msg(e.to_string()))
+            .set_range(start, v.clone().into_iter())
+            .map_err(|e| Report::msg(e.to_string()))?;
+        for i in start..v.len() {
+            self.cached_leaves_indices[i] = 1
+        }
+        Ok(())
     }
 
     fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>> {
@@ -208,6 +219,17 @@ impl ZerokitMerkleTree for PmTree {
         }
     }
 
+    fn get_empty_leaves_indices(&self) -> Vec<usize> {
+        let next_idx = self.leaves_set();
+        self.cached_leaves_indices
+            .iter()
+            .take(next_idx)
+            .enumerate()
+            .filter(|&(_, &v)| v == 0u8)
+            .map(|(idx, _)| idx)
+            .collect()
+    }
+
     fn override_range<I: IntoIterator<Item = FrOf<Self::Hasher>>, J: IntoIterator<Item = usize>>(
         &mut self,
         start: usize,
@@ -222,7 +244,7 @@ impl ZerokitMerkleTree for PmTree {
             (0, 0) => Err(Report::msg("no leaves or indices to be removed")),
             (1, 0) => self.set(start, leaves[0]),
             (0, 1) => self.delete(indices[0]),
-            (_, 0) => self.set_range_with_leaves(start, leaves),
+            (_, 0) => self.set_range(start, leaves),
             (0, _) => self.remove_indices(&indices),
             (_, _) => self.remove_indices_and_set_leaves(start, leaves, &indices),
         }
@@ -237,7 +259,9 @@ impl ZerokitMerkleTree for PmTree {
     fn delete(&mut self, index: usize) -> Result<()> {
         self.tree
             .delete(index)
-            .map_err(|e| Report::msg(e.to_string()))
+            .map_err(|e| Report::msg(e.to_string()))?;
+        self.cached_leaves_indices[index] = 0;
+        Ok(())
     }
 
     fn proof(&self, index: usize) -> Result<Self::Proof> {
@@ -282,12 +306,6 @@ type PmTreeHasher = <PmTree as ZerokitMerkleTree>::Hasher;
 type FrOfPmTreeHasher = FrOf<PmTreeHasher>;
 
 impl PmTree {
-    fn set_range_with_leaves(&mut self, start: usize, leaves: Vec<FrOfPmTreeHasher>) -> Result<()> {
-        self.tree
-            .set_range(start, leaves)
-            .map_err(|e| Report::msg(e.to_string()))
-    }
-
     fn remove_indices(&mut self, indices: &[usize]) -> Result<()> {
         let start = indices[0];
         let end = indices.last().unwrap() + 1;
@@ -296,7 +314,12 @@ impl PmTree {
 
         self.tree
             .set_range(start, new_leaves)
-            .map_err(|e| Report::msg(e.to_string()))
+            .map_err(|e| Report::msg(e.to_string()))?;
+
+        for i in start..end {
+            self.cached_leaves_indices[i] = 0
+        }
+        Ok(())
     }
 
     fn remove_indices_and_set_leaves(
@@ -322,8 +345,17 @@ impl PmTree {
         }
 
         self.tree
-            .set_range(min_index, set_values)
-            .map_err(|e| Report::msg(e.to_string()))
+            .set_range(start, set_values)
+            .map_err(|e| Report::msg(e.to_string()))?;
+
+        for i in indices {
+            self.cached_leaves_indices[*i] = 0;
+        }
+
+        for i in start..(max_index - min_index) {
+            self.cached_leaves_indices[i] = 1
+        }
+        Ok(())
     }
 }
 

rln/src/public.rs

@@ -603,6 +603,44 @@ impl RLN<'_> {
         Ok(())
     }
 
+    /// Returns indices of leaves in the tree are set to zero (upto the final leaf that was set).
+    ///
+    /// Output values are:
+    /// - `output_data`: a writer receiving the serialization of the indices of leaves.
+    ///
+    /// Example
+    /// ```
+    /// use rln::circuit::Fr;
+    /// use rln::utils::*;
+    ///
+    /// let start_index = 5;
+    /// let no_of_leaves = 256;
+    ///
+    /// // We generate a vector of random leaves
+    /// let mut leaves: Vec<Fr> = Vec::new();
+    /// let mut rng = thread_rng();
+    /// for _ in 0..no_of_leaves {
+    ///     let (_, id_commitment) = keygen();
+    ///     let rate_commitment = poseidon_hash(&[id_commitment, 1.into()]);
+    ///     leaves.push(rate_commitment);
+    /// }
+    ///
+    /// // We add leaves in a batch into the tree
+    /// let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
+    /// rln.set_leaves_from(index, &mut buffer).unwrap();
+    ///
+    /// // Get indices of first empty leaves upto start_index
+    /// let mut buffer = Cursor::new(Vec::<u8>::new());
+    /// rln.get_empty_leaves_indices(&mut buffer).unwrap();
+    /// let idxs = bytes_le_to_vec_usize(&buffer.into_inner()).unwrap();
+    /// assert_eq!(idxs, [0, 1, 2, 3, 4]);
+    /// ```
+    pub fn get_empty_leaves_indices<W: Write>(&self, mut output_data: W) -> Result<()> {
+        let idxs = self.tree.get_empty_leaves_indices();
+        idxs.serialize_compressed(&mut output_data)?;
+        Ok(())
+    }
+
     ////////////////////////////////////////////////////////
     // zkSNARK APIs
     ////////////////////////////////////////////////////////

rln/src/utils.rs

@@ -181,6 +181,19 @@ pub fn normalize_usize(input: usize) -> Vec<u8> {
     normalized_usize
 }
 
+pub fn bytes_le_to_vec_usize(input: &[u8]) -> Result<Vec<usize>> {
+    let nof_elem = usize::try_from(u64::from_le_bytes(input[0..8].try_into()?))?;
+    if nof_elem == 0 {
+        Ok(vec![])
+    } else {
+        let elements: Vec<usize> = input[8..]
+            .chunks(8)
+            .map(|ch| usize::from_le_bytes(ch[0..8].try_into().unwrap()))
+            .collect();
+        Ok(elements)
+    }
+}
+
 // using for test
 pub fn generate_input_buffer() -> Cursor<String> {
     Cursor::new(json!({}).to_string())

rln/tests/poseidon_tree.rs

@@ -68,4 +68,72 @@ mod test {
             }
         }
     }
+
+    #[test]
+    fn test_get_empty_leaves_indices() {
+        let depth = 4;
+        let nof_leaves: usize = 1 << (depth - 1);
+
+        let mut tree = PoseidonTree::default(depth).unwrap();
+        let leaves: Vec<Fr> = (0..nof_leaves).map(|s| Fr::from(s as i32)).collect();
+
+        // check set_range
+        let _ = tree.set_range(0, leaves.clone());
+        assert!(tree.get_empty_leaves_indices().is_empty());
+
+        let mut vec_idxs = Vec::new();
+        // check delete function
+        for i in 0..nof_leaves {
+            vec_idxs.push(i);
+            let _ = tree.delete(i);
+            assert_eq!(tree.get_empty_leaves_indices(), vec_idxs);
+        }
+        // check set function
+        for i in (0..nof_leaves).rev() {
+            vec_idxs.pop();
+            let _ = tree.set(i, leaves[i]);
+            assert_eq!(tree.get_empty_leaves_indices(), vec_idxs);
+        }
+
+        // check remove_indices_and_set_leaves inside override_range function
+        assert!(tree.get_empty_leaves_indices().is_empty());
+        let leaves_2: Vec<Fr> = (0..2).map(|s| Fr::from(s as i32)).collect();
+        tree.override_range(0, leaves_2.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree.get_empty_leaves_indices(), vec![2, 3]);
+
+        // check remove_indices inside override_range function
+        tree.override_range(0, [], [0, 1]).unwrap();
+        assert_eq!(tree.get_empty_leaves_indices(), vec![0, 1, 2, 3]);
+
+        // check set_range inside override_range function
+        tree.override_range(0, leaves_2.clone(), []).unwrap();
+        assert_eq!(tree.get_empty_leaves_indices(), vec![2, 3]);
+
+        let leaves_4: Vec<Fr> = (0..4).map(|s| Fr::from(s as i32)).collect();
+        // check if the indexes for write and delete are the same
+        tree.override_range(0, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert!(tree.get_empty_leaves_indices().is_empty());
+
+        // check if indexes for deletion are before indexes for overwriting
+        tree.override_range(4, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        // The result will be like this, because in the set_range function in pmtree
+        // the next_index value is increased not by the number of elements to insert,
+        // but by the union of indices for deleting and inserting.
+        assert_eq!(
+            tree.get_empty_leaves_indices(),
+            vec![0, 1, 2, 3, 8, 9, 10, 11]
+        );
+
+        // check if the indices for write and delete do not overlap completely
+        tree.override_range(2, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        // The result will be like this, because in the set_range function in pmtree
+        // the next_index value is increased not by the number of elements to insert,
+        // but by the union of indices for deleting and inserting.
+        // + we've already set to 6 and 7 in previous test
+        assert_eq!(tree.get_empty_leaves_indices(), vec![0, 1, 8, 9, 10, 11]);
+    }
 }

rln/tests/public.rs

@@ -23,10 +23,22 @@ mod test {
         let id_commitment = utils_poseidon_hash(&vec![identity_secret_hash]);
         let rate_commitment = utils_poseidon_hash(&[id_commitment, user_message_limit.into()]);
 
+        // check that leaves indices is empty
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_empty_leaves_indices(&mut buffer).unwrap();
+        let idxs = bytes_le_to_vec_usize(&buffer.into_inner()).unwrap();
+        assert!(idxs.is_empty());
+
         // We pass rate_commitment as Read buffer to RLN's set_leaf
         let mut buffer = Cursor::new(fr_to_bytes_le(&rate_commitment));
         rln.set_leaf(leaf_index, &mut buffer).unwrap();
 
+        // check that leaves before leaf_index is set to zero
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_empty_leaves_indices(&mut buffer).unwrap();
+        let idxs = bytes_le_to_vec_usize(&buffer.into_inner()).unwrap();
+        assert_eq!(idxs, [0, 1, 2]);
+
         // We check correct computation of the root
         let mut buffer = Cursor::new(Vec::<u8>::new());
         rln.get_root(&mut buffer).unwrap();

utils/benches/merkle_tree_benchmark.rs

@@ -97,6 +97,12 @@ pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) {
             tree.get_subtree_root(1, 0).unwrap();
         })
     });
+
+    c.bench_function("OptimalMerkleTree::get_empty_leaves_indices", |b| {
+        b.iter(|| {
+            tree.get_empty_leaves_indices();
+        })
+    });
 }
 
 pub fn full_merkle_tree_benchmark(c: &mut Criterion) {
@@ -139,6 +145,12 @@ pub fn full_merkle_tree_benchmark(c: &mut Criterion) {
             tree.get_subtree_root(1, 0).unwrap();
         })
     });
+
+    c.bench_function("FullMerkleTree::get_empty_leaves_indices", |b| {
+        b.iter(|| {
+            tree.get_empty_leaves_indices();
+        })
+    });
 }
 
 criterion_group!(