Have the prover use the new MerkleTree API

Before it was storing leaf data and Merkle roots, but nothing in between, since it wasn't yet interacting with intermediate layers (but it will once we hook up the FRI code).

src/circuit_builder.rs

@@ -13,11 +13,12 @@ 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::polynomial::polynomial::PolynomialValues;
 use crate::target::Target;
 use crate::util::{log2_strict, transpose, transpose_poly_values};
 use crate::wire::Wire;
+use crate::merkle_tree::MerkleTree;
 
 pub struct CircuitBuilder<F: Field> {
     pub(crate) config: CircuitConfig,
@@ -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

@@ -5,6 +5,7 @@ use crate::proof::{Hash, HashTarget, Proof};
 use crate::prover::prove;
 use crate::verifier::verify;
 use crate::witness::PartialWitness;
+use crate::merkle_tree::MerkleTree;
 
 #[derive(Copy, Clone)]
 pub struct CircuitConfig {
@@ -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,6 @@ 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::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;
@@ -17,6 +16,7 @@ use crate::util::{transpose, transpose_poly_values};
 use crate::vars::EvaluationVars;
 use crate::wire::Wire;
 use crate::witness::PartialWitness;
+use crate::merkle_tree::MerkleTree;
 
 pub(crate) fn prove<F: Field>(
     prover_data: &ProverOnlyCircuitData<F>,
@@ -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.clone(), 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.clone(), 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);
 
@@ -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,
     }
@@ -182,11 +182,11 @@ fn compute_vanishing_polys<F: Field>(
             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_constants = &prover_data.constants_tree.leaves[i];
+            let next_constants = &prover_data.constants_tree.leaves[i_next]; // TODO: "next" is deprecated
             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 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!()