Expose a public function to fetch the root of a subtree at level n (#247

)

* add get_subroot function

* update test

* update pmtree dependecy

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() {