feat(rln): function for checking indices of leaves which are set to zero (#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,7 +175,9 @@ 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>>>(
@@ -179,9 +185,14 @@ impl ZerokitMerkleTree for PmTree {
         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!(

utils/src/merkle_tree/full_merkle_tree.rs

@@ -26,6 +26,10 @@ pub struct FullMerkleTree<H: Hasher> {
     /// The tree nodes
     nodes: Vec<H::Fr>,
 
+    /// 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>,
+
     // 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,
@@ -96,6 +100,7 @@ where
             depth,
             cached_nodes,
             nodes,
+            cached_leaves_indices: vec![0; 1 << depth],
             next_index,
             metadata: Vec::new(),
         })
@@ -116,7 +121,7 @@ where
     }
 
     // Returns the total number of leaves set
-    fn leaves_set(&mut self) -> usize {
+    fn leaves_set(&self) -> usize {
         self.next_index
     }
 
@@ -167,6 +172,15 @@ where
             }
         }
     }
+    fn get_empty_leaves_indices(&self) -> Vec<usize> {
+        self.cached_leaves_indices
+            .iter()
+            .take(self.next_index)
+            .enumerate()
+            .filter(|&(_, &v)| v == 0u8)
+            .map(|(idx, _)| idx)
+            .collect()
+    }
 
     // Sets tree nodes, starting from start index
     // Function proper of FullMerkleTree implementation
@@ -185,6 +199,7 @@ where
         }
         hashes.into_iter().for_each(|hash| {
             self.nodes[index + count] = hash;
+            self.cached_leaves_indices[start + count] = 1;
             count += 1;
         });
         if count != 0 {
@@ -194,37 +209,36 @@ where
         Ok(())
     }
 
