Merkle subtree proofs - WIP

src/hash.rs

@@ -273,6 +273,7 @@ pub(crate) fn merkle_root_inner<F: Field>(vecs: Vec<Vec<F>>) -> Hash<F> {
         .map(|leaf_set| hash_or_noop(leaf_set))
         .collect::<Vec<_>>();
     while hashes.len() > 1 {
+        dbg!(&hashes);
         hashes = hashes
             .chunks(2)
             .map(|pair| compress(pair[0], pair[1]))

src/merkle_proofs.rs

@@ -1,10 +1,11 @@
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::field::Field;
 use crate::gates::gmimc::GMiMCGate;
-use crate::hash::GMIMC_ROUNDS;
 use crate::hash::{compress, hash_or_noop};
+use crate::hash::{merkle_root_inner, GMIMC_ROUNDS};
 use crate::proof::{Hash, HashTarget};
 use crate::target::Target;
+use crate::util::reverse_index_bits_in_place;
 use crate::wire::Wire;
 use anyhow::{ensure, Result};
 
@@ -47,6 +48,37 @@ pub(crate) fn verify_merkle_proof<F: Field>(
     Ok(())
 }
 
+/// Verifies that the given subtree is present at the given index in the Merkle tree with the
+/// given root.
+pub(crate) fn verify_merkle_proof_subtree<F: Field>(
+    mut subtree_leaves_data: Vec<Vec<F>>,
+    subtree_index: usize,
+    merkle_root: Hash<F>,
+    proof: &MerkleProof<F>,
+    reverse_bits: bool,
+) -> Result<()> {
+    let index = if reverse_bits {
+        // reverse_index_bits_in_place(&mut subtree_leaves_data);
+        crate::util::reverse_bits(subtree_index, proof.siblings.len())
+    } else {
+        subtree_index
+    };
+    dbg!(&subtree_leaves_data);
+    let mut current_digest = merkle_root_inner(subtree_leaves_data);
+    dbg!(current_digest);
+    for (i, &sibling_digest) in proof.siblings.iter().enumerate() {
+        let bit = (index >> i & 1) == 1;
+        current_digest = if bit {
+            compress(sibling_digest, current_digest)
+        } else {
+            compress(current_digest, sibling_digest)
+        }
+    }
+    ensure!(current_digest == merkle_root, "Invalid Merkle proof.");
+
+    Ok(())
+}
+
 impl<F: Field> CircuitBuilder<F> {
     /// Verifies that the given leaf data is present at the given index in the Merkle tree with the
     /// given root.

src/merkle_tree.rs

@@ -76,24 +76,65 @@ impl<F: Field> MerkleTree<F> {
                 .collect(),
         }
     }
+
+    /// Create a Merkle proof for an entire subtree.
+    /// Example:
+    /// ```
+    ///         G
+    ///        / \
+    ///       /   \
+    ///      /     \
+    ///     E       F
+    ///    / \     / \
+    ///   A   B   C   D
+    /// ```
+    /// `self.prove_subtree(0, 1)` gives a Merkle proof for the subtree E->(A,B), i.e., the
+    /// path (F,).
+    pub fn prove_subtree(&self, subtree_index: usize, subtree_height: usize) -> MerkleProof<F> {
+        let index = if self.reverse_bits {
+            reverse_bits(
+                subtree_index,
+                log2_strict(self.leaves.len()) - subtree_height,
+            )
+        } else {
+            subtree_index
+        };
+        MerkleProof {
+            siblings: self
+                .layers
+                .iter()
+                .skip(subtree_height)
+                .scan(index, |acc, layer| {
+                    let index = *acc ^ 1;
+                    *acc >>= 1;
+                    Some(layer[index])
+                })
+                .collect(),
+        }
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::field::crandall_field::CrandallField;
-    use crate::merkle_proofs::verify_merkle_proof;
+    use crate::merkle_proofs::{verify_merkle_proof, verify_merkle_proof_subtree};
     use crate::polynomial::division::divide_by_z_h;
     use anyhow::Result;
 
+    fn random_data<F: Field>(n: usize, k: usize) -> Vec<Vec<F>> {
+        (0..n)
+            .map(|_| (0..k).map(|_| F::rand()).collect())
+            .collect()
+    }
+
     #[test]
     fn test_merkle_trees() -> Result<()> {
         type F = CrandallField;
 
-        let n = 1 << 10;
-        let leaves: Vec<Vec<F>> = (0..n)
-            .map(|_| (0..10).map(|_| F::rand()).collect())
-            .collect();
+        let log_n = 3;
+        let n = 1 << log_n;
+        let leaves = random_data::<F>(n, 7);
 
         let tree = MerkleTree::new(leaves.clone(), false);
         for i in 0..n {
@@ -101,12 +142,54 @@ mod tests {
             verify_merkle_proof(tree.leaves[i].clone(), i, tree.root, &proof, false)?;
         }
 
+        for height in 0..=log_n {
+            for i in 0..(n >> height) {
+                let subtree_proof = tree.prove_subtree(i, height);
+                verify_merkle_proof_subtree(
+                    tree.leaves[i << height..(i + 1) << height].to_vec(),
+                    i,
+                    tree.root,
+                    &subtree_proof,
+                    false,
+                )?;
+            }
+        }
+
         let tree_reversed_bits = MerkleTree::new(leaves.clone(), true);
         for i in 0..n {
             let proof = tree_reversed_bits.prove(i);
             verify_merkle_proof(leaves[i].clone(), i, tree_reversed_bits.root, &proof, true)?;
         }
 
+        let (height, i) = (1, 0);
+        dbg!(height, i);
+        let subtree_proof = tree_reversed_bits.prove_subtree(i, height);
+        dbg!(&tree_reversed_bits, &subtree_proof);
+        verify_merkle_proof_subtree(
+            (i << height..(i + 1) << height)
+                .map(|j| tree_reversed_bits.leaves[j].to_vec())
+                .collect(),
+            i,
+            tree_reversed_bits.root,
+            &subtree_proof,
+            true,
+        )?;
+        for height in 1..=log_n {
+            for i in 0..(n >> height) {
+                dbg!(height, i);
+                let subtree_proof = tree_reversed_bits.prove_subtree(i, height);
+                verify_merkle_proof_subtree(
+                    (i << height..(i + 1) << height)
+                        .map(|j| tree_reversed_bits.leaves[j].to_vec())
+                        .collect(),
+                    i,
+                    tree_reversed_bits.root,
+                    &subtree_proof,
+                    true,
+                )?;
+            }
+        }
+
         Ok(())
     }
 }