Merge branch 'main' into digest-generic-mt

CHANGELOG.md

@@ -69,6 +69,7 @@ and follow [semantic versioning](https://semver.org/) for our releases.
 - [#85](https://github.com/EspressoSystems/jellyfish/pull/85), [#87](https://github.com/EspressoSystems/jellyfish/pull/87) (all) Added `no_std` compliance
 - [#116](https://github.com/EspressoSystems/jellyfish/pull/116) (`jf-primitives`) Introduced new `PolynomialCommitmentScheme` trait
 - [#117](https://github.com/EspressoSystems/jellyfish/pull/117) (`jf-relation`) Added gadgets for comparison with constant values
+- [#176](https://github.com/EspressoSystems/jellyfish/pull/176) (`jf-primitives`) Added implementation for light weight merkle tree -- an append-only merkle tree who only keeps its frontier.
 
 ### Changed
 

primitives/src/merkle_tree/append_only.rs

@@ -22,7 +22,7 @@ use num_traits::pow::pow;
 use serde::{Deserialize, Serialize};
 use typenum::Unsigned;
 
-impl_merkle_tree_scheme!(MerkleTree);
+impl_merkle_tree_scheme!(MerkleTree, build_tree_internal);
 impl_forgetable_merkle_tree_scheme!(MerkleTree);
 
 impl<E, H, I, Arity, T> AppendableMerkleTreeScheme for MerkleTree<E, H, I, Arity, T>

primitives/src/merkle_tree/internal.rs

@@ -66,6 +66,11 @@ where
             Self::ForgettenSubtree { value } => *value,
         }
     }
+
+    #[inline]
+    pub(crate) fn is_forgotten(&self) -> bool {
+        matches!(self, Self::ForgettenSubtree { .. })
+    }
 }
 
 /// A merkle commitment consists a root hash value, a tree height and number of
@@ -226,6 +231,92 @@ where
     }
 }
 
+#[allow(clippy::type_complexity)]
+pub(crate) fn build_light_weight_tree_internal<E, H, I, Arity, T>(
+    height: usize,
+    elems: impl IntoIterator<Item = impl Borrow<E>>,
+) -> Result<(Box<MerkleNode<E, I, T>>, u64), PrimitivesError>
+where
+    E: Element,
+    H: DigestAlgorithm<E, I, T>,
+    I: Index + From<u64>,
+    Arity: Unsigned,
+    T: NodeValue,
+{
+    let leaves: Vec<_> = elems.into_iter().collect();
+    let num_leaves = leaves.len() as u64;
+    let capacity = num_traits::checked_pow(Arity::to_u64(), height).ok_or(
+        PrimitivesError::ParameterError("Merkle tree size too large.".to_string()),
+    )?;
+
+    if num_leaves > capacity {
+        Err(PrimitivesError::ParameterError(
+            "Too many data for merkle tree".to_string(),
+        ))
+    } else if num_leaves > 0 {
+        let mut cur_nodes = leaves
+            .into_iter()
+            .enumerate()
+            .chunks(Arity::to_usize())
+            .into_iter()
+            .map(|chunk| {
+                let children = chunk
+                    .map(|(pos, elem)| {
+                        if (pos as u64) < num_leaves - 1 {
+                            Box::new(MerkleNode::ForgettenSubtree {
+                                value: H::digest_leaf(&I::from(pos as u64), elem.borrow()),
+                            })
+                        } else {
+                            let pos = I::from(pos as u64);
+                            Box::new(MerkleNode::Leaf {
+                                value: H::digest_leaf(&pos, elem.borrow()),
+                                pos,
+                                elem: elem.borrow().clone(),
+                            })
+                        }
+                    })
+                    .pad_using(Arity::to_usize(), |_| Box::new(MerkleNode::Empty))
+                    .collect_vec();
+                Box::new(MerkleNode::<E, I, T>::Branch {
+                    value: digest_branch::<E, H, I, T>(&children),
+                    children,
+                })
+            })
+            .collect_vec();
+        for i in 1..cur_nodes.len() - 1 {
+            cur_nodes[i] = Box::new(MerkleNode::ForgettenSubtree {
+                value: cur_nodes[i].value(),
+            })
+        }
+        for _ in 1..height {
+            cur_nodes = cur_nodes
+                .into_iter()
+                .chunks(Arity::to_usize())
+                .into_iter()
+                .map(|chunk| {
+                    let children = chunk
+                        .pad_using(Arity::to_usize(), |_| {
+                            Box::new(MerkleNode::<E, I, T>::Empty)
+                        })
+                        .collect_vec();
+                    Box::new(MerkleNode::<E, I, T>::Branch {
+                        value: digest_branch::<E, H, I, T>(&children),
+                        children,
+                    })
+                })
+                .collect_vec();
+            for i in 1..cur_nodes.len() - 1 {
+                cur_nodes[i] = Box::new(MerkleNode::ForgettenSubtree {
+                    value: cur_nodes[i].value(),
+                })
+            }
+        }
+        Ok((cur_nodes[0].clone(), num_leaves))
+    } else {
+        Ok((Box::new(MerkleNode::<E, I, T>::Empty), 0))
+    }
+}
+
 pub(crate) fn digest_branch<E, H, I, T>(data: &[Box<MerkleNode<E, I, T>>]) -> T
 where
     E: Element,
