Make PoseidonBranch store the root separately

This also required a slight change in the algorithm computing the
branch, to handle the root not being stored at the last position of the
path.

Co-authored-by: moCello <moana@dusk.network>
Resolves #189

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Change `PoseidonBranch` to have two fields - `root` and `path`. The path is
+  now a fixed length array. [#189]
 - Change `PoseidonTree` to build only with the `alloc` feature [#180]
 - Change `PoseidonTree` to take a generic `Keyed` over the leaf type [#180]
   instead of a `PoseidonAnnotation` [#180]
@@ -326,6 +328,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Poseidon252 Sponge-hash impl with BulletProofs.
 - Variants of sponge for `Scalar` & `Gadget(Variable/LC)`.
 
+[#189]: https://github.com/dusk-network/poseidon252/issues/189
 [#181]: https://github.com/dusk-network/poseidon252/issues/181
 [#180]: https://github.com/dusk-network/poseidon252/issues/180
 [#175]: https://github.com/dusk-network/poseidon252/issues/175

src/tree/branch.rs

@@ -8,10 +8,7 @@ use super::PoseidonLeaf;
 
 use crate::tree::PoseidonAnnotation;
 
-use alloc::vec::Vec;
-
 use core::borrow::Borrow;
-use core::iter;
 use core::ops::Deref;
 
 use dusk_bls12_381::BlsScalar;
@@ -34,6 +31,7 @@ use rkyv::{Archive, Deserialize, Serialize};
 )]
 pub struct PoseidonLevel {
     pub(crate) level: [BlsScalar; dusk_hades::WIDTH],
+    // TODO figure out if we can make this smaller
     index: u64,
 }
 
@@ -79,7 +77,8 @@ impl AsRef<[BlsScalar]> for PoseidonLevel {
     archive_attr(derive(CheckBytes))
 )]
 pub struct PoseidonBranch<const DEPTH: usize> {
-    pub(crate) path: Vec<PoseidonLevel>,
+    path: [PoseidonLevel; DEPTH],
+    root: BlsScalar,
 }
 
 impl<const DEPTH: usize> PoseidonBranch<DEPTH> {
@@ -88,10 +87,16 @@ impl<const DEPTH: usize> PoseidonBranch<DEPTH> {
 
     /// Represents the root for a given path of an opening over a subtree
     pub fn root(&self) -> &BlsScalar {
-        self.path
-            .last()
-            .map(|l| l.deref())
-            .unwrap_or(&Self::NULL_ROOT)
+        &self.root
+    }
+
+    /// Returns a mutable reference to the level with the given `index` in the
+    /// path.
+    ///
+    /// NOTE: This should **never** be exposed to the consumer of the crate.
+    #[allow(dead_code)]
+    pub(crate) fn level_mut(&mut self, index: usize) -> &mut PoseidonLevel {
+        &mut self.path[index]
     }
 }
 
@@ -103,15 +108,16 @@ impl<const DEPTH: usize> Deref for PoseidonBranch<DEPTH> {
     }
 }
 
-impl<const DEPTH: usize> Default for PoseidonBranch<DEPTH> {
-    fn default() -> Self {
-        let path = iter::repeat(PoseidonLevel::default())
-            .take(DEPTH + 1)
-            .collect();
-
-        Self { path }
-    }
-}
+// // Having a 'Default' branch probably makes no sense
+// impl<const DEPTH: usize> Default for PoseidonBranch<DEPTH> {
+//     fn default() -> Self {
+//         let path = iter::repeat(PoseidonLevel::default())
+//             .take(DEPTH + 1)
+//             .collect();
+//
+//         Self { path }
+//     }
+// }
 
 impl<const DEPTH: usize> AsRef<[PoseidonLevel]> for PoseidonBranch<DEPTH> {
     fn as_ref(&self) -> &[PoseidonLevel] {
@@ -129,24 +135,19 @@ where
     fn from(
         b: &Branch<'_, NStack<L, PoseidonAnnotation<K>>, PoseidonAnnotation<K>>,
     ) -> Self {
-        let mut branch = PoseidonBranch::default();
-        let mut depth = 0;
+        let mut path = [PoseidonLevel::default(); DEPTH];
 
-        b.levels()
-            .iter()
-            .rev()
-            .zip(branch.path.iter_mut())
-            .for_each(|(l, b)| {
-                depth += 1;
-                b.index = l.index() as u64 + 1;
+        b.levels().iter().rev().zip(path.iter_mut()).for_each(
+            |(nstack_level, poseidon_level)| {
+                poseidon_level.index = nstack_level.index() as u64 + 1;
 
                 let mut flag = 1;
                 let mut mask = 0;
 
-                match &**l {
+                match &**nstack_level {
                     NStack::Leaf(l) => l
                         .iter()
-                        .zip(b.level.iter_mut().skip(1))
+                        .zip(poseidon_level.level.iter_mut().skip(1))
                         .for_each(|(leaf, l)| {
                             if let Some(leaf) = leaf {
                                 mask |= flag;
@@ -157,7 +158,7 @@ where
                         }),
                     NStack::Node(n) => n
                         .iter()
-                        .zip(b.level.iter_mut().skip(1))
+                        .zip(poseidon_level.level.iter_mut().skip(1))
                         .for_each(|(node, l)| {
                             if let Some(annotated) = node {
                                 let anno = annotated.anno();
@@ -172,29 +173,47 @@ where
                         }),
                 }
 
-                b.level[0] = BlsScalar::from(mask);
-            });
+                poseidon_level.level[0] = BlsScalar::from(mask);
+            },
+        );
+
+        // If the nstack is smaller than the poseidon tree then the we need to
+        // populate the remaining levels of the tree.
+        let nstack_depth = b.levels().len();
+        let flag = BlsScalar::one();
 
-        if depth >= DEPTH {
-            return branch;
+        if nstack_depth < DEPTH {
+            let level = path[nstack_depth - 1].level;
+            let mut perm = [BlsScalar::zero(); dusk_hades::WIDTH];
+
+            let mut h = ScalarStrategy::new();
+
+            path.iter_mut().skip(nstack_depth).fold(level, |l, b| {
+                perm.copy_from_slice(&l);
+                h.perm(&mut perm);
+
+                b.index = 1;
+                b.level[0] = flag;
+                b.level[1] = perm[1];
+
+                b.level
+            });
         }
 
-        let flag = BlsScalar::one();
-        let level = branch.path[depth - 1].level;
-        let mut perm = [BlsScalar::zero(); dusk_hades::WIDTH];
+        // TODO: The amount of repetition here hints at the fact that hashing
+        //  should be more ergonomic.
 
+        // Calculate the root
+        let mut perm = [BlsScalar::zero(); dusk_hades::WIDTH];
         let mut h = ScalarStrategy::new();
-        branch.path.iter_mut().skip(depth).fold(level, |l, b| {
-            perm.copy_from_slice(&l);
-            h.perm(&mut perm);
 
-            b.index = 1;
-            b.level[0] = flag;
-            b.level[1] = perm[1];
+        perm.copy_from_slice(&path[DEPTH - 1].level);
+        perm[0] = flag;
+        h.perm(&mut perm);
 
-            b.level
-        });
-
-        branch
+        PoseidonBranch {
+            path,
+            root: perm[1],
+        }
     }
 }

src/tree/zk.rs

@@ -213,7 +213,7 @@ mod tests {
         let label = b"dusk-network";
         let (pk, _, pp, mut circuit) = init_valid_opening_setup();
 
-        circuit.branch.path[10].level[3] = BlsScalar::random(&mut OsRng);
+        circuit.branch.level_mut(3).level[3] = BlsScalar::random(&mut OsRng);
 
         // With an incorrect path we can not generate a valid proof
         circuit
@@ -226,9 +226,9 @@ mod tests {
         let label = b"dusk-network";
         let (pk, _, pp, mut circuit) = init_valid_opening_setup();
 
-        for depth in 0..DEPTH + 1 {
+        for depth in 0..DEPTH {
             for offset in 1..dusk_hades::WIDTH {
-                circuit.branch.path[depth].level[offset] =
+                circuit.branch.level_mut(depth).level[offset] =
                     BlsScalar::random(&mut OsRng);
             }
         }