From 0df1545f0c6474d7cbc05b2f548c87ab354080b0 Mon Sep 17 00:00:00 2001
From: Jakub Nabaglo <jakub@mirprotocol.org>
Date: Fri, 4 Feb 2022 07:59:05 -0800
Subject: [PATCH] Merkle tree bugfixes + tests (#467)

* Merkle tree bugfixes + tests

* Minor: Clippy + lints
---
 plonky2/src/hash/merkle_tree.rs | 59 ++++++++++++++++++++++++++++++---
 1 file changed, 54 insertions(+), 5 deletions(-)

diff --git a/plonky2/src/hash/merkle_tree.rs b/plonky2/src/hash/merkle_tree.rs
index e10b5019..54c62eeb 100644
--- a/plonky2/src/hash/merkle_tree.rs
+++ b/plonky2/src/hash/merkle_tree.rs
@@ -90,6 +90,20 @@ fn fill_digests_buf<F: RichField, H: Hasher<F>>(
     leaves: &[Vec<F>],
     cap_height: usize,
 ) {
+    // Special case of a tree that's all cap. The usual case will panic because we'll try to split
+    // an empty slice into chunks of `0`. (We would not need this if there was a way to split into
+    // `blah` chunks as opposed to chunks _of_ `blah`.)
+    if digests_buf.is_empty() {
+        debug_assert_eq!(cap_buf.len(), leaves.len());
+        cap_buf
+            .par_iter_mut()
+            .zip(leaves)
+            .for_each(|(cap_buf, leaf)| {
+                cap_buf.write(H::hash(leaf, false));
+            });
+        return;
+    }
+
     let subtree_digests_len = digests_buf.len() >> cap_height;
     let subtree_leaves_len = leaves.len() >> cap_height;
     let digests_chunks = digests_buf.par_chunks_exact_mut(subtree_digests_len);
@@ -108,6 +122,12 @@ fn fill_digests_buf<F: RichField, H: Hasher<F>>(
 
 impl<F: RichField, H: Hasher<F>> MerkleTree<F, H> {
     pub fn new(leaves: Vec<Vec<F>>, cap_height: usize) -> Self {
+        let log2_leaves_len = log2_strict(leaves.len());
+        assert!(
+            cap_height <= log2_leaves_len,
+            "cap height should be at most log2(leaves.len())"
+        );
+
         let num_digests = 2 * (leaves.len() - (1 << cap_height));
         let mut digests = Vec::with_capacity(num_digests);
 
@@ -194,16 +214,45 @@ mod tests {
         const D: usize,
     >(
         leaves: Vec<Vec<F>>,
-        n: usize,
+        cap_height: usize,
     ) -> Result<()> {
-        let tree = MerkleTree::<F, C::Hasher>::new(leaves.clone(), 1);
-        for i in 0..n {
+        let tree = MerkleTree::<F, C::Hasher>::new(leaves.clone(), cap_height);
+        for (i, leaf) in leaves.into_iter().enumerate() {
             let proof = tree.prove(i);
-            verify_merkle_proof(leaves[i].clone(), i, &tree.cap, &proof)?;
+            verify_merkle_proof(leaf, i, &tree.cap, &proof)?;
         }
         Ok(())
     }
 
+    #[test]
+    #[should_panic]
+    fn test_cap_height_too_big() {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+
+        let log_n = 8;
+        let cap_height = log_n + 1; // Should panic if `cap_height > len_n`.
+
+        let leaves = random_data::<F>(1 << log_n, 7);
+        let _ = MerkleTree::<F, <C as GenericConfig<D>>::Hasher>::new(leaves, cap_height);
+    }
+
+    #[test]
+    fn test_cap_height_eq_log2_len() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+
+        let log_n = 8;
+        let n = 1 << log_n;
+        let leaves = random_data::<F>(n, 7);
+
+        verify_all_leaves::<F, C, D>(leaves, log_n)?;
+
+        Ok(())
+    }
+
     #[test]
     fn test_merkle_trees() -> Result<()> {
         const D: usize = 2;
@@ -214,7 +263,7 @@ mod tests {
         let n = 1 << log_n;
         let leaves = random_data::<F>(n, 7);
 
-        verify_all_leaves::<F, C, D>(leaves, n)?;
+        verify_all_leaves::<F, C, D>(leaves, 1)?;
 
         Ok(())
     }