Merge pull request #19 from mir-protocol/prover_use_trees

Have the prover use the new MerkleTree API

src/circuit_builder.rs

@@ -13,7 +13,8 @@ use crate::gates::constant::ConstantGate;
 use crate::gates::gate::{GateInstance, GateRef};
 use crate::gates::noop::NoopGate;
 use crate::generator::{CopyGenerator, WitnessGenerator};
-use crate::hash::{hash_n_to_hash, merkle_root_bit_rev_order};
+use crate::hash::hash_n_to_hash;
+use crate::merkle_tree::MerkleTree;
 use crate::polynomial::polynomial::PolynomialValues;
 use crate::target::Target;
 use crate::util::{log2_strict, transpose, transpose_poly_values};
@@ -236,20 +237,26 @@ impl<F: Field> CircuitBuilder<F> {
         let constant_vecs = self.constant_polys();
         let constant_ldes = PolynomialValues::lde_multiple(constant_vecs, self.config.rate_bits);
         let constant_ldes_t = transpose_poly_values(constant_ldes);
-        let constants_root = merkle_root_bit_rev_order(constant_ldes_t.clone());
+        let constants_tree = MerkleTree::new(constant_ldes_t, true);
 
         let sigma_vecs = self.sigma_vecs();
         let sigma_ldes = PolynomialValues::lde_multiple(sigma_vecs, self.config.rate_bits);
         let sigma_ldes_t = transpose_poly_values(sigma_ldes);
-        let sigmas_root = merkle_root_bit_rev_order(sigma_ldes_t.clone());
+        let sigmas_tree = MerkleTree::new(sigma_ldes_t, true);
+
+        let constants_root = constants_tree.root;
+        let sigmas_root = sigmas_tree.root;
+        let verifier_only = VerifierOnlyCircuitData {
+            constants_root,
+            sigmas_root,
+        };
 
         let generators = self.generators;
         let prover_only = ProverOnlyCircuitData {
             generators,
-            constant_ldes_t,
-            sigma_ldes_t,
+            constants_tree,
+            sigmas_tree,
         };
-        let verifier_only = VerifierOnlyCircuitData {};
 
         // The HashSet of gates will have a non-deterministic order. When converting to a Vec, we
         // sort by ID to make the ordering deterministic.
@@ -274,8 +281,6 @@ impl<F: Field> CircuitBuilder<F> {
             degree_bits,
             gates,
             num_gate_constraints,
-            constants_root,
-            sigmas_root,
             k_is,
             circuit_digest,
         };

src/circuit_data.rs

@@ -1,6 +1,7 @@
 use crate::field::field::Field;
 use crate::gates::gate::GateRef;
 use crate::generator::WitnessGenerator;
+use crate::merkle_tree::MerkleTree;
 use crate::proof::{Hash, HashTarget, Proof};
 use crate::prover::prove;
 use crate::verifier::verify;
@@ -37,7 +38,7 @@ impl CircuitConfig {
 /// Circuit data required by the prover or the verifier.
 pub struct CircuitData<F: Field> {
     pub(crate) prover_only: ProverOnlyCircuitData<F>,
-    pub(crate) verifier_only: VerifierOnlyCircuitData,
+    pub(crate) verifier_only: VerifierOnlyCircuitData<F>,
     pub(crate) common: CommonCircuitData<F>,
 }
 
@@ -71,7 +72,7 @@ impl<F: Field> ProverCircuitData<F> {
 
 /// Circuit data required by the prover.
 pub struct VerifierCircuitData<F: Field> {
-    pub(crate) verifier_only: VerifierOnlyCircuitData,
+    pub(crate) verifier_only: VerifierOnlyCircuitData<F>,
     pub(crate) common: CommonCircuitData<F>,
 }
 
@@ -84,13 +85,20 @@ impl<F: Field> VerifierCircuitData<F> {
 /// Circuit data required by the prover, but not the verifier.
 pub(crate) struct ProverOnlyCircuitData<F: Field> {
     pub generators: Vec<Box<dyn WitnessGenerator<F>>>,
-    pub constant_ldes_t: Vec<Vec<F>>,
-    /// Transpose of LDEs of sigma polynomials (in the context of Plonk's permutation argument).
-    pub sigma_ldes_t: Vec<Vec<F>>,
+    /// Merkle tree containing LDEs of each constant polynomial.
+    pub constants_tree: MerkleTree<F>,
+    /// Merkle tree containing LDEs of each sigma polynomial.
+    pub sigmas_tree: MerkleTree<F>,
 }
 
 /// Circuit data required by the verifier, but not the prover.
-pub(crate) struct VerifierOnlyCircuitData {}
+pub(crate) struct VerifierOnlyCircuitData<F: Field> {
+    /// A commitment to each constant polynomial.
+    pub(crate) constants_root: Hash<F>,
+
+    /// A commitment to each permutation polynomial.
+    pub(crate) sigmas_root: Hash<F>,
+}
 
 /// Circuit data required by both the prover and the verifier.
 pub(crate) struct CommonCircuitData<F: Field> {
@@ -104,12 +112,6 @@ pub(crate) struct CommonCircuitData<F: Field> {
     /// The largest number of constraints imposed by any gate.
     pub(crate) num_gate_constraints: usize,
 
-    /// A commitment to each constant polynomial.
-    pub(crate) constants_root: Hash<F>,
-
-    /// A commitment to each permutation polynomial.
-    pub(crate) sigmas_root: Hash<F>,
-
     /// The `{k_i}` valued used in `S_ID_i` in Plonk's permutation argument.
     pub(crate) k_is: Vec<F>,
 

src/fri.rs

@@ -8,7 +8,6 @@ use crate::polynomial::polynomial::{PolynomialCoeffs, PolynomialValues};
 use crate::proof::{FriEvaluations, FriMerkleProofs, FriProof, FriQueryRound, Hash};
 use crate::util::log2_strict;
 use anyhow::{ensure, Result};
-use std::iter::FromIterator;
 
 /// Somewhat arbitrary. Smaller values will increase delta, but with diminishing returns,
 /// while increasing L, potentially requiring more challenge points.

src/hash.rs

@@ -1,13 +1,10 @@
 //! Concrete instantiation of a hash function.
 
-use rayon::prelude::*;
-
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::field::Field;
 use crate::gmimc::gmimc_permute_array;
 use crate::proof::{Hash, HashTarget};
 use crate::target::Target;
-use crate::util::reverse_index_bits_in_place;
 
 pub(crate) const SPONGE_RATE: usize = 8;
 pub(crate) const SPONGE_CAPACITY: usize = 4;
@@ -245,33 +242,3 @@ pub fn hash_n_to_hash<F: Field>(inputs: Vec<F>, pad: bool) -> Hash<F> {
 pub fn hash_n_to_1<F: Field>(inputs: Vec<F>, pad: bool) -> F {
     hash_n_to_m(inputs, 1, pad)[0]
 }
-
-/// Like `merkle_root`, but first reorders each vector so that `new[i] = old[i.reverse_bits()]`.
-pub(crate) fn merkle_root_bit_rev_order<F: Field>(mut vecs: Vec<Vec<F>>) -> Hash<F> {
-    reverse_index_bits_in_place(&mut vecs);
-    merkle_root(vecs)
-}
-
-/// Given `n` vectors, each of length `l`, constructs a Merkle tree with `l` leaves, where each leaf
-/// is a hash obtained by hashing a "leaf set" consisting of `n` elements. If `n <= 4`, this hashing
-/// is skipped, as there is no need to compress leaf data.
-pub(crate) fn merkle_root<F: Field>(vecs: Vec<Vec<F>>) -> Hash<F> {
-    let elems_per_leaf = vecs[0].len();
-    let leaves_per_chunk = (ELEMS_PER_CHUNK / elems_per_leaf).next_power_of_two();
-    let subtree_roots: Vec<Vec<F>> = vecs
-        .par_chunks(leaves_per_chunk)
-        .map(|chunk| merkle_root_inner(chunk.to_vec()).elements.to_vec())
-        .collect();
-    merkle_root_inner(subtree_roots)
-}
-
-pub(crate) fn merkle_root_inner<F: Field>(vecs: Vec<Vec<F>>) -> Hash<F> {
-    let mut hashes = vecs.into_iter().map(hash_or_noop).collect::<Vec<_>>();
-    while hashes.len() > 1 {
-        hashes = hashes
-            .chunks(2)
-            .map(|pair| compress(pair[0], pair[1]))
-            .collect();
-    }
-    hashes[0]
-}

src/merkle_tree.rs

@@ -1,3 +1,5 @@
+use rayon::prelude::*;
+
 use crate::field::field::Field;
 use crate::hash::{compress, hash_or_noop};
 use crate::merkle_proofs::MerkleProof;
@@ -26,7 +28,7 @@ impl<F: Field> MerkleTree<F> {
             reverse_index_bits_in_place(&mut leaves);
         }
         let mut layers = vec![leaves
-            .iter()
+            .par_iter()
             .map(|l| hash_or_noop(l.clone()))
             .collect::<Vec<_>>()];
         while let Some(l) = layers.last() {
@@ -34,7 +36,7 @@ impl<F: Field> MerkleTree<F> {
                 break;
             }
             let next_layer = l
-                .chunks(2)
+                .par_chunks(2)
                 .map(|chunk| compress(chunk[0], chunk[1]))
                 .collect::<Vec<_>>();
             layers.push(next_layer);
@@ -80,11 +82,13 @@ impl<F: Field> MerkleTree<F> {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use anyhow::Result;
+
     use crate::field::crandall_field::CrandallField;
     use crate::merkle_proofs::verify_merkle_proof;
     use crate::polynomial::division::divide_by_z_h;
-    use anyhow::Result;
+
+    use super::*;
 
     #[test]
     fn test_merkle_trees() -> Result<()> {

src/prover.rs

@@ -7,7 +7,7 @@ use crate::circuit_data::{CommonCircuitData, ProverOnlyCircuitData};
 use crate::field::fft::{fft, ifft};
 use crate::field::field::Field;
 use crate::generator::generate_partial_witness;
-use crate::hash::merkle_root_bit_rev_order;
+use crate::merkle_tree::MerkleTree;
 use crate::plonk_challenger::Challenger;
 use crate::plonk_common::{eval_l_1, evaluate_gate_constraints, reduce_with_powers_multi};
 use crate::polynomial::division::divide_by_z_h;
@@ -61,7 +61,7 @@ pub(crate) fn prove<F: Field>(
 
     // TODO: Could avoid cloning if it's significant?
     let start_wires_root = Instant::now();
-    let wires_root = merkle_root_bit_rev_order(wire_ldes_t.clone());
+    let wires_tree = MerkleTree::new(wire_ldes_t, true);
     info!(
         "{:.3}s to Merklize wire LDEs",
         start_wires_root.elapsed().as_secs_f32()
@@ -72,7 +72,7 @@ pub(crate) fn prove<F: Field>(
     // TODO: Need to include public inputs as well.
     challenger.observe_hash(&common_data.circuit_digest);
 
-    challenger.observe_hash(&wires_root);
+    challenger.observe_hash(&wires_tree.root);
     let betas = challenger.get_n_challenges(num_checks);
     let gammas = challenger.get_n_challenges(num_checks);
 
@@ -86,13 +86,13 @@ pub(crate) fn prove<F: Field>(
     );
 
     let start_plonk_z_root = Instant::now();
-    let plonk_zs_root = merkle_root_bit_rev_order(plonk_z_ldes_t.clone());
+    let plonk_zs_tree = MerkleTree::new(plonk_z_ldes_t, true);
     info!(
         "{:.3}s to Merklize Z's",
         start_plonk_z_root.elapsed().as_secs_f32()
     );
 
-    challenger.observe_hash(&plonk_zs_root);
+    challenger.observe_hash(&plonk_zs_tree.root);
 
     let alphas = challenger.get_n_challenges(num_checks);
 
@@ -100,8 +100,8 @@ pub(crate) fn prove<F: Field>(
     let vanishing_polys = compute_vanishing_polys(
         common_data,
         prover_data,
-        wire_ldes_t,
-        plonk_z_ldes_t,
+        &wires_tree,
+        &plonk_zs_tree,
         &betas,
         &gammas,
         &alphas,
@@ -125,8 +125,8 @@ pub(crate) fn prove<F: Field>(
             quotient_poly_coeff_ldes.into_par_iter().map(fft).collect();
         all_quotient_poly_chunk_ldes.extend(quotient_poly_chunk_ldes);
     }
-    let quotient_polys_root =
-        merkle_root_bit_rev_order(transpose_poly_values(all_quotient_poly_chunk_ldes));
+    let quotient_polys_tree =
+        MerkleTree::new(transpose_poly_values(all_quotient_poly_chunk_ldes), true);
     info!(
         "{:.3}s to compute quotient polys and their LDEs",
         quotient_polys_start.elapsed().as_secs_f32()
@@ -142,9 +142,9 @@ pub(crate) fn prove<F: Field>(
     );
 
     Proof {
-        wires_root,
-        plonk_zs_root,
-        quotient_polys_root,
+        wires_root: wires_tree.root,
+        plonk_zs_root: plonk_zs_tree.root,
+        quotient_polys_root: quotient_polys_tree.root,
         openings,
         fri_proofs,
     }
@@ -164,8 +164,8 @@ fn compute_z<F: Field>(common_data: &CommonCircuitData<F>, i: usize) -> Polynomi
 fn compute_vanishing_polys<F: Field>(
     common_data: &CommonCircuitData<F>,
     prover_data: &ProverOnlyCircuitData<F>,
-    wire_ldes_t: Vec<Vec<F>>,
-    plonk_z_lde_t: Vec<Vec<F>>,
+    wires_tree: &MerkleTree<F>,
+    plonk_zs_tree: &MerkleTree<F>,
     betas: &[F],
     gammas: &[F],
     alphas: &[F],
@@ -180,13 +180,11 @@ fn compute_vanishing_polys<F: Field>(
         .enumerate()
         .map(|(i, x)| {
             let i_next = (i + 1) % lde_size;
-            let local_wires = &wire_ldes_t[i];
-            let next_wires = &wire_ldes_t[i_next];
-            let local_constants = &prover_data.constant_ldes_t[i];
-            let next_constants = &prover_data.constant_ldes_t[i_next];
-            let local_plonk_zs = &plonk_z_lde_t[i];
-            let next_plonk_zs = &plonk_z_lde_t[i_next];
-            let s_sigmas = &prover_data.sigma_ldes_t[i];
+            let local_wires = &wires_tree.leaves[i];
+            let local_constants = &prover_data.constants_tree.leaves[i];
+            let local_plonk_zs = &plonk_zs_tree.leaves[i];
+            let next_plonk_zs = &plonk_zs_tree.leaves[i_next];
+            let s_sigmas = &prover_data.sigmas_tree.leaves[i];
 
             debug_assert_eq!(local_wires.len(), common_data.config.num_wires);
             debug_assert_eq!(local_plonk_zs.len(), num_checks);

src/verifier.rs

@@ -2,7 +2,7 @@ use crate::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData};
 use crate::field::field::Field;
 
 pub(crate) fn verify<F: Field>(
-    verifier_data: &VerifierOnlyCircuitData,
+    verifier_data: &VerifierOnlyCircuitData<F>,
     common_data: &CommonCircuitData<F>,
 ) {
     todo!()