From 0e40b7299b383dc4462ef747984a2187d928f1c2 Mon Sep 17 00:00:00 2001 From: seemenkina Date: Thu, 16 May 2024 19:38:38 +0700 Subject: [PATCH] add function for empty leaves --- rln/benches/pmtree_benchmark.rs | 6 + rln/src/pm_tree_adapter.rs | 59 +++++-- rln/src/public.rs | 38 +++++ rln/src/utils.rs | 13 ++ rln/tests/poseidon_tree.rs | 62 ++++++- rln/tests/public.rs | 12 ++ utils/benches/merkle_tree_benchmark.rs | 12 ++ utils/src/merkle_tree/full_merkle_tree.rs | 20 ++- utils/src/merkle_tree/merkle_tree.rs | 3 +- utils/src/merkle_tree/optimal_merkle_tree.rs | 22 ++- utils/tests/merkle_tree.rs | 171 +++++++++++++++---- 11 files changed, 362 insertions(+), 56 deletions(-) diff --git a/rln/benches/pmtree_benchmark.rs b/rln/benches/pmtree_benchmark.rs index 4a2df1a3..922e6783 100644 --- a/rln/benches/pmtree_benchmark.rs +++ b/rln/benches/pmtree_benchmark.rs @@ -37,6 +37,12 @@ pub fn pmtree_benchmark(c: &mut Criterion) { tree.get(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 3007c7d8..7118bad5 100644 --- a/rln/src/pm_tree_adapter.rs +++ b/rln/src/pm_tree_adapter.rs @@ -16,6 +16,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, } @@ -143,6 +146,7 @@ impl ZerokitMerkleTree for PmTree { Ok(PmTree { tree, + cached_leaves_indices: vec![0; 1 << depth], metadata: Vec::new(), }) } @@ -155,7 +159,7 @@ impl ZerokitMerkleTree for PmTree { self.tree.capacity() } - fn leaves_set(&mut self) -> usize { + fn leaves_set(&self) -> usize { self.tree.leaves_set() } @@ -170,7 +174,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>>( @@ -178,15 +184,31 @@ 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> { self.tree.get(index).map_err(|e| Report::msg(e.to_string())) } + 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, @@ -201,7 +223,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), } @@ -216,7 +238,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 { @@ -261,12 +285,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; @@ -275,7 +293,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( @@ -302,7 +325,15 @@ impl PmTree { self.tree .set_range(min_index, set_values) - .map_err(|e| Report::msg(e.to_string())) + .map_err(|e| Report::msg(e.to_string()))?; + + for i in indices { + self.cached_leaves_indices[*i] = 0; + } + for i in min_index..(max_index - min_index) { + self.cached_leaves_indices[i] = 1 + } + Ok(()) } } diff --git a/rln/src/public.rs b/rln/src/public.rs index 1d4ba3e4..bdfd38f6 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -576,6 +576,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 93382cb2..fcc3c934 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 c6bb5d9f..2ece65c5 100644 --- a/rln/tests/poseidon_tree.rs +++ b/rln/tests/poseidon_tree.rs @@ -4,12 +4,12 @@ #[cfg(test)] mod test { - use rln::circuit::*; use rln::hashers::PoseidonHash; + use rln::{circuit::*, poseidon_tree::PoseidonTree}; use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree}; #[test] - /// The test is checked correctness for `FullMerkleTree` and `OptimalMerkleTree` with Poseidon hash + // The test is checked correctness for `FullMerkleTree` and `OptimalMerkleTree` with Poseidon hash fn test_zerokit_merkle_implementations() { let sample_size = 100; let leaves: Vec = (0..sample_size).map(|s| Fr::from(s)).collect(); @@ -33,4 +33,62 @@ mod test { assert_eq!(tree_full_root, tree_opt_root); } + + #[test] + fn test_subtree_root() { + 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(); + // assert_eq!(tree.get_empty_leaves_indices(), vec![0, 1, 2, 3]); + + // check if the indices for write and delete do not overlap completely + // tree.override_range(2, leaves_4.clone(), [0, 1, 2, 3]) + // .unwrap(); + // assert_eq!(tree.get_empty_leaves_indices(), vec![0, 1]); + } } diff --git a/rln/tests/public.rs b/rln/tests/public.rs index 5b8dde38..8e77946c 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 64d81ef7..b24ddc5b 100644 --- a/utils/benches/merkle_tree_benchmark.rs +++ b/utils/benches/merkle_tree_benchmark.rs @@ -90,6 +90,12 @@ pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) { tree.get(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) { @@ -125,6 +131,12 @@ pub fn full_merkle_tree_benchmark(c: &mut Criterion) { tree.get(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 15ef4220..76e81263 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 } @@ -141,6 +146,16 @@ where Ok(self.nodes[self.capacity() + leaf - 1]) } + 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 fn set_range>>( @@ -158,6 +173,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 { @@ -190,6 +206,7 @@ where // insert new leaves for hash in leaves { self.nodes[index + count] = hash; + self.cached_leaves_indices[start + count] = 1; count += 1; } @@ -211,6 +228,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 17d107bd..147a3cf0 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 set(&mut self, index: usize, leaf: FrOf) -> Result<()>; @@ -55,6 +55,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 48470e7d..5798a410 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 } @@ -116,6 +121,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(()) } @@ -127,6 +133,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::>(); @@ -136,6 +152,7 @@ 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()); @@ -162,6 +179,7 @@ where // add leaves 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)?; } @@ -183,6 +201,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(()) } @@ -266,6 +285,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 8d491869..f15ed3cc 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,25 +46,118 @@ pub mod test { } } - lazy_static! { - static ref LEAVES: [TestFr; 4] = [ - hex!("0000000000000000000000000000000000000000000000000000000000000001"), - hex!("0000000000000000000000000000000000000000000000000000000000000002"), - hex!("0000000000000000000000000000000000000000000000000000000000000003"), - hex!("0000000000000000000000000000000000000000000000000000000000000004"), - ] - .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()) + } } - fn default_full_merkle_tree() -> FullMerkleTree { - FullMerkleTree::::new(2, TestFr([0; 32]), FullMerkleConfig::default()).unwrap() + const DEFAULT_DEPTH: usize = 2; + + fn default_full_merkle_tree(depth: usize) -> FullMerkleTree { + FullMerkleTree::::new(depth, TestFr([0; 32]), FullMerkleConfig::default()) + .unwrap() } - fn default_optimal_merkle_tree() -> OptimalMerkleTree { - OptimalMerkleTree::::new(2, TestFr([0; 32]), OptimalMerkleConfig::default()) + fn default_optimal_merkle_tree(depth: usize) -> OptimalMerkleTree { + OptimalMerkleTree::::new(depth, TestFr([0; 32]), OptimalMerkleConfig::default()) .unwrap() } + #[test] + 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(); + + 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(); + + assert!(tree_full.get_empty_leaves_indices().is_empty()); + // 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(); + + assert!(tree_opt.get_empty_leaves_indices().is_empty()); + // 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_root() { let default_tree_root = TestFr(hex!( @@ -80,28 +172,34 @@ pub mod test { ] .map(TestFr); - let mut tree = default_full_merkle_tree(); + 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.len() { - tree.set(i, LEAVES[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(); + let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH); assert_eq!(tree.root(), default_tree_root); - for i in 0..LEAVES.len() { - tree.set(i, LEAVES[i]).unwrap(); + for i in 0..nof_leaves { + tree.set(i, leaves[i]).unwrap(); assert_eq!(tree.root(), roots[i]); } } #[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(); - for i in 0..LEAVES.len() { + let mut tree = default_full_merkle_tree(DEFAULT_DEPTH); + for i in 0..nof_leaves { // We set the leaves - tree.set(i, LEAVES[i]).unwrap(); + tree.set(i, leaves[i]).unwrap(); // We compute a merkle proof let proof = tree.proof(i).expect("index should be set"); @@ -110,22 +208,20 @@ pub mod test { assert_eq!(proof.leaf_index(), i); // We verify the proof - assert!(tree.verify(&LEAVES[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[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[(i + 1) % LEAVES.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(); - for i in 0..LEAVES.len() { + let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH); + for i in 0..nof_leaves { // We set the leaves - tree.set(i, LEAVES[i]).unwrap(); + tree.set(i, leaves[i]).unwrap(); // We compute a merkle proof let proof = tree.proof(i).expect("index should be set"); @@ -134,24 +230,25 @@ pub mod test { assert_eq!(proof.leaf_index(), i); // We verify the proof - assert!(tree.verify(&LEAVES[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[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[(i + 1) % LEAVES.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(); + 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.iter().cloned()).unwrap(); + tree.set_range(0, leaves.iter().cloned()).unwrap(); let new_leaves = [ hex!("0000000000000000000000000000000000000000000000000000000000000005"),