Proper ZK treatment in plonky2

plonky2/src/batch_fri/oracle.rs

@@ -1,5 +1,5 @@
 #[cfg(not(feature = "std"))]
-use alloc::{format, vec::Vec};
+use alloc::{format, vec, vec::Vec};
 
 use itertools::Itertools;
 use plonky2_field::extension::Extendable;
@@ -148,12 +148,16 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             // For each batch, we compute the composition polynomial `F_i = sum alpha^j f_ij`,
             // where `alpha` is a random challenge in the extension field.
             // The final polynomial is then computed as `final_poly = sum_i alpha^(k_i) (F_i(X) - F_i(z_i))/(X-z_i)`
+            // (or `final_poly = R(X) + sum_i alpha^(k_i) (F_i(X) - F_i(z_i))/(X-z_i)` in the case of zero-knowledge),
             // where the `k_i`s are chosen such that each power of `alpha` appears only once in the final sum.
             // There are usually two batches for the openings at `zeta` and `g * zeta`.
             // The oracles used in Plonky2 are given in `FRI_ORACLES` in `plonky2/src/plonk/plonk_common.rs`.
-            for FriBatchInfo { point, polynomials } in &instance.batches {
+            for (idx, FriBatchInfo { point, polynomials }) in instance.batches.iter().enumerate() {
+                let is_zk = fri_params.hiding;
+                let nb_r_polys: usize = is_zk as usize;
+                let last_poly = polynomials.len() - nb_r_polys * (idx == 0) as usize;
                 // Collect the coefficients of all the polynomials in `polynomials`.
-                let polys_coeff = polynomials.iter().map(|fri_poly| {
+                let polys_coeff = polynomials[..last_poly].iter().map(|fri_poly| {
                     &oracles[fri_poly.oracle_index].polynomials[fri_poly.polynomial_index]
                 });
                 let composition_poly = timed!(
@@ -165,6 +169,16 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
                 quotient.coeffs.push(F::Extension::ZERO); // pad back to power of two
                 alpha.shift_poly(&mut final_poly);
                 final_poly += quotient;
+
+                // If we are in the zk case, we still have to add `R(X)` to the batch.
+                if is_zk && idx == 0 {
+                    // There should be only one R polynomial left.
+                    polynomials[last_poly..].iter().for_each(|fri_poly| {
+                        final_poly += oracles[fri_poly.oracle_index].polynomials
+                            [fri_poly.polynomial_index]
+                            .to_extension();
+                    });
+                }
             }
 
             assert_eq!(final_poly.len(), 1 << degree_bits[i]);

plonky2/src/batch_fri/prover.rs

@@ -31,7 +31,7 @@ pub fn batch_fri_proof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>,
     timing: &mut TimingTree,
 ) -> FriProof<F, C::Hasher, D> {
     let n = lde_polynomial_coeffs.len();
-    assert_eq!(lde_polynomial_values[0].len(), n);
+    assert_eq!(lde_polynomial_values[0].len(), lde_polynomial_coeffs.len());
     // The polynomial vectors should be sorted by degree, from largest to smallest, with no duplicate degrees.
     assert!(lde_polynomial_values
         .windows(2)

plonky2/src/batch_fri/recursive_verifier.rs

@@ -62,7 +62,8 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
                 PrecomputedReducedOpeningsTarget::from_os_and_alpha(
                     opn,
                     challenges.fri_alpha,
-                    self
+                    self,
+                    params.hiding,
                 )
             );
             precomputed_reduced_evals.push(pre);
@@ -165,13 +166,20 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let mut alpha = ReducingFactorTarget::new(alpha);
         let mut sum = self.zero_extension();
 
-        for (batch, reduced_openings) in instance[index]
+        for (idx, (batch, reduced_openings)) in instance[index]
             .batches
             .iter()
             .zip(&precomputed_reduced_evals.reduced_openings_at_point)
+            .enumerate()
         {
+            // If we are in the zk case, the `R` polynomial (the last polynomials in the first batch) is added to
+            // the batch polynomial independently, without being quotiented. So the final polynomial becomes:
+            // `final_poly = R(X) + sum_i alpha^(k_i) (F_i(X) - F_i(z_i))/(X-z_i)`.
             let FriBatchInfoTarget { point, polynomials } = batch;
-            let evals = polynomials
+            let is_zk = params.hiding;
+            let nb_r_polys = is_zk as usize;
+            let last_poly = polynomials.len() - nb_r_polys * (idx == 0) as usize;
+            let evals = polynomials[..last_poly]
                 .iter()
                 .map(|p| {
                     let poly_blinding = instance[index].oracles[p.oracle_index].blinding;
@@ -184,6 +192,18 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             let denominator = self.sub_extension(subgroup_x, *point);
             sum = alpha.shift(sum, self);
             sum = self.div_add_extension(numerator, denominator, sum);
+
+            // If we are in the zk case, we still have to add `R(X)` to the batch.
+            if is_zk && idx == 0 {
+                polynomials[last_poly..].iter().for_each(|p| {
+                    let poly_blinding = instance[index].oracles[p.oracle_index].blinding;
+                    let salted = params.hiding && poly_blinding;
+                    let eval = proof.unsalted_eval(p.oracle_index, p.polynomial_index, salted);
+
+                    let eval_extension = eval.to_ext_target(self.zero());
+                    sum = self.add_extension(sum, eval_extension);
+                });
+            }
         }
 
         sum

plonky2/src/batch_fri/verifier.rs

@@ -46,7 +46,8 @@ pub fn verify_batch_fri_proof<
 
     let mut precomputed_reduced_evals = Vec::with_capacity(openings.len());
     for opn in openings {
-        let pre = PrecomputedReducedOpenings::from_os_and_alpha(opn, challenges.fri_alpha);
+        let pre =
+            PrecomputedReducedOpenings::from_os_and_alpha(opn, challenges.fri_alpha, params.hiding);
         precomputed_reduced_evals.push(pre);
     }
     let degree_bits = degree_bits
@@ -123,13 +124,20 @@ fn batch_fri_combine_initial<
     let mut alpha = ReducingFactor::new(alpha);
     let mut sum = F::Extension::ZERO;
 
-    for (batch, reduced_openings) in instances[index]
+    // If we are in the zk case, the `R` polynomial (the last polynomials in the first batch) is added to
+    // the batch polynomial independently, without being quotiented. So the final polynomial becomes:
+    // `final_poly = R(X) + sum_i alpha^(k_i) (F_i(X) - F_i(z_i))/(X-z_i)`.
+    for (idx, (batch, reduced_openings)) in instances[index]
         .batches
         .iter()
         .zip(&precomputed_reduced_evals.reduced_openings_at_point)
+        .enumerate()
     {
         let FriBatchInfo { point, polynomials } = batch;
-        let evals = polynomials
+        let is_zk = params.hiding;
+        let nb_r_polys = is_zk as usize;
+        let last_poly = polynomials.len() - nb_r_polys * (idx == 0) as usize;
+        let evals = polynomials[..last_poly]
             .iter()
             .map(|p| {
                 let poly_blinding = instances[index].oracles[p.oracle_index].blinding;
@@ -142,6 +150,17 @@ fn batch_fri_combine_initial<
         let denominator = subgroup_x - *point;
         sum = alpha.shift(sum);
         sum += numerator / denominator;
+
+        // If we are in the zk case, we still have to add `R(X)` to the batch.
+        if is_zk && idx == 0 {
+            polynomials[last_poly..].iter().for_each(|p| {
+                let poly_blinding = instances[index].oracles[p.oracle_index].blinding;
+                let salted = params.hiding && poly_blinding;
+                let eval = proof.unsalted_eval(p.oracle_index, p.polynomial_index, salted);
+
+                sum += F::Extension::from_basefield(eval);
+            });
+        }
     }
 
     sum

plonky2/src/fri/oracle.rs

@@ -1,5 +1,5 @@
 #[cfg(not(feature = "std"))]
-use alloc::{format, vec::Vec};
+use alloc::{format, vec, vec::Vec};
 
 use itertools::Itertools;
 use plonky2_field::types::Field;
@@ -194,9 +194,19 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
         // where the `k_i`s are chosen such that each power of `alpha` appears only once in the final sum.
         // There are usually two batches for the openings at `zeta` and `g * zeta`.
         // The oracles used in Plonky2 are given in `FRI_ORACLES` in `plonky2/src/plonk/plonk_common.rs`.
-        for FriBatchInfo { point, polynomials } in &instance.batches {
-            // Collect the coefficients of all the polynomials in `polynomials`.
-            let polys_coeff = polynomials.iter().map(|fri_poly| {
+        //
+        // If we are in the zk case, the `R` polynomial (the last polynomials in the first batch) is added to
+        // the batch polynomial independently, without being quotiented. So the final polynomial becomes:
+        // `final_poly = R(X) + sum_i alpha^(k_i) (F_i(X) - F_i(z_i))/(X-z_i)`.
+        // Then, since the degree of `R` is double that of the batch polynomial in our cimplementation, we need to
+        // compute one extra step in FRI to reach the correct degree.
+        let is_zk = fri_params.hiding;
+
+        for (idx, FriBatchInfo { point, polynomials }) in instance.batches.iter().enumerate() {
+            let nb_r_polys = is_zk as usize;
+            let last_poly = polynomials.len() - nb_r_polys * (idx == 0) as usize;
+            // Collect the coefficients of all the polynomials in `polynomials` until `last_poly`.
+            let polys_coeff = polynomials[..last_poly].iter().map(|fri_poly| {
                 &oracles[fri_poly.oracle_index].polynomials[fri_poly.polynomial_index]
             });
             let composition_poly = timed!(
@@ -208,6 +218,15 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             quotient.coeffs.push(F::Extension::ZERO); // pad back to power of two
             alpha.shift_poly(&mut final_poly);
             final_poly += quotient;
+
+            // If we are in the zk case, we still have to add `R(X)` to the batch.
+            if is_zk && idx == 0 {
+                polynomials[last_poly..].iter().for_each(|fri_poly| {
+                    final_poly += oracles[fri_poly.oracle_index].polynomials
+                        [fri_poly.polynomial_index]
+                        .to_extension();
+                });
+            }
         }
 
         let lde_final_poly = final_poly.lde(fri_params.config.rate_bits);

plonky2/src/fri/proof.rs

@@ -144,6 +144,7 @@ impl<F: RichField + Extendable<D>, H: Hasher<F>, const D: usize> FriProof<F, H,
             ..
         } = self;
         let cap_height = params.config.cap_height;
+
         let reduction_arity_bits = &params.reduction_arity_bits;
         let num_reductions = reduction_arity_bits.len();
         let num_initial_trees = query_round_proofs[0].initial_trees_proof.evals_proofs.len();
@@ -215,7 +216,7 @@ impl<F: RichField + Extendable<D>, H: Hasher<F>, const D: usize> FriProof<F, H,
                 .entry(index)
                 .or_insert(initial_proof);
             for j in 0..num_reductions {
-                index >>= reduction_arity_bits[j];
+                index >>= params.reduction_arity_bits[j];
                 let query_step = FriQueryStep {
                     evals: steps_evals[j][i].clone(),
                     merkle_proof: steps_proofs[j][i].clone(),
@@ -256,6 +257,7 @@ impl<F: RichField + Extendable<D>, H: Hasher<F>, const D: usize> CompressedFriPr
         } = &challenges.fri_challenges;
         let mut fri_inferred_elements = fri_inferred_elements.0.into_iter();
         let cap_height = params.config.cap_height;
+
         let reduction_arity_bits = &params.reduction_arity_bits;
         let num_reductions = reduction_arity_bits.len();
         let num_initial_trees = query_round_proofs
@@ -274,7 +276,8 @@ impl<F: RichField + Extendable<D>, H: Hasher<F>, const D: usize> CompressedFriPr
         let mut steps_evals = vec![vec![]; num_reductions];
         let mut steps_proofs = vec![vec![]; num_reductions];
         let height = params.degree_bits + params.config.rate_bits;
-        let heights = reduction_arity_bits
+        let heights = params
+            .reduction_arity_bits
             .iter()
             .scan(height, |acc, &bits| {
                 *acc -= bits;
@@ -283,8 +286,7 @@ impl<F: RichField + Extendable<D>, H: Hasher<F>, const D: usize> CompressedFriPr
             .collect::<Vec<_>>();
 
         // Holds the `evals` vectors that have already been reconstructed at each reduction depth.
-        let mut evals_by_depth =
-            vec![HashMap::<usize, Vec<_>>::new(); params.reduction_arity_bits.len()];
+        let mut evals_by_depth = vec![HashMap::<usize, Vec<_>>::new(); reduction_arity_bits.len()];
         for &(mut index) in indices {
             let initial_trees_proof = query_round_proofs.initial_trees_proofs[&index].clone();
             for (i, (leaves_data, proof)) in

plonky2/src/fri/prover.rs

@@ -29,7 +29,7 @@ pub fn fri_proof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const
     timing: &mut TimingTree,
 ) -> FriProof<F, C::Hasher, D> {
     let n = lde_polynomial_values.len();
-    assert_eq!(lde_polynomial_coeffs.len(), n);
+    assert_eq!(lde_polynomial_coeffs.len(), lde_polynomial_values.len());
 
     // Commit phase
     let (trees, final_coeffs) = timed!(
@@ -76,6 +76,7 @@ fn fri_committed_trees<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>,
     let mut trees = Vec::with_capacity(fri_params.reduction_arity_bits.len());
 
     let mut shift = F::MULTIPLICATIVE_GROUP_GENERATOR;
+
     for arity_bits in &fri_params.reduction_arity_bits {
         let arity = 1 << arity_bits;
 
@@ -196,22 +197,33 @@ fn fri_prover_query_round<
     fri_params: &FriParams,
 ) -> FriQueryRound<F, C::Hasher, D> {
     let mut query_steps = Vec::new();
+
     let initial_proof = initial_merkle_trees
         .iter()
-        .map(|t| (t.get(x_index).to_vec(), t.prove(x_index)))
+        .filter_map(|t| {
+            if !t.leaves.is_empty() {
+                Some((t.get(x_index).to_vec(), t.prove(x_index)))
+            } else {
+                None
+            }
+        })
         .collect::<Vec<_>>();
-    for (i, tree) in trees.iter().enumerate() {
-        let arity_bits = fri_params.reduction_arity_bits[i];
-        let evals = unflatten(tree.get(x_index >> arity_bits));
-        let merkle_proof = tree.prove(x_index >> arity_bits);
 
-        query_steps.push(FriQueryStep {
-            evals,
-            merkle_proof,
-        });
-
-        x_index >>= arity_bits;
+    for (i, tree) in trees.iter().enumerate() {
+        if !tree.leaves.is_empty() {
+            let arity_bits = fri_params.reduction_arity_bits[i];
+            let evals = unflatten(tree.get(x_index >> arity_bits));
+            let merkle_proof = tree.prove(x_index >> arity_bits);
+
+            query_steps.push(FriQueryStep {
+                evals,
+                merkle_proof,
+            });
+
+            x_index >>= arity_bits;
+        }
     }
+
     FriQueryRound {
         initial_trees_proof: FriInitialTreeProof {
             evals_proofs: initial_proof,