From 0005b1d61fa44255030619e3b2efcc9fbbbf7a56 Mon Sep 17 00:00:00 2001 From: Ekaterina Broslavskaya Date: Fri, 17 May 2024 16:35:18 +0700 Subject: [PATCH] Expose a public function to fetch the root of a subtree at level n (#247) * add get_subroot function * update test * update pmtree dependecy --- Cargo.lock | 21 +-- rln/benches/pmtree_benchmark.rs | 7 + rln/src/pm_tree_adapter.rs | 21 +++ rln/src/public.rs | 27 ++++ rln/tests/poseidon_tree.rs | 41 +++++- rln/tests/public.rs | 23 ++++ utils/Cargo.toml | 3 +- utils/benches/merkle_tree_benchmark.rs | 14 ++ utils/src/merkle_tree/full_merkle_tree.rs | 27 ++++ utils/src/merkle_tree/merkle_tree.rs | 1 + utils/src/merkle_tree/optimal_merkle_tree.rs | 16 +++ utils/tests/merkle_tree.rs | 130 +++++++++++++++---- 12 files changed, 289 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0ccdaed..d73092c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2041,15 +2041,6 @@ dependencies = [ "plotters-backend", ] -[[package]] -name = "pmtree" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e054322ee96d2ccd86cd47b87797166682e45f5d67571c48eaa864668d26f510" -dependencies = [ - "rayon", -] - [[package]] name = "ppv-lite86" version = "0.2.17" @@ -3003,6 +2994,15 @@ version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4dad5567ad0cf5b760e5665964bec1b47dfd077ba8a2544b513f3556d3d239a2" +[[package]] +name = "vacp2p_pmtree" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "632293f506ca10d412dbe1d427295317b4c794fa9ddfd66fbd2fa971de88c1f6" +dependencies = [ + "rayon", +] + [[package]] name = "valuable" version = "0.1.0" @@ -3595,12 +3595,13 @@ dependencies = [ "ark-ff", "color-eyre", "criterion 0.4.0", + "hex", "hex-literal", "lazy_static 1.4.0", "num-bigint", "num-traits", - "pmtree", "serde", "sled", "tiny-keccak", + "vacp2p_pmtree", ] diff --git a/rln/benches/pmtree_benchmark.rs b/rln/benches/pmtree_benchmark.rs index 4a2df1a..0bfd10b 100644 --- a/rln/benches/pmtree_benchmark.rs +++ b/rln/benches/pmtree_benchmark.rs @@ -37,6 +37,13 @@ pub fn pmtree_benchmark(c: &mut Criterion) { tree.get(0).unwrap(); }) }); + + // check intermediate node getter which required additional computation of sub root index + c.bench_function("Pmtree::get_subtree_root", |b| { + b.iter(|| { + tree.get_subtree_root(1, 0).unwrap(); + }) + }); } criterion_group!(benches, pmtree_benchmark); diff --git a/rln/src/pm_tree_adapter.rs b/rln/src/pm_tree_adapter.rs index 3007c7d..990f353 100644 --- a/rln/src/pm_tree_adapter.rs +++ b/rln/src/pm_tree_adapter.rs @@ -5,6 +5,7 @@ use std::str::FromStr; use color_eyre::{Report, Result}; use serde_json::Value; +use utils::pmtree::tree::Key; use utils::pmtree::{Database, Hasher}; use utils::*; @@ -187,6 +188,26 @@ impl ZerokitMerkleTree for PmTree { self.tree.get(index).map_err(|e| Report::msg(e.to_string())) } + fn get_subtree_root(&self, n: usize, index: usize) -> Result> { + if n > self.depth() { + return Err(Report::msg("level exceeds depth size")); + } + if index >= self.capacity() { + return Err(Report::msg("index exceeds set size")); + } + if n == 0 { + Ok(self.root()) + } else if n == self.depth() { + self.get(index) + } else { + let node = self + .tree + .get_elem(Key::new(n, index >> (self.depth() - n))) + .unwrap(); + Ok(node) + } + } + fn override_range>, J: IntoIterator>( &mut self, start: usize, diff --git a/rln/src/public.rs b/rln/src/public.rs index 1d4ba3e..f20b1d5 100644 --- a/rln/src/public.rs +++ b/rln/src/public.rs @@ -544,6 +544,33 @@ impl RLN<'_> { Ok(()) } + /// Returns the root of subtree in the Merkle tree + /// + /// Output values are: + /// - `output_data`: a writer receiving the serialization of the node value (serialization done with [`rln::utils::fr_to_bytes_le`](crate::utils::fr_to_bytes_le)) + /// + /// Example + /// ``` + /// use rln::utils::*; + /// + /// let mut buffer = Cursor::new(Vec::::new()); + /// let level = 1; + /// let index = 2; + /// rln.get_subtree_root(level, index, &mut buffer).unwrap(); + /// let (subroot, _) = bytes_le_to_fr(&buffer.into_inner()); + /// ``` + pub fn get_subtree_root( + &self, + level: usize, + index: usize, + mut output_data: W, + ) -> Result<()> { + let subroot = self.tree.get_subtree_root(level, index)?; + output_data.write_all(&fr_to_bytes_le(&subroot))?; + + Ok(()) + } + /// Returns the Merkle proof of the leaf at position index /// /// Input values are: diff --git a/rln/tests/poseidon_tree.rs b/rln/tests/poseidon_tree.rs index c6bb5d9..50b7e15 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::hashers::{poseidon_hash, 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,39 @@ mod test { assert_eq!(tree_full_root, tree_opt_root); } + + #[test] + fn test_subtree_root() { + const DEPTH: usize = 3; + const LEAVES_LEN: usize = 6; + + let mut tree = PoseidonTree::default(DEPTH).unwrap(); + let leaves: Vec = (0..LEAVES_LEN).map(|s| Fr::from(s as i32)).collect(); + let _ = tree.set_range(0, leaves); + + for i in 0..LEAVES_LEN { + // check leaves + assert_eq!( + tree.get(i).unwrap(), + tree.get_subtree_root(DEPTH, i).unwrap() + ); + // check root + assert_eq!(tree.root(), tree.get_subtree_root(0, i).unwrap()); + } + + // check intermediate nodes + for n in (1..=DEPTH).rev() { + for i in (0..(1 << n)).step_by(2) { + let idx_l = i * (1 << (DEPTH - n)); + let idx_r = (i + 1) * (1 << (DEPTH - n)); + let idx_sr = idx_l; + + let prev_l = tree.get_subtree_root(n, idx_l).unwrap(); + let prev_r = tree.get_subtree_root(n, idx_r).unwrap(); + let subroot = tree.get_subtree_root(n - 1, idx_sr).unwrap(); + + assert_eq!(poseidon_hash(&[prev_l, prev_r]), subroot); + } + } + } } diff --git a/rln/tests/public.rs b/rln/tests/public.rs index 5b8dde3..8f32307 100644 --- a/rln/tests/public.rs +++ b/rln/tests/public.rs @@ -82,6 +82,29 @@ mod test { assert_eq!(path_elements, expected_path_elements); assert_eq!(identity_path_index, expected_identity_path_index); + // check subtree root computation for leaf 0 for all corresponding node until the root + let l_idx = 0; + for n in (1..=TEST_TREE_HEIGHT).rev() { + let idx_l = l_idx * (1 << (TEST_TREE_HEIGHT - n)); + let idx_r = (l_idx + 1) * (1 << (TEST_TREE_HEIGHT - n)); + let idx_sr = idx_l; + + let mut buffer = Cursor::new(Vec::::new()); + rln.get_subtree_root(n, idx_l, &mut buffer).unwrap(); + let (prev_l, _) = bytes_le_to_fr(&buffer.into_inner()); + + let mut buffer = Cursor::new(Vec::::new()); + rln.get_subtree_root(n, idx_r, &mut buffer).unwrap(); + let (prev_r, _) = bytes_le_to_fr(&buffer.into_inner()); + + let mut buffer = Cursor::new(Vec::::new()); + rln.get_subtree_root(n - 1, idx_sr, &mut buffer).unwrap(); + let (subroot, _) = bytes_le_to_fr(&buffer.into_inner()); + + let res = utils_poseidon_hash(&[prev_l, prev_r]); + assert_eq!(res, subroot); + } + // We double check that the proof computed from public API is correct let root_from_proof = compute_tree_root( &identity_secret_hash, diff --git a/utils/Cargo.toml b/utils/Cargo.toml index 345f981..3cce646 100644 --- a/utils/Cargo.toml +++ b/utils/Cargo.toml @@ -15,10 +15,11 @@ bench = false 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 = { package = "pmtree", version = "=2.0.0", optional = true} +pmtree = { package = "vacp2p_pmtree", version = "=2.0.2", optional = true} sled = "=0.34.7" serde = "=1.0.163" lazy_static = "1.4.0" +hex = "0.4" [dev-dependencies] ark-bn254 = "=0.4.0" diff --git a/utils/benches/merkle_tree_benchmark.rs b/utils/benches/merkle_tree_benchmark.rs index 64d81ef..356c2b8 100644 --- a/utils/benches/merkle_tree_benchmark.rs +++ b/utils/benches/merkle_tree_benchmark.rs @@ -90,6 +90,13 @@ pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) { tree.get(0).unwrap(); }) }); + + // check intermediate node getter which required additional computation of sub root index + c.bench_function("OptimalMerkleTree::get_subtree_root", |b| { + b.iter(|| { + tree.get_subtree_root(1, 0).unwrap(); + }) + }); } pub fn full_merkle_tree_benchmark(c: &mut Criterion) { @@ -125,6 +132,13 @@ pub fn full_merkle_tree_benchmark(c: &mut Criterion) { tree.get(0).unwrap(); }) }); + + // check intermediate node getter which required additional computation of sub root index + c.bench_function("FullMerkleTree::get_subtree_root", |b| { + b.iter(|| { + tree.get_subtree_root(1, 0).unwrap(); + }) + }); } criterion_group!( diff --git a/utils/src/merkle_tree/full_merkle_tree.rs b/utils/src/merkle_tree/full_merkle_tree.rs index 15ef422..9b16a24 100644 --- a/utils/src/merkle_tree/full_merkle_tree.rs +++ b/utils/src/merkle_tree/full_merkle_tree.rs @@ -141,6 +141,33 @@ where Ok(self.nodes[self.capacity() + leaf - 1]) } + fn get_subtree_root(&self, n: usize, index: usize) -> Result { + if n > self.depth() { + return Err(Report::msg("level exceeds depth size")); + } + if index >= self.capacity() { + return Err(Report::msg("index exceeds set size")); + } + if n == 0 { + Ok(self.root()) + } else if n == self.depth { + self.get(index) + } else { + let mut idx = self.capacity() + index - 1; + let mut nd = self.depth; + loop { + let parent = self.parent(idx).unwrap(); + nd -= 1; + if nd == n { + return Ok(self.nodes[parent]); + } else { + idx = parent; + continue; + } + } + } + } + // Sets tree nodes, starting from start index // Function proper of FullMerkleTree implementation fn set_range>>( diff --git a/utils/src/merkle_tree/merkle_tree.rs b/utils/src/merkle_tree/merkle_tree.rs index 17d107b..7b52413 100644 --- a/utils/src/merkle_tree/merkle_tree.rs +++ b/utils/src/merkle_tree/merkle_tree.rs @@ -50,6 +50,7 @@ pub trait ZerokitMerkleTree { fn leaves_set(&mut self) -> usize; fn root(&self) -> FrOf; fn compute_root(&mut self) -> Result>; + fn get_subtree_root(&self, n: usize, index: usize) -> Result>; fn set(&mut self, index: usize, leaf: FrOf) -> Result<()>; fn set_range(&mut self, start: usize, leaves: I) -> Result<()> where diff --git a/utils/src/merkle_tree/optimal_merkle_tree.rs b/utils/src/merkle_tree/optimal_merkle_tree.rs index 48470e7..d21d2a0 100644 --- a/utils/src/merkle_tree/optimal_merkle_tree.rs +++ b/utils/src/merkle_tree/optimal_merkle_tree.rs @@ -108,6 +108,22 @@ where self.get_node(0, 0) } + fn get_subtree_root(&self, n: usize, index: usize) -> Result { + if n > self.depth() { + return Err(Report::msg("level exceeds depth size")); + } + if index >= self.capacity() { + return Err(Report::msg("index exceeds set size")); + } + if n == 0 { + Ok(self.root()) + } else if n == self.depth { + self.get(index) + } else { + Ok(self.get_node(n, index >> (self.depth - n))) + } + } + // Sets a leaf at the specified tree index fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()> { if index >= self.capacity() { diff --git a/utils/tests/merkle_tree.rs b/utils/tests/merkle_tree.rs index 8d49186..35733ad 100644 --- a/utils/tests/merkle_tree.rs +++ b/utils/tests/merkle_tree.rs @@ -35,7 +35,7 @@ pub mod test { impl Display for TestFr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", String::from_utf8_lossy(self.0.as_slice())) + write!(f, "{}", hex::encode(self.0.as_slice())) } } @@ -48,21 +48,33 @@ pub mod test { } lazy_static! { - static ref LEAVES: [TestFr; 4] = [ + 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); + } + const DEPTH_2: usize = 2; + const DEPTH_3: usize = 3; + + fn default_full_merkle_tree(depth: usize) -> FullMerkleTree { + FullMerkleTree::::new(depth, TestFr([0; 32]), FullMerkleConfig::default()) + .unwrap() } - fn default_full_merkle_tree() -> FullMerkleTree { - FullMerkleTree::::new(2, 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() } @@ -80,28 +92,90 @@ pub mod test { ] .map(TestFr); - let mut tree = default_full_merkle_tree(); + let mut tree = default_full_merkle_tree(DEPTH_2); assert_eq!(tree.root(), default_tree_root); - for i in 0..LEAVES.len() { - tree.set(i, LEAVES[i]).unwrap(); + for i in 0..LEAVES_D2.len() { + tree.set(i, LEAVES_D2[i]).unwrap(); assert_eq!(tree.root(), roots[i]); } - let mut tree = default_optimal_merkle_tree(); + let mut tree = default_optimal_merkle_tree(DEPTH_2); assert_eq!(tree.root(), default_tree_root); - for i in 0..LEAVES.len() { - tree.set(i, LEAVES[i]).unwrap(); + for i in 0..LEAVES_D2.len() { + tree.set(i, LEAVES_D2[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()); + + for i in 0..LEAVES_D3.len() { + // check leaves + assert_eq!( + tree_full.get(i).unwrap(), + tree_full.get_subtree_root(DEPTH_3, i).unwrap() + ); + + // check root + assert_eq!(tree_full.root(), tree_full.get_subtree_root(0, i).unwrap()); + } + + // check intermediate nodes + for n in (1..=DEPTH_3).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_sr = idx_l; + + let prev_l = tree_full.get_subtree_root(n, idx_l).unwrap(); + let prev_r = tree_full.get_subtree_root(n, idx_r).unwrap(); + let subroot = tree_full.get_subtree_root(n - 1, idx_sr).unwrap(); + + // check intermediate nodes + assert_eq!(Keccak256::hash(&[prev_l, prev_r]), subroot); + } + } + + let mut tree_opt = default_full_merkle_tree(DEPTH_3); + let _ = tree_opt.set_range(0, LEAVES_D3.iter().cloned()); + + for i in 0..LEAVES_D3.len() { + // check leaves + assert_eq!( + tree_opt.get(i).unwrap(), + tree_opt.get_subtree_root(DEPTH_3, 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 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_sr = idx_l; + + let prev_l = tree_opt.get_subtree_root(n, idx_l).unwrap(); + let prev_r = tree_opt.get_subtree_root(n, idx_r).unwrap(); + let subroot = tree_opt.get_subtree_root(n - 1, idx_sr).unwrap(); + + // check intermediate nodes + assert_eq!(Keccak256::hash(&[prev_l, prev_r]), subroot); + } + } + } + #[test] fn test_proof() { // 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(DEPTH_2); + for i in 0..LEAVES_D2.len() { // We set the leaves - tree.set(i, LEAVES[i]).unwrap(); + tree.set(i, LEAVES_D2[i]).unwrap(); // We compute a merkle proof let proof = tree.proof(i).expect("index should be set"); @@ -110,22 +184,22 @@ pub mod test { assert_eq!(proof.leaf_index(), i); // We verify the proof - assert!(tree.verify(&LEAVES[i], &proof).unwrap()); + assert!(tree.verify(&LEAVES_D2[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_D2[i]), tree.root()); // We check that the proof is not valid for another leaf assert!(!tree - .verify(&LEAVES[(i + 1) % LEAVES.len()], &proof) + .verify(&LEAVES_D2[(i + 1) % LEAVES_D2.len()], &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(DEPTH_2); + for i in 0..LEAVES_D2.len() { // We set the leaves - tree.set(i, LEAVES[i]).unwrap(); + tree.set(i, LEAVES_D2[i]).unwrap(); // We compute a merkle proof let proof = tree.proof(i).expect("index should be set"); @@ -134,24 +208,24 @@ pub mod test { assert_eq!(proof.leaf_index(), i); // We verify the proof - assert!(tree.verify(&LEAVES[i], &proof).unwrap()); + assert!(tree.verify(&LEAVES_D2[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_D2[i]), tree.root()); // We check that the proof is not valid for another leaf assert!(!tree - .verify(&LEAVES[(i + 1) % LEAVES.len()], &proof) + .verify(&LEAVES_D2[(i + 1) % LEAVES_D2.len()], &proof) .unwrap()); } } #[test] fn test_override_range() { - let mut tree = default_optimal_merkle_tree(); + let mut tree = default_optimal_merkle_tree(DEPTH_2); // We set the leaves - tree.set_range(0, LEAVES.iter().cloned()).unwrap(); + tree.set_range(0, LEAVES_D2.iter().cloned()).unwrap(); let new_leaves = [ hex!("0000000000000000000000000000000000000000000000000000000000000005"),