From 422e72954c13534769fec72702be5b402da7118d Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Mon, 20 Sep 2021 14:37:28 +0200 Subject: [PATCH] Working path (de)compression --- src/hash/merkle_proofs.rs | 2 +- src/hash/merkle_tree.rs | 4 + src/hash/mod.rs | 1 + src/hash/path_compression.rs | 150 +++++++++++++++++++++++++++++++++++ 4 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 src/hash/path_compression.rs diff --git a/src/hash/merkle_proofs.rs b/src/hash/merkle_proofs.rs index 7a176dd9..4b7fdf63 100644 --- a/src/hash/merkle_proofs.rs +++ b/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 { /// The Merkle digest of each sibling subtree, staying from the bottommost layer. diff --git a/src/hash/merkle_tree.rs b/src/hash/merkle_tree.rs index 1bc28d5a..f512a260 100644 --- a/src/hash/merkle_tree.rs +++ b/src/hash/merkle_tree.rs @@ -13,6 +13,10 @@ use crate::hash::merkle_proofs::MerkleProof; pub struct MerkleCap(pub Vec>); impl MerkleCap { + pub fn len(&self) -> usize { + self.0.len() + } + pub fn flatten(&self) -> Vec { self.0.iter().flat_map(|h| h.elements).collect() } diff --git a/src/hash/mod.rs b/src/hash/mod.rs index 7ba7c42c..e33023ec 100644 --- a/src/hash/mod.rs +++ b/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; diff --git a/src/hash/path_compression.rs b/src/hash/path_compression.rs new file mode 100644 index 00000000..cb6cd8e1 --- /dev/null +++ b/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( + cap_height: usize, + proofs: Vec<(usize, MerkleProof)>, +) -> Vec> { + 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( + leaves_data: &[Vec], + leaves_indices: &[usize], + compressed_proofs: &[MerkleProof], + height: usize, + cap: &MerkleCap, +) -> Vec> { + 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::>(); + 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::>(); + 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::>(), + &indices, + &compressed_proofs, + h, + &mt.cap, + ); + + let proofs = proofs.into_iter().map(|(_, p)| p).collect::>(); + 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()); + } +}