-    fn override_range<I, J>(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()>
+    fn override_range<I, J>(&mut self, start: usize, leaves: I, indices: J) -> Result<()>
     where
         I: IntoIterator<Item = FrOf<Self::Hasher>>,
         J: IntoIterator<Item = usize>,
     {
-        let index = self.capacity() + start - 1;
-        let mut count = 0;
-        let leaves = leaves.into_iter().collect::<Vec<_>>();
-        let to_remove_indices = to_remove_indices.into_iter().collect::<Vec<_>>();
-        // first count number of hashes, and check that they fit in the tree
-        // then insert into the tree
-        if leaves.len() + start - to_remove_indices.len() > self.capacity() {
-            return Err(Report::msg("provided hashes do not fit in the tree"));
+        let indices = indices.into_iter().collect::<Vec<_>>();
+        let min_index = *indices.first().unwrap();
+        let leaves_vec = leaves.into_iter().collect::<Vec<_>>();
+
+        let max_index = start + leaves_vec.len();
+
+        let mut set_values = vec![Self::Hasher::default_leaf(); max_index - min_index];
+
+        for i in min_index..start {
+            if !indices.contains(&i) {
+                let value = self.get(i)?;
+                set_values[i - min_index] = value;
+            }
         }
 
-        // remove leaves
-        for i in &to_remove_indices {
-            self.delete(*i)?;
+        for i in 0..leaves_vec.len() {
+            set_values[start - min_index + i] = leaves_vec[i];
         }
 
-        // insert new leaves
-        for hash in leaves {
-            self.nodes[index + count] = hash;
-            count += 1;
+        for i in indices {
+            self.cached_leaves_indices[i] = 0;
         }
 
-        if count != 0 {
-            self.update_nodes(index, index + (count - 1))?;
-            self.next_index = max(self.next_index, start + count - to_remove_indices.len());
-        }
-        Ok(())
+        self.set_range(start, set_values)
+            .map_err(|e| Report::msg(e.to_string()))
     }
 
     // Sets a leaf at the next available index
@@ -238,6 +252,7 @@ where
         // 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())?;
+            self.cached_leaves_indices[index] = 0;
         }
         Ok(())
     }

utils/src/merkle_tree/merkle_tree.rs

@@ -47,7 +47,7 @@ pub trait ZerokitMerkleTree {
         Self: Sized;
     fn depth(&self) -> usize;
     fn capacity(&self) -> usize;
-    fn leaves_set(&mut self) -> usize;
+    fn leaves_set(&self) -> usize;
     fn root(&self) -> FrOf<Self::Hasher>;
     fn compute_root(&mut self) -> Result<FrOf<Self::Hasher>>;
     fn get_subtree_root(&self, n: usize, index: usize) -> Result<FrOf<Self::Hasher>>;
@@ -56,6 +56,7 @@ pub trait ZerokitMerkleTree {
     where
         I: IntoIterator<Item = FrOf<Self::Hasher>>;
     fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>>;
+    fn get_empty_leaves_indices(&self) -> Vec<usize>;
     fn override_range<I, J>(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()>
     where
         I: IntoIterator<Item = FrOf<Self::Hasher>>,

utils/src/merkle_tree/optimal_merkle_tree.rs

@@ -27,6 +27,10 @@ where
     /// The tree nodes
     nodes: HashMap<(usize, usize), H::Fr>,
 
+    /// 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>,
+
     // 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,
@@ -78,6 +82,7 @@ where
             cached_nodes: cached_nodes.clone(),
             depth,
             nodes: HashMap::new(),
+            cached_leaves_indices: vec![0; 1 << depth],
             next_index: 0,
             metadata: Vec::new(),
         })
@@ -98,7 +103,7 @@ where
     }
 
     // Returns the total number of leaves set
-    fn leaves_set(&mut self) -> usize {
+    fn leaves_set(&self) -> usize {
         self.next_index
     }
 
@@ -132,6 +137,7 @@ where
         self.nodes.insert((self.depth, index), leaf);
         self.recalculate_from(index)?;
         self.next_index = max(self.next_index, index + 1);
+        self.cached_leaves_indices[index] = 1;
         Ok(())
     }
 
@@ -143,6 +149,16 @@ where
         Ok(self.get_node(self.depth, index))
     }
 
+    fn get_empty_leaves_indices(&self) -> Vec<usize> {
+        self.cached_leaves_indices
+            .iter()
+            .take(self.next_index)
+            .enumerate()
+            .filter(|&(_, &v)| v == 0u8)
+            .map(|(idx, _)| idx)
+            .collect()
+    }
+
     // 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<_>>();
@@ -152,40 +168,43 @@ where
         }
         for (i, leaf) in leaves.iter().enumerate() {
             self.nodes.insert((self.depth, start + i), *leaf);
+            self.cached_leaves_indices[start + i] = 1;
             self.recalculate_from(start + i)?;
         }
         self.next_index = max(self.next_index, start + leaves.len());
         Ok(())
     }
 
