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

@@ -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",
 ]

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);

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<FrOf<Self::Hasher>> {
+        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<I: IntoIterator<Item = FrOf<Self::Hasher>>, J: IntoIterator<Item = usize>>(
         &mut self,
         start: usize,

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::<u8>::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<W: Write>(
+        &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:

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<Fr> = (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<Fr> = (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);
+            }
+        }
+    }
 }

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::<u8>::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::<u8>::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::<u8>::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,

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"

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!(

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<H::Fr> {
+        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<I: IntoIterator<Item = FrOf<Self::Hasher>>>(

utils/src/merkle_tree/merkle_tree.rs

@@ -50,6 +50,7 @@ pub trait ZerokitMerkleTree {
     fn leaves_set(&mut self) -> usize;
     fn root(&self) -> FrOf<Self::Hasher>;
     fn compute_root(&mut self) -> Result<FrOf<Self::Hasher>>;
+    fn get_subtree_root(&self, n: usize, index: usize) -> Result<FrOf<Self::Hasher>>;
     fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>) -> Result<()>;
     fn set_range<I>(&mut self, start: usize, leaves: I) -> Result<()>
     where

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<H::Fr> {
+        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() {

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<Keccak256> {
+        FullMerkleTree::<Keccak256>::new(depth, TestFr([0; 32]), FullMerkleConfig::default())
+            .unwrap()
     }
 
-    fn default_full_merkle_tree() -> FullMerkleTree<Keccak256> {
-        FullMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), FullMerkleConfig::default()).unwrap()
-    }
-
-    fn default_optimal_merkle_tree() -> OptimalMerkleTree<Keccak256> {
-        OptimalMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), OptimalMerkleConfig::default())
+    fn default_optimal_merkle_tree(depth: usize) -> OptimalMerkleTree<Keccak256> {
+        OptimalMerkleTree::<Keccak256>::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"),