feat(rln): public, ffi for atomic ops (#162)

rln/src/ffi.rs

@@ -244,6 +244,17 @@ pub extern "C" fn init_tree_with_leaves(ctx: *mut RLN, input_buffer: *const Buff
     call!(ctx, init_tree_with_leaves, input_buffer)
 }
 
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn atomic_operation(
+    ctx: *mut RLN,
+    index: usize,
+    leaves_buffer: *const Buffer,
+    indices_buffer: *const Buffer,
+) -> bool {
+    call!(ctx, atomic_operation, index, leaves_buffer, indices_buffer)
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
 pub extern "C" fn get_root(ctx: *const RLN, output_buffer: *mut Buffer) -> bool {

rln/src/public.rs

@@ -256,7 +256,7 @@ impl RLN<'_> {
 
         // We set the leaves
         self.tree
-            .set_range(index, leaves)
+            .override_range(index, leaves, [])
             .map_err(|_| Report::msg("Could not set leaves"))?;
         Ok(())
     }
@@ -275,6 +275,72 @@ impl RLN<'_> {
         self.set_leaves_from(0, input_data)
     }
 
+    /// Sets multiple leaves starting from position index in the internal Merkle tree.
+    /// Also accepts an array of indices to remove from the tree.
+    ///
+    /// If n leaves are passed as input, these will be set at positions `index`, `index+1`, ..., `index+n-1` respectively.
+    /// If m indices are passed as input, these will be removed from the tree.
+    ///
+    /// This function updates the internal Merkle tree `next_index value indicating the next available index corresponding to a never-set leaf as `next_index = max(next_index, index + n)`.
+    ///
+    /// Input values are:
+    /// - `index`: the index of the first leaf to be set
+    /// - `input_leaves`: a reader for the serialization of multiple leaf values (serialization done with [`rln::utils::vec_fr_to_bytes_le`](crate::utils::vec_fr_to_bytes_le))
+    /// - `input_indices`: a reader for the serialization of multiple indices to remove (serialization done with [`rln::utils::vec_u8_to_bytes_le`](crate::utils::vec_u8_to_bytes_le))
+    ///
+    /// Example:
+    /// ```
+    /// use rln::circuit::Fr;
+    /// use rln::utils::*;
+    ///
+    /// let start_index = 10;
+    /// 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();
+    ///     leaves.push(id_commitment);
+    /// }
+    ///
+    /// let mut indices: Vec<u8> = Vec::new();
+    /// for i in 0..no_of_leaves {
+    ///    if i % 2 == 0 {
+    ///       indices.push(i as u8);
+    ///   }
+    /// }
+    ///
+    /// // We atomically add leaves and remove indices from the tree
+    /// let mut leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
+    /// let mut indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices));
+    /// rln.set_leaves_from(index, &mut leaves_buffer, indices_buffer).unwrap();
+    /// ```
+    pub fn atomic_operation<R: Read>(
+        &mut self,
+        index: usize,
+        mut input_leaves: R,
+        mut input_indices: R,
+    ) -> Result<()> {
+        // We read input
+        let mut leaves_byte: Vec<u8> = Vec::new();
+        input_leaves.read_to_end(&mut leaves_byte)?;
+
+        let (leaves, _) = bytes_le_to_vec_fr(&leaves_byte)?;
+
+        let mut indices_byte: Vec<u8> = Vec::new();
+        input_indices.read_to_end(&mut indices_byte)?;
+
+        let (indices, _) = bytes_le_to_vec_u8(&indices_byte)?;
+        let indices: Vec<usize> = indices.iter().map(|x| *x as usize).collect();
+
+        // We set the leaves
+        self.tree
+            .override_range(index, leaves, indices)
+            .map_err(|_| Report::msg("Could not perform the batch operation"))?;
+        Ok(())
+    }
+
     /// Sets a leaf value at the next available never-set leaf index.
     ///
     /// This function updates the internal Merkle tree `next_index` value indicating the next available index corresponding to a never-set leaf as `next_index = next_index + 1`.
@@ -1252,6 +1318,57 @@ mod test {
         assert_eq!(root_batch_with_init, root_single_additions);
     }
 
+    #[test]
+    // Tests the atomic_operation fn, which set_leaves_from uses internally
+    fn test_atomic_operation() {
+        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));
+        }
+
+        // We create a new tree
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        // We add leaves in a batch into the tree
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).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);
+
+        // 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_after_insertion, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        let last_leaf = leaves.last().unwrap();
+        let last_leaf_index = no_of_leaves - 1;
+        let indices = vec![last_leaf_index as u8];
+        let last_leaf = vec![*last_leaf];
+        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices).unwrap());
+        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&last_leaf).unwrap());
+
+        rln.atomic_operation(no_of_leaves, leaves_buffer, indices_buffer)
+            .unwrap();
+
+        // We get the root of the tree obtained after a no-op
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_noop, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        assert_eq!(root_after_insertion, root_after_noop);
+    }
+
     #[allow(unused_must_use)]
     #[test]
     // This test checks if `set_leaves_from` throws an error when the index is out of bounds

rln/tests/ffi.rs

@@ -218,6 +218,71 @@ mod test {
         assert_eq!(root_batch_with_init, root_single_additions);
     }
 
+    #[test]
+    // This test is similar to the one in public.rs but it uses the RLN object as a pointer
+    fn test_atomic_operation_ffi() {
+        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_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.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));
+        }
+
+        // We add leaves in a batch into the tree
+        let leaves_ser = vec_fr_to_bytes_le(&leaves).unwrap();
+        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_after_insertion, _) = bytes_le_to_fr(&result_data);
+
+        let last_leaf = leaves.last().unwrap();
+        let last_leaf_index = no_of_leaves - 1;
+        let indices = vec![last_leaf_index as u8];
+        let last_leaf = vec![*last_leaf];
+        let indices = vec_u8_to_bytes_le(&indices).unwrap();
+        let indices_buffer = &Buffer::from(indices.as_ref());
+        let leaves = vec_fr_to_bytes_le(&last_leaf).unwrap();
+        let leaves_buffer = &Buffer::from(leaves.as_ref());
+
+        let success = atomic_operation(
+            rln_pointer,
+            no_of_leaves as usize,
+            leaves_buffer,
+            indices_buffer,
+        );
+        assert!(success, "atomic operation call failed");
+
+        // We get the root of the tree obtained after a no-op
+        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_after_noop, _) = bytes_le_to_fr(&result_data);
+
+        assert_eq!(root_after_insertion, root_after_noop);
+    }
+
     #[test]
     // This test is similar to the one in public.rs but it uses the RLN object as a pointer
     fn test_set_leaves_bad_index_ffi() {