feat(rln): ability to set leaves from a given index (#63)

* feat(rln): ability to set leaves from a given index

* style(rln): fmt

* fix(rln): bubble the error

* revert: "fix(rln): bubble the error"

This reverts commit 40265082ba.

* fix(rln): replace ark::std::Zero with default

* fix(rln): ordering of ffi api

* fix(rln): use PoseidonHash::default_leaf

* chore(rln): add ffi test for batching

* docs(rln): add entry to changelog

* fix(rln): address reviews

* fix(changelog): fmt

* fix(rln): safe conversion

* fix(rln): reset tree with init func,  add comments

* chore(rln): add comment about init_tree_with_leaves

* fix(rln): comment

CHANGELOG.md

@@ -1,8 +1,24 @@
+## Upcoming release
+
+Release highlights:
+- Allows consumers of zerokit RLN to set leaves to the Merkle Tree from an arbitrary index. Useful for batching updates to the Merkle Tree.
+
+The full list of changes is below.
+
+### Features
+- Creation of `set_leaves_from`, which allows consumers to add leaves to a tree from a given starting index. `init_tree_with_leaves` internally uses `set_leaves_from`, with index 0.
+
+### Changes
+- Renaming of `set_leaves` to `init_tree_with_leaves`, which is a more accurate representation of the function's utility.
+
+### Fixes
+- None
+
 ## 2022-09-19 v0.1
 
 Initial beta release.
 
-This release contain:
+This release contains:
 
 - RLN Module with API to manage, compute and verify [RLN](https://rfc.vac.dev/spec/32/) zkSNARK proofs and RLN primitives.
 - This can be consumed either as a Rust API or as a C FFI. The latter means it can be easily consumed through other environments, such as [Go](https://github.com/status-im/go-zerokit-rln/blob/master/rln/librln.h) or [Nim](https://github.com/status-im/nwaku/blob/4745c7872c69b5fd5c6ddab36df9c5c3d55f57c3/waku/v2/protocol/waku_rln_relay/waku_rln_relay_types.nim).

rln/src/ffi.rs

@@ -104,10 +104,22 @@ pub extern "C" fn set_next_leaf(ctx: *mut RLN, input_buffer: *const Buffer) -> b
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
-pub extern "C" fn set_leaves(ctx: *mut RLN, input_buffer: *const Buffer) -> bool {
+pub extern "C" fn set_leaves_from(
+    ctx: *mut RLN,
+    index: usize,
+    input_buffer: *const Buffer,
+) -> bool {
     let rln = unsafe { &mut *ctx };
     let input_data = <&[u8]>::from(unsafe { &*input_buffer });
-    rln.set_leaves(input_data).is_ok()
+    rln.set_leaves_from(index, input_data).is_ok()
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn init_tree_with_leaves(ctx: *mut RLN, input_buffer: *const Buffer) -> bool {
+    let rln = unsafe { &mut *ctx };
+    let input_data = <&[u8]>::from(unsafe { &*input_buffer });
+    rln.init_tree_with_leaves(input_data).is_ok()
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -387,7 +399,7 @@ mod test {
         // We add leaves in a batch into the tree
         let leaves_ser = vec_fr_to_bytes_le(&leaves);
         let input_buffer = &Buffer::from(leaves_ser.as_ref());
-        let success = set_leaves(rln_pointer, input_buffer);
+        let success = init_tree_with_leaves(rln_pointer, input_buffer);
         assert!(success, "set leaves call failed");
 
         // We get the root of the tree obtained adding leaves in batch
@@ -433,6 +445,95 @@ mod test {
         assert_eq!(root_delete, root_empty);
     }
 
+    #[test]
+    // This test is similar to the one in public.rs but it uses the RLN object as a pointer
+    // Uses `set_leaves_from` to set leaves in a batch
+    fn test_leaf_setting_with_index_ffi() {
+        // We create a new tree
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 256;
+
+        // We create a RLN instance
+        let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
+        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
+        assert!(success, "RLN object creation failed");
+        let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
+
+        // 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 {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // set_index is the index from which we start setting leaves
+        // random number between 0..no_of_leaves
+        let set_index = rng.gen_range(0..no_of_leaves) as usize;
+
+        // We add leaves in a batch into the tree
+        let leaves_ser = vec_fr_to_bytes_le(&leaves);
+        let input_buffer = &Buffer::from(leaves_ser.as_ref());
+        let success = init_tree_with_leaves(rln_pointer, input_buffer);
+        assert!(success, "init tree with leaves call failed");
+
+        // We get the root of the tree obtained adding leaves in batch
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_root(rln_pointer, output_buffer.as_mut_ptr());
+        assert!(success, "get root call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (root_batch_with_init, _) = bytes_le_to_fr(&result_data);
+
+        // `init_tree_with_leaves` resets the tree to the height it was initialized with, using `set_tree`
+
+        // We add leaves in a batch starting from index 0..set_index
+        let leaves_m = vec_fr_to_bytes_le(&leaves[0..set_index]);
+        let buffer = &Buffer::from(leaves_m.as_ref());
+        let success = init_tree_with_leaves(rln_pointer, buffer);
+        assert!(success, "init tree with leaves call failed");
+
+        // We add the remaining n leaves in a batch starting from index set_index
+        let leaves_n = vec_fr_to_bytes_le(&leaves[set_index..]);
+        let buffer = &Buffer::from(leaves_n.as_ref());
+        let success = set_leaves_from(rln_pointer, set_index, buffer);
+        assert!(success, "set leaves from call failed");
+
+        // We get the root of the tree obtained adding leaves in batch
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_root(rln_pointer, output_buffer.as_mut_ptr());
+        assert!(success, "get root call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (root_batch_with_custom_index, _) = bytes_le_to_fr(&result_data);
+
+        assert_eq!(root_batch_with_init, root_batch_with_custom_index);
+
+        // We reset the tree to default
+        let success = set_tree(rln_pointer, tree_height);
+        assert!(success, "set tree call failed");
+
+        // We add leaves one by one using the internal index (new leaves goes in next available position)
+        for leaf in &leaves {
+            let leaf_ser = fr_to_bytes_le(&leaf);
+            let input_buffer = &Buffer::from(leaf_ser.as_ref());
+            let success = set_next_leaf(rln_pointer, input_buffer);
+            assert!(success, "set next leaf call failed");
+        }
+
+        // We get the root of the tree obtained adding leaves using the internal index
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_root(rln_pointer, output_buffer.as_mut_ptr());
+        assert!(success, "get root call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (root_single_additions, _) = bytes_le_to_fr(&result_data);
+
+        assert_eq!(root_batch_with_init, root_single_additions);
+    }
     #[test]
     // This test is similar to the one in lib, but uses only public C API
     fn test_merkle_proof_ffi() {
@@ -739,7 +840,7 @@ mod test {
         // We add leaves in a batch into the tree
         let leaves_ser = vec_fr_to_bytes_le(&leaves);
         let input_buffer = &Buffer::from(leaves_ser.as_ref());
-        let success = set_leaves(rln_pointer, input_buffer);
+        let success = init_tree_with_leaves(rln_pointer, input_buffer);
         assert!(success, "set leaves call failed");
 
         // We generate a new identity pair
@@ -824,7 +925,7 @@ mod test {
         // We add leaves in a batch into the tree
         let leaves_ser = vec_fr_to_bytes_le(&leaves);
         let input_buffer = &Buffer::from(leaves_ser.as_ref());
-        let success = set_leaves(rln_pointer, input_buffer);
+        let success = init_tree_with_leaves(rln_pointer, input_buffer);
         assert!(success, "set leaves call failed");
 
         // We generate a new identity pair

rln/src/public.rs

@@ -118,8 +118,7 @@ impl RLN<'_> {
         Ok(())
     }
 
-    //TODO: change to set_leaves_from(index, input_data)
-    pub fn set_leaves<R: Read>(&mut self, mut input_data: R) -> io::Result<()> {
+    pub fn set_leaves_from<R: Read>(&mut self, index: usize, mut input_data: R) -> io::Result<()> {
         // We read input
         let mut leaves_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut leaves_byte)?;
@@ -128,12 +127,20 @@ impl RLN<'_> {
 
         // We set the leaves
         for (i, leaf) in leaves.iter().enumerate() {
-            self.tree.set(i, *leaf)?;
+            self.tree.set(i + index, *leaf)?;
         }
 
         Ok(())
     }
 
+    pub fn init_tree_with_leaves<R: Read>(&mut self, input_data: R) -> io::Result<()> {
+        // reset the tree
+        // NOTE: this requires the tree to be initialized with the correct height initially
+        // TODO: accept tree_height as a parameter and initialize the tree with that height
+        self.set_tree(self.tree.depth())?;
+        return self.set_leaves_from(0, input_data);
+    }
+
     // Set input leaf to the next available index
     pub fn set_next_leaf<R: Read>(&mut self, mut input_data: R) -> io::Result<()> {
         // We read input
@@ -515,7 +522,7 @@ mod test {
 
         // We add leaves in a batch into the tree
         let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
-        rln.set_leaves(&mut buffer).unwrap();
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
         assert_eq!(rln.tree.leaves_set(), no_of_leaves);
@@ -550,6 +557,80 @@ mod test {
         assert_eq!(root_delete, root_empty);
     }
 
+    #[test]
+    // We test leaf setting with a custom index, to enable batch updates to the root
+    // Uses `set_leaves_from` to set leaves in a batch, from index `start_index`
+    fn test_leaf_setting_with_index() {
+        let tree_height = TEST_TREE_HEIGHT;
+        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 {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // set_index is the index from which we start setting leaves
+        // random number between 0..no_of_leaves
+        let set_index = rng.gen_range(0..no_of_leaves) as usize;
+
+        // We create a new tree
+        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let mut rln = RLN::new(tree_height, input_buffer);
+
+        // We add leaves in a batch into the tree
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
+
+        // We check if number of leaves set is consistent
+        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());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_batch_with_init, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        // `init_tree_with_leaves` resets the tree to the height it was initialized with, using `set_tree`
+
+        // We add leaves in a batch starting from index 0..set_index
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves[0..set_index]));
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
+
+        // We add the remaining n leaves in a batch starting from index m
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves[set_index..]));
+        rln.set_leaves_from(set_index, &mut buffer).unwrap();
+
+        // We check if number of leaves set is consistent
+        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());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_batch_with_custom_index, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        assert_eq!(root_batch_with_init, root_batch_with_custom_index);
+
+        // We reset the tree to default
+        rln.set_tree(tree_height).unwrap();
+
+        // We add leaves one by one using the internal index (new leaves goes in next available position)
+        for leaf in &leaves {
+            let mut buffer = Cursor::new(fr_to_bytes_le(&leaf));
+            rln.set_next_leaf(&mut buffer).unwrap();
+        }
+
+        // We check if numbers of leaves set is consistent
+        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());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_single_additions, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        assert_eq!(root_batch_with_init, root_single_additions);
+    }
+
     #[test]
     // This test is similar to the one in lib, but uses only public API
     fn test_merkle_proof() {
@@ -772,7 +853,7 @@ mod test {
 
         // We add leaves in a batch into the tree
         let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
-        rln.set_leaves(&mut buffer).unwrap();
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // Generate identity pair
         let (identity_secret, id_commitment) = keygen();
@@ -837,7 +918,7 @@ mod test {
 
         // We add leaves in a batch into the tree
         let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
-        rln.set_leaves(&mut buffer).unwrap();
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // Generate identity pair
         let (identity_secret, id_commitment) = keygen();
@@ -982,7 +1063,7 @@ mod test {
 
         // We add leaves in a batch into the tree
         let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
-        rln.set_leaves(&mut buffer).unwrap();
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // Generate identity pair
         let (identity_secret, id_commitment) = keygen();