@@ -299,6 +390,7 @@ where
         }
     }
 
+    /// Re-insert a forgotten leaf to the Merkle tree if the proof is valid.
     pub(crate) fn remember_internal<H, Arity>(
         &mut self,
         height: usize,
@@ -382,6 +474,9 @@ where
         }
     }
 
+    /// Query the given index at the current Merkle node. Return the element
+    /// with a membership proof if presence, otherwise return a non-membership
+    /// proof.
     #[allow(clippy::type_complexity)]
     pub(crate) fn lookup_internal(
         &self,
@@ -444,6 +539,7 @@ where
         }
     }
 
+    /// Update the element at the given index.
     pub(crate) fn update_internal<H, Arity>(
         &mut self,
         height: usize,
@@ -511,12 +607,13 @@ where
         }
     }
 
+    /// Batch insertion for the given Merkle node.
     pub(crate) fn extend_internal<H, Arity>(
         &mut self,
         height: usize,
         pos: &I,
         traversal_path: &[usize],
-        tight_frontier: bool,
+        at_frontier: bool,
         data: &mut Peekable<impl Iterator<Item = impl Borrow<E>>>,
     ) -> Result<u64, PrimitivesError>
     where
@@ -531,7 +628,7 @@ where
         match self {
             MerkleNode::Branch { value, children } => {
                 let mut cnt = 0u64;
-                let mut frontier = if tight_frontier {
+                let mut frontier = if at_frontier {
                     traversal_path[height - 1]
                 } else {
                     0
@@ -542,7 +639,7 @@ where
                         height - 1,
                         &cur_pos,
                         traversal_path,
-                        tight_frontier && frontier == traversal_path[height - 1],
+                        at_frontier && frontier == traversal_path[height - 1],
                         data,
                     )?;
                     cnt += increment;
@@ -564,7 +661,7 @@ where
                     Ok(1)
                 } else {
                     let mut cnt = 0u64;
-                    let mut frontier = if tight_frontier {
+                    let mut frontier = if at_frontier {
                         traversal_path[height - 1]
                     } else {
                         0
@@ -576,7 +673,109 @@ where
                             height - 1,
                             &cur_pos,
                             traversal_path,
-                            tight_frontier && frontier == traversal_path[height - 1],
+                            at_frontier && frontier == traversal_path[height - 1],
+                            data,
+                        )?;
+                        cnt += increment;
+                        cur_pos += I::from(increment);
+                        frontier += 1;
+                    }
+                    *self = MerkleNode::Branch {
+                        value: digest_branch::<E, H, I, T>(&children),
+                        children,
+                    };
+                    Ok(cnt)
+                }
+            },
+            MerkleNode::Leaf { .. } => Err(PrimitivesError::ParameterError(
+                "Incompatible merkle tree: index already occupied".to_string(),
+            )),
+            MerkleNode::ForgettenSubtree { .. } => Err(PrimitivesError::ParameterError(
+                "Given part of merkle tree is not in memory".to_string(),
+            )),
+        }
+    }
+
+    /// Similar to [`extend_internal`], but this function will automatically
+    /// forget every leaf except for the Merkle tree frontier.
+    pub(crate) fn extend_and_forget_internal<H, Arity>(
+        &mut self,
+        height: usize,
+        pos: &I,
+        traversal_path: &[usize],
+        at_frontier: bool,
+        data: &mut Peekable<impl Iterator<Item = impl Borrow<E>>>,
+    ) -> Result<u64, PrimitivesError>
+    where
+        H: DigestAlgorithm<E, I, T>,
+        Arity: Unsigned,
+        I: AddAssign,
+    {
+        if data.peek().is_none() {
+            return Ok(0);
+        }
+        let mut cur_pos = pos.clone();
+        match self {
+            MerkleNode::Branch { value, children } => {
+                let mut cnt = 0u64;
+                let mut frontier = if at_frontier {
+                    traversal_path[height - 1]
+                } else {
+                    0
+                };
+                let cap = Arity::to_usize();
+                while data.peek().is_some() && frontier < cap {
+                    if frontier > 0 && !children[frontier - 1].is_forgotten() {
+                        children[frontier - 1] =
+                            Box::new(MerkleNode::<E, I, T>::ForgettenSubtree {
+                                value: children[frontier - 1].value(),
+                            });
+                    }
+                    let increment = children[frontier].extend_internal::<H, Arity>(
+                        height - 1,
+                        &cur_pos,
+                        traversal_path,
+                        at_frontier && frontier == traversal_path[height - 1],
+                        data,
+                    )?;
+                    cnt += increment;
+                    cur_pos += I::from(increment);
+                    frontier += 1;
+                }
+                *value = digest_branch::<E, H, I, T>(children);
+                Ok(cnt)
+            },
+            MerkleNode::Empty => {
+                if height == 0 {
+                    let elem = data.next().unwrap();
+                    let elem = elem.borrow();
+                    *self = MerkleNode::Leaf {
+                        value: H::digest_leaf(pos, elem),
+                        pos: pos.clone(),
+                        elem: elem.clone(),
+                    };
+                    Ok(1)
+                } else {
+                    let mut cnt = 0u64;
+                    let mut frontier = if at_frontier {
+                        traversal_path[height - 1]
+                    } else {
+                        0
+                    };
+                    let cap = Arity::to_usize();
+                    let mut children = vec![Box::new(MerkleNode::Empty); cap];
+                    while data.peek().is_some() && frontier < cap {
+                        if frontier > 0 && !children[frontier - 1].is_forgotten() {
+                            children[frontier - 1] =
+                                Box::new(MerkleNode::<E, I, T>::ForgettenSubtree {
+                                    value: children[frontier - 1].value(),
+                                });
+                        }
+                        let increment = children[frontier].extend_internal::<H, Arity>(
+                            height - 1,
+                            &cur_pos,
+                            traversal_path,
+                            at_frontier && frontier == traversal_path[height - 1],
                             data,
                         )?;
                         cnt += increment;
@@ -607,6 +806,8 @@ where
     T: NodeValue,
     Arity: Unsigned,
 {
+    /// Verify a membership proof by comparing the computed root value to the
+    /// expected one.
     pub(crate) fn verify_membership_proof<H>(
         &self,
         expected_root: &T,
@@ -654,6 +855,8 @@ where
         }
     }
 
+    /// Verify a non membership proof by comparing the computed root value
+    /// to the expected one.
     pub(crate) fn verify_non_membership_proof<H>(
         &self,
         expected_root: &T,