From d3d85c3e3c042a9fa0a41d1d07e591978c0d39ea Mon Sep 17 00:00:00 2001 From: Ekaterina Broslavskaya Date: Fri, 17 May 2024 23:10:42 +0700 Subject: [PATCH] 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 | 6 + rln/src/pm_tree_adapter.rs | 62 ++++-- rln/src/public.rs | 38 ++++ rln/src/utils.rs | 13 ++ rln/tests/poseidon_tree.rs | 68 ++++++ rln/tests/public.rs | 12 ++ utils/benches/merkle_tree_benchmark.rs | 12 ++ utils/src/merkle_tree/full_merkle_tree.rs | 59 ++++-- utils/src/merkle_tree/merkle_tree.rs | 3 +- utils/src/merkle_tree/optimal_merkle_tree.rs | 59 ++++-- utils/tests/merkle_tree.rs | 205 +++++++++++++------ 11 files changed, 421 insertions(+), 116 deletions(-) diff --git a/rln/benches/pmtree_benchmark.rs b/rln/benches/pmtree_benchmark.rs index 0bfd10b..6fa3d87 100644 --- a/rln/benches/pmtree_benchmark.rs +++ b/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); diff --git a/rln/src/pm_tree_adapter.rs b/rln/src/pm_tree_adapter.rs index 990f353..7901f4f 100644 --- a/rln/src/pm_tree_adapter.rs +++ b/rln/src/pm_tree_adapter.rs @@ -17,6 +17,9 @@ const METADATA_KEY: [u8; 8] = *b"metadata"; pub struct PmTree { tree: pmtree::MerkleTree, + /// 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, // metadata that an application may use to store additional information metadata: Vec, } @@ -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) -> 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>>( @@ -179,9 +185,14 @@ impl ZerokitMerkleTree for PmTree { start: usize, values: I, ) -> Result<()> { + let v = values.into_iter().collect::>(); 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> { @@ -208,6 +219,17 @@ impl ZerokitMerkleTree for PmTree { } } + fn get_empty_leaves_indices(&self) -> Vec { + 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>, J: IntoIterator>( &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 { @@ -282,12 +306,6 @@ type PmTreeHasher = ::Hasher; type FrOfPmTreeHasher = FrOf; impl PmTree { - fn set_range_with_leaves(&mut self, start: usize, leaves: Vec) -> 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(()) } } diff --git a/rln/src/public.rs b/rln/src/public.rs index f20b1d5..54fd037 100644 --- a/rln/src/public.rs +++ b/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 = 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::::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(&self, mut output_data: W) -> Result<()> { + let idxs = self.tree.get_empty_leaves_indices(); + idxs.serialize_compressed(&mut output_data)?; + Ok(()) + } + //////////////////////////////////////////////////////// // zkSNARK APIs //////////////////////////////////////////////////////// diff --git a/rln/src/utils.rs b/rln/src/utils.rs index 93382cb..6b444c9 100644 --- a/rln/src/utils.rs +++ b/rln/src/utils.rs @@ -181,6 +181,19 @@ pub fn normalize_usize(input: usize) -> Vec { normalized_usize } +pub fn bytes_le_to_vec_usize(input: &[u8]) -> Result> { + 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 = 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 { Cursor::new(json!({}).to_string()) diff --git a/rln/tests/poseidon_tree.rs b/rln/tests/poseidon_tree.rs index 50b7e15..69749b1 100644 --- a/rln/tests/poseidon_tree.rs +++ b/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 = (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 = (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 = (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]); + } } diff --git a/rln/tests/public.rs b/rln/tests/public.rs index 8f32307..6458e9a 100644 --- a/rln/tests/public.rs +++ b/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::::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::::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::::new()); rln.get_root(&mut buffer).unwrap(); diff --git a/utils/benches/merkle_tree_benchmark.rs b/utils/benches/merkle_tree_benchmark.rs index 356c2b8..b1adc90 100644 --- a/utils/benches/merkle_tree_benchmark.rs +++ b/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!( diff --git a/utils/src/merkle_tree/full_merkle_tree.rs b/utils/src/merkle_tree/full_merkle_tree.rs index 9b16a24..eea6c10 100644 --- a/utils/src/merkle_tree/full_merkle_tree.rs +++ b/utils/src/merkle_tree/full_merkle_tree.rs @@ -26,6 +26,10 @@ pub struct FullMerkleTree { /// The tree nodes nodes: Vec, + /// 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, + // 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 { + 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(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()> + fn override_range(&mut self, start: usize, leaves: I, indices: J) -> Result<()> where I: IntoIterator>, J: IntoIterator, { - let index = self.capacity() + start - 1; - let mut count = 0; - let leaves = leaves.into_iter().collect::>(); - let to_remove_indices = to_remove_indices.into_iter().collect::>(); - // 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::>(); + let min_index = *indices.first().unwrap(); + let leaves_vec = leaves.into_iter().collect::>(); + + 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(()) } diff --git a/utils/src/merkle_tree/merkle_tree.rs b/utils/src/merkle_tree/merkle_tree.rs index 7b52413..65ad1a7 100644 --- a/utils/src/merkle_tree/merkle_tree.rs +++ b/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; fn compute_root(&mut self) -> Result>; fn get_subtree_root(&self, n: usize, index: usize) -> Result>; @@ -56,6 +56,7 @@ pub trait ZerokitMerkleTree { where I: IntoIterator>; fn get(&self, index: usize) -> Result>; + fn get_empty_leaves_indices(&self) -> Vec; fn override_range(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()> where I: IntoIterator>, diff --git a/utils/src/merkle_tree/optimal_merkle_tree.rs b/utils/src/merkle_tree/optimal_merkle_tree.rs index d21d2a0..6ee4578 100644 --- a/utils/src/merkle_tree/optimal_merkle_tree.rs +++ b/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, + // 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 { + 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>(&mut self, start: usize, leaves: I) -> Result<()> { let leaves = leaves.into_iter().collect::>(); @@ -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(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()> + fn override_range(&mut self, start: usize, leaves: I, indices: J) -> Result<()> where I: IntoIterator>, J: IntoIterator, { - let leaves = leaves.into_iter().collect::>(); - let to_remove_indices = to_remove_indices.into_iter().collect::>(); - // 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::>(); + let min_index = *indices.first().unwrap(); + let leaves_vec = leaves.into_iter().collect::>(); + + 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; } diff --git a/utils/tests/merkle_tree.rs b/utils/tests/merkle_tree.rs index 35733ad..00f46a6 100644 --- a/utils/tests/merkle_tree.rs +++ b/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 for TestFr { + fn from(value: u32) -> Self { + let mut bytes: Vec = 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 { FullMerkleTree::::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 = (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 = (0..nof_leaves as u32).map(TestFr::from).collect(); + let leaves_2: Vec = (0u32..2).map(TestFr::from).collect(); + let leaves_4: Vec = (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 = (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 = (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 = (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"),