From 490206aa440fb524f0cdf47019ea9604466aae75 Mon Sep 17 00:00:00 2001 From: Aaryamann Challani <43716372+rymnc@users.noreply.github.com> Date: Thu, 3 Nov 2022 17:49:30 +0530 Subject: [PATCH] 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 40265082ba888f0a583c4e4d0b7abab9da9b7da8. * 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 | 18 +++++++- rln/src/ffi.rs | 111 +++++++++++++++++++++++++++++++++++++++++++--- rln/src/public.rs | 95 ++++++++++++++++++++++++++++++++++++--- 3 files changed, 211 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 94532cd..f46aa4a 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/rln/src/ffi.rs b/rln/src/ffi.rs index fef6342..4e12f87 100644 --- a/rln/src/ffi.rs +++ b/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 = 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::::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::::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::::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 diff --git a/rln/src/public.rs b/rln/src/public.rs index 48b6450..ad35f4e 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -118,8 +118,7 @@ impl RLN<'_> { Ok(()) } - //TODO: change to set_leaves_from(index, input_data) - pub fn set_leaves(&mut self, mut input_data: R) -> io::Result<()> { + pub fn set_leaves_from(&mut self, index: usize, mut input_data: R) -> io::Result<()> { // We read input let mut leaves_byte: Vec = 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(&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(&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 = 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::::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::::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::::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();