From 8f2c9e3586a5148ffd12c28f4059feb975781b9e Mon Sep 17 00:00:00 2001 From: Aaryamann Challani <43716372+rymnc@users.noreply.github.com> Date: Mon, 15 May 2023 07:11:43 +0530 Subject: [PATCH] combined batch operations (insert + delete) (#160) * fix(rln): clippy error * feat: batch ops in ZerokitMerkleTree * chore: bump pmtree * fix: upstream root calc --- Cargo.lock | 2 +- rln/src/pm_tree_adapter.rs | 46 ++++++++++++++++++++ utils/Cargo.toml | 2 +- utils/src/merkle_tree/full_merkle_tree.rs | 45 +++++++++++++++++++ utils/src/merkle_tree/merkle_tree.rs | 6 +++ utils/src/merkle_tree/optimal_merkle_tree.rs | 44 +++++++++++++++++++ utils/src/pm_tree/sled_adapter.rs | 5 +-- utils/tests/merkle_tree.rs | 37 ++++++++++++++++ 8 files changed, 182 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3d08f7e..793eb79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2110,7 +2110,7 @@ dependencies = [ [[package]] name = "pmtree" version = "1.0.0" -source = "git+https://github.com/Rate-Limiting-Nullifier/pmtree?rev=b3a02216cece3e9c24e1754ea381bf784fd1df48#b3a02216cece3e9c24e1754ea381bf784fd1df48" +source = "git+https://github.com/Rate-Limiting-Nullifier/pmtree?rev=bf27d4273b34a7f4ec3c77cdf67fb17a5602504a#bf27d4273b34a7f4ec3c77cdf67fb17a5602504a" dependencies = [ "rayon", ] diff --git a/rln/src/pm_tree_adapter.rs b/rln/src/pm_tree_adapter.rs index f7a3b84..e9bb730 100644 --- a/rln/src/pm_tree_adapter.rs +++ b/rln/src/pm_tree_adapter.rs @@ -3,9 +3,11 @@ use crate::hashers::{poseidon_hash, PoseidonHash}; use crate::utils::{bytes_le_to_fr, fr_to_bytes_le}; use color_eyre::{Report, Result}; use serde_json::Value; +use std::collections::HashSet; use std::fmt::Debug; use std::path::PathBuf; use std::str::FromStr; +use utils::pmtree::Hasher; use utils::*; pub struct PmTree { @@ -127,6 +129,10 @@ impl ZerokitMerkleTree for PmTree { .map_err(|e| Report::msg(e.to_string())) } + fn get(&self, index: usize) -> Result> { + self.tree.get(index).map_err(|e| Report::msg(e.to_string())) + } + fn set_range>>( &mut self, start: usize, @@ -137,6 +143,42 @@ impl ZerokitMerkleTree for PmTree { .map_err(|e| Report::msg(e.to_string())) } + fn override_range>, J: IntoIterator>( + &mut self, + start: usize, + leaves: I, + indices: J, + ) -> Result<()> { + let leaves = leaves.into_iter().collect::>(); + let indices = indices.into_iter().collect::>(); + let end = start + leaves.len(); + + if leaves.len() + start - indices.len() > self.capacity() { + return Err(Report::msg("index out of bounds")); + } + + // extend the range to include indices to be removed + let min_index = indices.iter().min().unwrap_or(&start); + let max_index = indices.iter().max().unwrap_or(&end); + + let mut new_leaves = Vec::new(); + + // insert leaves into new_leaves + for i in *min_index..*max_index { + if indices.contains(&i) { + // insert 0 + new_leaves.push(Self::Hasher::default_leaf()); + } else { + // insert leaf + new_leaves.push(leaves[i - start]); + } + } + + self.tree + .set_range(start, new_leaves) + .map_err(|e| Report::msg(e.to_string())) + } + fn update_next(&mut self, leaf: FrOf) -> Result<()> { self.tree .update_next(leaf) @@ -161,6 +203,10 @@ impl ZerokitMerkleTree for PmTree { Err(Report::msg("verify failed")) } } + + fn compute_root(&mut self) -> Result> { + Ok(self.tree.root()) + } } impl ZerokitMerkleProof for PmTreeProof { diff --git a/utils/Cargo.toml b/utils/Cargo.toml index 462a02a..02416ed 100644 --- a/utils/Cargo.toml +++ b/utils/Cargo.toml @@ -8,7 +8,7 @@ license = "MIT OR Apache-2.0" ark-ff = { version = "=0.4.1", default-features = false, features = ["asm"] } num-bigint = { version = "=0.4.3", default-features = false, features = ["rand"] } color-eyre = "=0.6.2" -pmtree = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree", rev = "b3a02216cece3e9c24e1754ea381bf784fd1df48", optional = true} +pmtree = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree", rev = "bf27d4273b34a7f4ec3c77cdf67fb17a5602504a", optional = true} sled = "=0.34.7" serde = "1.0.44" diff --git a/utils/src/merkle_tree/full_merkle_tree.rs b/utils/src/merkle_tree/full_merkle_tree.rs index 5a248ff..cad0a34 100644 --- a/utils/src/merkle_tree/full_merkle_tree.rs +++ b/utils/src/merkle_tree/full_merkle_tree.rs @@ -125,6 +125,14 @@ where Ok(()) } + // Get a leaf from the specified tree index + fn get(&self, leaf: usize) -> Result> { + if leaf >= self.capacity() { + return Err(Report::msg("leaf index out of bounds")); + } + Ok(self.nodes[self.capacity() + leaf - 1]) + } + // Sets tree nodes, starting from start index // Function proper of FullMerkleTree implementation fn set_range>>( @@ -151,6 +159,39 @@ where Ok(()) } + fn override_range(&mut self, start: usize, leaves: I, to_remove_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")); + } + + // remove leaves + for i in &to_remove_indices { + self.delete(*i)?; + } + + // insert new leaves + for hash in leaves { + self.nodes[index + count] = hash; + count += 1; + } + + if count != 0 { + self.update_nodes(index, index + (count - 1))?; + self.next_index = max(self.next_index, start + count - to_remove_indices.len()); + } + Ok(()) + } + // Sets a leaf at the next available index fn update_next(&mut self, leaf: FrOf) -> Result<()> { self.set(self.next_index, leaf)?; @@ -189,6 +230,10 @@ where fn verify(&self, hash: &FrOf, proof: &FullMerkleProof) -> Result { Ok(proof.compute_root_from(hash) == self.root()) } + + fn compute_root(&mut self) -> Result> { + Ok(self.root()) + } } impl FullMerkleTree diff --git a/utils/src/merkle_tree/merkle_tree.rs b/utils/src/merkle_tree/merkle_tree.rs index d2a112f..a6e6fed 100644 --- a/utils/src/merkle_tree/merkle_tree.rs +++ b/utils/src/merkle_tree/merkle_tree.rs @@ -49,10 +49,16 @@ pub trait ZerokitMerkleTree { fn capacity(&self) -> usize; fn leaves_set(&mut self) -> usize; fn root(&self) -> FrOf; + fn compute_root(&mut self) -> Result>; fn set(&mut self, index: usize, leaf: FrOf) -> Result<()>; fn set_range(&mut self, start: usize, leaves: I) -> Result<()> where I: IntoIterator>; + fn get(&self, index: usize) -> Result>; + fn override_range(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()> + where + I: IntoIterator>, + J: IntoIterator; fn update_next(&mut self, leaf: FrOf) -> Result<()>; fn delete(&mut self, index: usize) -> Result<()>; fn proof(&self, index: usize) -> Result; diff --git a/utils/src/merkle_tree/optimal_merkle_tree.rs b/utils/src/merkle_tree/optimal_merkle_tree.rs index 989da97..3d39412 100644 --- a/utils/src/merkle_tree/optimal_merkle_tree.rs +++ b/utils/src/merkle_tree/optimal_merkle_tree.rs @@ -1,4 +1,5 @@ use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree}; +use crate::FrOf; use color_eyre::{Report, Result}; use std::collections::HashMap; use std::str::FromStr; @@ -110,6 +111,14 @@ where Ok(()) } + // Get a leaf from the specified tree index + fn get(&self, index: usize) -> Result { + if index >= self.capacity() { + return Err(Report::msg("index exceeds set size")); + } + Ok(self.get_node(self.depth, index)) + } + // Sets multiple leaves from the specified tree index fn set_range>(&mut self, start: usize, leaves: I) -> Result<()> { let leaves = leaves.into_iter().collect::>(); @@ -125,6 +134,36 @@ where Ok(()) } + fn override_range(&mut self, start: usize, leaves: I, to_remove_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")); + } + + // remove leaves + for i in &to_remove_indices { + self.delete(*i)?; + } + + // add leaves + for (i, leaf) in leaves.iter().enumerate() { + self.nodes.insert((self.depth, start + i), *leaf); + self.recalculate_from(start + i)?; + } + + self.next_index = max( + self.next_index, + start + leaves.len() - to_remove_indices.len(), + ); + Ok(()) + } + // Sets a leaf at the next available index fn update_next(&mut self, leaf: H::Fr) -> Result<()> { self.set(self.next_index, leaf)?; @@ -172,6 +211,11 @@ where let expected_root = witness.compute_root_from(leaf); Ok(expected_root.eq(&self.root())) } + + fn compute_root(&mut self) -> Result> { + self.recalculate_from(0)?; + Ok(self.root()) + } } impl OptimalMerkleTree diff --git a/utils/src/pm_tree/sled_adapter.rs b/utils/src/pm_tree/sled_adapter.rs index ba155da..c368b32 100644 --- a/utils/src/pm_tree/sled_adapter.rs +++ b/utils/src/pm_tree/sled_adapter.rs @@ -14,8 +14,7 @@ impl Database for SledDB { Err(e) => { return Err(PmtreeErrorKind::DatabaseError( DatabaseErrorKind::CustomError(format!( - "Cannot create database: {} {:#?}", - e, config + "Cannot create database: {e} {config:#?}", )), )) } @@ -29,7 +28,7 @@ impl Database for SledDB { Ok(db) => db, Err(e) => { return Err(PmtreeErrorKind::DatabaseError( - DatabaseErrorKind::CustomError(format!("Cannot load database: {}", e)), + DatabaseErrorKind::CustomError(format!("Cannot load database: {e}")), )) } }; diff --git a/utils/tests/merkle_tree.rs b/utils/tests/merkle_tree.rs index 1c99046..dd9a500 100644 --- a/utils/tests/merkle_tree.rs +++ b/utils/tests/merkle_tree.rs @@ -125,4 +125,41 @@ mod test { .unwrap()); } } + + #[test] + fn test_override_range() { + let initial_leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + hex!("0000000000000000000000000000000000000000000000000000000000000003"), + hex!("0000000000000000000000000000000000000000000000000000000000000004"), + ]; + + let mut tree = + OptimalMerkleTree::::new(2, [0; 32], OptimalMerkleConfig::default()) + .unwrap(); + + // We set the leaves + tree.set_range(0, initial_leaves.iter().cloned()).unwrap(); + + let new_leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000005"), + hex!("0000000000000000000000000000000000000000000000000000000000000006"), + ]; + + let to_delete_indices: [usize; 2] = [0, 1]; + + // We override the leaves + tree.override_range( + 0, // start from the end of the initial leaves + new_leaves.iter().cloned(), + to_delete_indices.iter().cloned(), + ) + .unwrap(); + + // ensure that the leaves are set correctly + for i in 0..new_leaves.len() { + assert_eq!(tree.get_leaf(i), new_leaves[i]); + } + } }