Working path (de)compression

src/hash/merkle_proofs.rs

@@ -13,7 +13,7 @@ use crate::iop::target::{BoolTarget, Target};
 use crate::iop::wire::Wire;
 use crate::plonk::circuit_builder::CircuitBuilder;
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
 #[serde(bound = "")]
 pub struct MerkleProof<F: Field> {
     /// The Merkle digest of each sibling subtree, staying from the bottommost layer.

src/hash/merkle_tree.rs

@@ -13,6 +13,10 @@ use crate::hash::merkle_proofs::MerkleProof;
 pub struct MerkleCap<F: Field>(pub Vec<HashOut<F>>);
 
 impl<F: Field> MerkleCap<F> {
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
     pub fn flatten(&self) -> Vec<F> {
         self.0.iter().flat_map(|h| h.elements).collect()
     }

src/hash/mod.rs

@@ -3,6 +3,7 @@ pub mod hash_types;
 pub mod hashing;
 pub mod merkle_proofs;
 pub mod merkle_tree;
+pub mod path_compression;
 pub mod poseidon;
 pub mod rescue;
 

src/hash/path_compression.rs

@@ -0,0 +1,150 @@
+use std::collections::HashMap;
+
+use anyhow::{ensure, Result};
+use num::Integer;
+use serde::{Deserialize, Serialize};
+
+use crate::field::field_types::{Field, RichField};
+use crate::hash::hash_types::HashOut;
+use crate::hash::hashing::{compress, hash_or_noop};
+use crate::hash::merkle_proofs::MerkleProof;
+use crate::hash::merkle_tree::MerkleCap;
+use crate::util::log2_strict;
+
+/// Compress multiple Merkle proofs on the same tree by removing redundancy in the Merkle paths.
+pub(crate) fn compress_merkle_proofs<F: Field>(
+    cap_height: usize,
+    proofs: Vec<(usize, MerkleProof<F>)>,
+) -> Vec<MerkleProof<F>> {
+    assert!(!proofs.is_empty());
+    let height = cap_height + proofs[0].1.siblings.len();
+    let num_leaves = 1 << height;
+    let mut compressed_proofs = Vec::with_capacity(proofs.len());
+    // Holds the known nodes in the tree at a given time. The root is at index 1.
+    let mut known = vec![false; 2 * num_leaves];
+    for (i, _) in &proofs {
+        // The leaves are known.
+        known[*i + num_leaves] = true;
+    }
+    // For each proof collect all the unknown proof elements.
+    for (i, p) in proofs {
+        let mut compressed_proof = MerkleProof {
+            siblings: Vec::new(),
+        };
+        let mut index = i + num_leaves;
+        for sibling in p.siblings {
+            let sibling_index = index ^ 1;
+            if !known[sibling_index] {
+                // If the sibling is not yet known, add it to the proof and set it to known.
+                compressed_proof.siblings.push(sibling);
+                known[sibling_index] = true;
+            }
+            // Go up the tree and set the parent to known.
+            index >>= 1;
+            known[index] = true;
+        }
+        compressed_proofs.push(compressed_proof);
+    }
+
+    compressed_proofs
+}
+
+/// Verify a compressed Merkle proof.
+/// Note: The data and indices must be in the same order as in `compress_merkle_proofs`.
+pub(crate) fn decompress_merkle_proof<F: RichField>(
+    leaves_data: &[Vec<F>],
+    leaves_indices: &[usize],
+    compressed_proofs: &[MerkleProof<F>],
+    height: usize,
+    cap: &MerkleCap<F>,
+) -> Vec<MerkleProof<F>> {
+    let num_leaves = 1 << height;
+    let cap_height = log2_strict(cap.len());
+    let mut compressed_proofs = compressed_proofs.to_vec();
+    let mut decompressed_proofs = Vec::with_capacity(compressed_proofs.len());
+    // Holds the already seen nodes in the tree along with their value.
+    let mut seen = HashMap::new();
+
+    for (&i, v) in leaves_indices.iter().zip(leaves_data) {
+        // Observe the leaves.
+        seen.insert(i + num_leaves, hash_or_noop(v.to_vec()));
+    }
+    // For every index, go up the tree by querying `seen` to get node values, or if they are unknown
+    // pop them from the compressed proof.
+    for (&i, p) in leaves_indices.iter().zip(compressed_proofs.iter_mut()) {
+        // We'll use `pop` to run through the Merkle paths, so we reverse it first.
+        p.siblings.reverse();
+        let mut decompressed_proof = MerkleProof {
+            siblings: Vec::new(),
+        };
+        let mut index = i + num_leaves;
+        let mut current_digest = seen[&index];
+        for _ in 0..height - cap_height {
+            let sibling_index = index ^ 1;
+            // Get the value of the sibling node by querying it or popping it from the proof.
+            let h = *seen
+                .entry(sibling_index)
+                .or_insert_with(|| p.siblings.pop().unwrap());
+            decompressed_proof.siblings.push(h);
+            // Update the current digest to the value of the parent.
+            current_digest = if index.is_even() {
+                compress(current_digest, h)
+            } else {
+                compress(h, current_digest)
+            };
+            // Observe the parent.
+            index >>= 1;
+            seen.insert(index, current_digest);
+        }
+
+        decompressed_proofs.push(decompressed_proof);
+    }
+
+    decompressed_proofs
+}
+
+#[cfg(test)]
+mod tests {
+    use rand::{thread_rng, Rng};
+
+    use super::*;
+    use crate::field::crandall_field::CrandallField;
+    use crate::field::field_types::Field;
+    use crate::hash::merkle_proofs::MerkleProof;
+    use crate::hash::merkle_tree::MerkleTree;
+
+    #[test]
+    fn test_path_compression() {
+        type F = CrandallField;
+        let h = 10;
+        let cap_height = 3;
+        let vs = (0..1 << h).map(|_| vec![F::rand()]).collect::<Vec<_>>();
+        let mt = MerkleTree::new(vs.clone(), cap_height);
+
+        let mut rng = thread_rng();
+        let k = rng.gen_range(1..=1 << h);
+        let indices = (0..k).map(|_| rng.gen_range(0..1 << h)).collect::<Vec<_>>();
+        let proofs: Vec<(usize, MerkleProof<_>)> =
+            indices.iter().map(|&i| (i, mt.prove(i))).collect();
+
+        let compressed_proofs = compress_merkle_proofs(cap_height, proofs.clone());
+        let decompressed_proofs = decompress_merkle_proof(
+            &indices.iter().map(|&i| vs[i].clone()).collect::<Vec<_>>(),
+            &indices,
+            &compressed_proofs,
+            h,
+            &mt.cap,
+        );
+
+        let proofs = proofs.into_iter().map(|(_, p)| p).collect::<Vec<_>>();
+        assert_eq!(proofs, decompressed_proofs);
+
+        let compressed_proof_bytes = serde_cbor::to_vec(&compressed_proofs).unwrap();
+        println!(
+            "Compressed proof length: {} bytes",
+            compressed_proof_bytes.len()
+        );
+        let proof_bytes = serde_cbor::to_vec(&proofs).unwrap();
+        println!("Proof length: {} bytes", proof_bytes.len());
+    }
+}