-    fn override_range<I, J>(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()>
+    fn override_range<I, J>(&mut self, start: usize, leaves: I, indices: J) -> Result<()>
     where
         I: IntoIterator<Item = FrOf<Self::Hasher>>,
         J: IntoIterator<Item = usize>,
     {
-        let leaves = leaves.into_iter().collect::<Vec<_>>();
-        let to_remove_indices = to_remove_indices.into_iter().collect::<Vec<_>>();
-        // check if the range is valid
-        if leaves.len() + start - to_remove_indices.len() > self.capacity() {
-            return Err(Report::msg("provided range exceeds set size"));
+        let indices = indices.into_iter().collect::<Vec<_>>();
+        let min_index = *indices.first().unwrap();
+        let leaves_vec = leaves.into_iter().collect::<Vec<_>>();
+
+        let max_index = start + leaves_vec.len();
+
+        let mut set_values = vec![Self::Hasher::default_leaf(); max_index - min_index];
+
+        for i in min_index..start {
+            if !indices.contains(&i) {
+                let value = self.get_leaf(i);
+                set_values[i - min_index] = value;
+            }
         }
 
-        // remove leaves
-        for i in &to_remove_indices {
-            self.delete(*i)?;
+        for i in 0..leaves_vec.len() {
+            set_values[start - min_index + i] = leaves_vec[i];
         }
 
-        // add leaves
-        for (i, leaf) in leaves.iter().enumerate() {
-            self.nodes.insert((self.depth, start + i), *leaf);
-            self.recalculate_from(start + i)?;
+        for i in indices {
+            self.cached_leaves_indices[i] = 0;
         }
 
-        self.next_index = max(
-            self.next_index,
-            start + leaves.len() - to_remove_indices.len(),
-        );
-        Ok(())
+        self.set_range(start, set_values)
+            .map_err(|e| Report::msg(e.to_string()))
     }
 
     // Sets a leaf at the next available index
@@ -199,6 +218,7 @@ where
         // 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())?;
+            self.cached_leaves_indices[index] = 0;
         }
         Ok(())
     }
@@ -282,6 +302,7 @@ where
             i >>= 1;
             depth -= 1;
             self.nodes.insert((depth, i), h);
+            self.cached_leaves_indices[index] = 1;
             if depth == 0 {
                 break;
             }

utils/tests/merkle_tree.rs

@@ -2,7 +2,6 @@
 #[cfg(test)]
 pub mod test {
     use hex_literal::hex;
-    use lazy_static::lazy_static;
     use std::{fmt::Display, str::FromStr};
     use tiny_keccak::{Hasher as _, Keccak};
     use zerokit_utils::{
@@ -47,26 +46,15 @@ pub mod test {
         }
     }
 
-    lazy_static! {
-        static ref LEAVES_D2: [TestFr; 4] = [
-            hex!("0000000000000000000000000000000000000000000000000000000000000001"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000002"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000003"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000004"),
-        ]
-        .map(TestFr);
-        static ref LEAVES_D3: [TestFr; 6] = [
-            hex!("0000000000000000000000000000000000000000000000000000000000000001"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000002"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000003"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000004"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000005"),
-            hex!("0000000000000000000000000000000000000000000000000000000000000006"),
-        ]
-        .map(TestFr);
+    impl From<u32> for TestFr {
+        fn from(value: u32) -> Self {
+            let mut bytes: Vec<u8> = vec![0; 28];
+            bytes.extend_from_slice(&value.to_be_bytes());
+            TestFr(bytes.as_slice().try_into().unwrap())
+        }
     }
-    const DEPTH_2: usize = 2;
-    const DEPTH_3: usize = 3;
+
+    const DEFAULT_DEPTH: usize = 2;
 
     fn default_full_merkle_tree(depth: usize) -> FullMerkleTree<Keccak256> {
         FullMerkleTree::<Keccak256>::new(depth, TestFr([0; 32]), FullMerkleConfig::default())
@@ -92,31 +80,128 @@ pub mod test {
         ]
         .map(TestFr);
 
-        let mut tree = default_full_merkle_tree(DEPTH_2);
+        let nof_leaves = 4;
+        let leaves: Vec<TestFr> = (1..=nof_leaves as u32).map(TestFr::from).collect();
+
+        let mut tree = default_full_merkle_tree(DEFAULT_DEPTH);
         assert_eq!(tree.root(), default_tree_root);
-        for i in 0..LEAVES_D2.len() {
-            tree.set(i, LEAVES_D2[i]).unwrap();
+        for i in 0..nof_leaves {
+            tree.set(i, leaves[i]).unwrap();
             assert_eq!(tree.root(), roots[i]);
         }
 
-        let mut tree = default_optimal_merkle_tree(DEPTH_2);
+        let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH);
         assert_eq!(tree.root(), default_tree_root);
-        for i in 0..LEAVES_D2.len() {
-            tree.set(i, LEAVES_D2[i]).unwrap();
+        for i in 0..nof_leaves {
+            tree.set(i, leaves[i]).unwrap();
             assert_eq!(tree.root(), roots[i]);
         }
     }
 
     #[test]
-    fn test_subtree_root() {
-        let mut tree_full = default_optimal_merkle_tree(DEPTH_3);
-        let _ = tree_full.set_range(0, LEAVES_D3.iter().cloned());
+    fn test_get_empty_leaves_indices() {
+        let depth = 4;
+        let nof_leaves: usize = 1 << (depth - 1);
+        let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
+        let leaves_2: Vec<TestFr> = (0u32..2).map(TestFr::from).collect();
+        let leaves_4: Vec<TestFr> = (0u32..4).map(TestFr::from).collect();
 
-        for i in 0..LEAVES_D3.len() {
+        let mut tree_full = default_full_merkle_tree(depth);
+        let _ = tree_full.set_range(0, leaves.clone());
+        assert!(tree_full.get_empty_leaves_indices().is_empty());
+
+        let mut vec_idxs = Vec::new();
+        for i in 0..nof_leaves {
+            vec_idxs.push(i);
+            let _ = tree_full.delete(i);
+            assert_eq!(tree_full.get_empty_leaves_indices(), vec_idxs);
+        }
+
+        for i in (0..nof_leaves).rev() {
+            vec_idxs.pop();
+            let _ = tree_full.set(i, leaves[i]);
+            assert_eq!(tree_full.get_empty_leaves_indices(), vec_idxs);
+        }
+
+        // Check situation when the number of items to insert is less than the number of items to delete
+        tree_full
+            .override_range(0, leaves_2.clone(), [0, 1, 2, 3])
+            .unwrap();
+
+        // check if the indexes for write and delete are the same
+        tree_full
+            .override_range(0, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree_full.get_empty_leaves_indices(), vec![]);
+
+        // check if indexes for deletion are before indexes for overwriting
+        tree_full
+            .override_range(4, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree_full.get_empty_leaves_indices(), vec![0, 1, 2, 3]);
+
+        // check if the indices for write and delete do not overlap completely
+        tree_full
+            .override_range(2, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree_full.get_empty_leaves_indices(), vec![0, 1]);
+
+        //// Optimal Merkle Tree Trest
+
+        let mut tree_opt = default_optimal_merkle_tree(depth);
+        let _ = tree_opt.set_range(0, leaves.clone());
+        assert!(tree_opt.get_empty_leaves_indices().is_empty());
+
+        let mut vec_idxs = Vec::new();
+        for i in 0..nof_leaves {
+            vec_idxs.push(i);
+            let _ = tree_opt.delete(i);
+            assert_eq!(tree_opt.get_empty_leaves_indices(), vec_idxs);
+        }
+        for i in (0..nof_leaves).rev() {
+            vec_idxs.pop();
+            let _ = tree_opt.set(i, leaves[i]);
+            assert_eq!(tree_opt.get_empty_leaves_indices(), vec_idxs);
+        }
+
+        // Check situation when the number of items to insert is less than the number of items to delete
+        tree_opt
+            .override_range(0, leaves_2.clone(), [0, 1, 2, 3])
+            .unwrap();
+
+        // check if the indexes for write and delete are the same
+        tree_opt
+            .override_range(0, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree_opt.get_empty_leaves_indices(), vec![]);
+
+        // check if indexes for deletion are before indexes for overwriting
+        tree_opt
+            .override_range(4, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree_opt.get_empty_leaves_indices(), vec![0, 1, 2, 3]);
+
+        // check if the indices for write and delete do not overlap completely
+        tree_opt
+            .override_range(2, leaves_4.clone(), [0, 1, 2, 3])
+            .unwrap();
+        assert_eq!(tree_opt.get_empty_leaves_indices(), vec![0, 1]);
+    }
+
+    #[test]
+    fn test_subtree_root() {
+        let depth = 3;
+        let nof_leaves: usize = 6;
+        let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
+
+        let mut tree_full = default_optimal_merkle_tree(depth);
+        let _ = tree_full.set_range(0, leaves.iter().cloned());
+
+        for i in 0..nof_leaves {
             // check leaves
             assert_eq!(
                 tree_full.get(i).unwrap(),
-                tree_full.get_subtree_root(DEPTH_3, i).unwrap()
+                tree_full.get_subtree_root(depth, i).unwrap()
             );
 
             // check root
@@ -124,10 +209,10 @@ pub mod test {
         }
 
         // check intermediate nodes
-        for n in (1..=DEPTH_3).rev() {
+        for n in (1..=depth).rev() {
             for i in (0..(1 << n)).step_by(2) {
-                let idx_l = i * (1 << (DEPTH_3 - n));
-                let idx_r = (i + 1) * (1 << (DEPTH_3 - n));
+                let idx_l = i * (1 << (depth - n));
+                let idx_r = (i + 1) * (1 << (depth - n));
                 let idx_sr = idx_l;
 
                 let prev_l = tree_full.get_subtree_root(n, idx_l).unwrap();
@@ -139,24 +224,24 @@ pub mod test {
             }
         }
 
-        let mut tree_opt = default_full_merkle_tree(DEPTH_3);
-        let _ = tree_opt.set_range(0, LEAVES_D3.iter().cloned());
+        let mut tree_opt = default_full_merkle_tree(depth);
+        let _ = tree_opt.set_range(0, leaves.iter().cloned());
 
-        for i in 0..LEAVES_D3.len() {
+        for i in 0..nof_leaves {
             // check leaves
             assert_eq!(
                 tree_opt.get(i).unwrap(),
-                tree_opt.get_subtree_root(DEPTH_3, i).unwrap()
+                tree_opt.get_subtree_root(depth, i).unwrap()
             );
             // check root
             assert_eq!(tree_opt.root(), tree_opt.get_subtree_root(0, i).unwrap());
         }
 
         // check intermediate nodes
-        for n in (1..=DEPTH_3).rev() {
+        for n in (1..=depth).rev() {
             for i in (0..(1 << n)).step_by(2) {
-                let idx_l = i * (1 << (DEPTH_3 - n));
-                let idx_r = (i + 1) * (1 << (DEPTH_3 - n));
+                let idx_l = i * (1 << (depth - n));
+                let idx_r = (i + 1) * (1 << (depth - n));
                 let idx_sr = idx_l;
 
                 let prev_l = tree_opt.get_subtree_root(n, idx_l).unwrap();
@@ -171,11 +256,14 @@ pub mod test {
 
     #[test]
     fn test_proof() {
+        let nof_leaves = 4;
+        let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
+
         // We thest the FullMerkleTree implementation
-        let mut tree = default_full_merkle_tree(DEPTH_2);
-        for i in 0..LEAVES_D2.len() {
+        let mut tree = default_full_merkle_tree(DEFAULT_DEPTH);
+        for i in 0..nof_leaves {
             // We set the leaves
-            tree.set(i, LEAVES_D2[i]).unwrap();
+            tree.set(i, leaves[i]).unwrap();
 
             // We compute a merkle proof
             let proof = tree.proof(i).expect("index should be set");
@@ -184,22 +272,20 @@ pub mod test {
             assert_eq!(proof.leaf_index(), i);
 
             // We verify the proof
-            assert!(tree.verify(&LEAVES_D2[i], &proof).unwrap());
+            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_D2[i]), tree.root());
+            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_D2[(i + 1) % LEAVES_D2.len()], &proof)
-                .unwrap());
+            assert!(!tree.verify(&leaves[(i + 1) % nof_leaves], &proof).unwrap());
         }
 
         // We test the OptimalMerkleTree implementation
-        let mut tree = default_optimal_merkle_tree(DEPTH_2);
-        for i in 0..LEAVES_D2.len() {
+        let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH);
+        for i in 0..nof_leaves {
             // We set the leaves
-            tree.set(i, LEAVES_D2[i]).unwrap();
+            tree.set(i, leaves[i]).unwrap();
 
             // We compute a merkle proof
             let proof = tree.proof(i).expect("index should be set");
@@ -208,24 +294,25 @@ pub mod test {
             assert_eq!(proof.leaf_index(), i);
 
             // We verify the proof
-            assert!(tree.verify(&LEAVES_D2[i], &proof).unwrap());
+            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_D2[i]), tree.root());
+            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_D2[(i + 1) % LEAVES_D2.len()], &proof)
-                .unwrap());
+            assert!(!tree.verify(&leaves[(i + 1) % nof_leaves], &proof).unwrap());
         }
     }
 
     #[test]
     fn test_override_range() {
-        let mut tree = default_optimal_merkle_tree(DEPTH_2);
+        let nof_leaves = 4;
+        let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
+
+        let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH);
 
         // We set the leaves
-        tree.set_range(0, LEAVES_D2.iter().cloned()).unwrap();
+        tree.set_range(0, leaves.iter().cloned()).unwrap();
 
         let new_leaves = [
             hex!("0000000000000000000000000000000000000000000000000000000000000005"),