PR feedback

starky/src/permutation.rs

@@ -4,6 +4,7 @@ use itertools::Itertools;
 use plonky2::field::batch_util::batch_multiply_inplace;
 use plonky2::field::extension_field::{Extendable, FieldExtension};
 use plonky2::field::field_types::Field;
+use plonky2::field::packed_field::PackedField;
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::challenger::{Challenger, RecursiveChallenger};
@@ -54,7 +55,6 @@ pub(crate) struct PermutationChallengeSet<T: Copy> {
 pub(crate) fn compute_permutation_z_polys<F, C, S, const D: usize>(
     stark: &S,
     config: &StarkConfig,
-    challenger: &mut Challenger<F, C::Hasher>,
     trace_poly_values: &[PolynomialValues<F>],
     permutation_challenge_sets: &[PermutationChallengeSet<F>],
 ) -> Vec<PolynomialValues<F>>
@@ -239,27 +239,28 @@ pub(crate) fn get_permutation_batches<'a, T: Copy>(
 }
 
 // TODO: Use slices.
-pub struct PermutationCheckData<F: Field, FE: FieldExtension<D2, BaseField = F>, const D2: usize> {
+pub struct PermutationCheckVars<F: Field, FE: FieldExtension<D2, BaseField = F>, const D2: usize> {
     pub(crate) local_zs: Vec<FE>,
     pub(crate) next_zs: Vec<FE>,
     pub(crate) permutation_challenge_sets: Vec<PermutationChallengeSet<F>>,
 }
 
-pub(crate) fn eval_permutation_checks<F, FE, C, S, const D: usize, const D2: usize>(
+pub(crate) fn eval_permutation_checks<F, FE, P, C, S, const D: usize, const D2: usize>(
     stark: &S,
     config: &StarkConfig,
     vars: StarkEvaluationVars<FE, FE, { S::COLUMNS }, { S::PUBLIC_INPUTS }>,
-    permutation_data: PermutationCheckData<F, FE, D2>,
+    permutation_data: PermutationCheckVars<F, FE, D2>,
     consumer: &mut ConstraintConsumer<FE>,
 ) where
     F: RichField + Extendable<D>,
     FE: FieldExtension<D2, BaseField = F>,
+    P: PackedField<Scalar = FE>,
     C: GenericConfig<D, F = F>,
     S: Stark<F, D>,
     [(); S::COLUMNS]:,
     [(); S::PUBLIC_INPUTS]:,
 {
-    let PermutationCheckData {
+    let PermutationCheckVars {
         local_zs,
         next_zs,
         permutation_challenge_sets,
@@ -350,7 +351,6 @@ pub(crate) fn eval_permutation_checks_recursively<F, S, const D: usize>(
 
     // Each zs value corresponds to a permutation batch.
     for (i, instances) in permutation_batches.iter().enumerate() {
-        // Z(gx) * down = Z x  * up
         let (reduced_lhs, reduced_rhs): (Vec<ExtensionTarget<D>>, Vec<ExtensionTarget<D>>) =
             instances
                 .iter()
@@ -359,7 +359,6 @@ pub(crate) fn eval_permutation_checks_recursively<F, S, const D: usize>(
                         pair: PermutationPair { column_pairs },
                         challenge: PermutationChallenge { beta, gamma },
                     } = instance;
-                    let zero = builder.zero_extension();
                     let beta_ext = builder.convert_to_ext(*beta);
                     let gamma_ext = builder.convert_to_ext(*gamma);
                     let mut factor = ReducingFactorTarget::new(beta_ext);

starky/src/proof.rs

@@ -32,6 +32,7 @@ pub struct StarkProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>,
 }
 
 impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> StarkProof<F, C, D> {
+    /// Recover the length of the trace from a STARK proof and a STARK config.
     pub(crate) fn recover_degree_bits(&self, config: &StarkConfig) -> usize {
         let initial_merkle_proof = &self.opening_proof.query_round_proofs[0]
             .initial_trees_proof
@@ -51,6 +52,7 @@ pub struct StarkProofTarget<const D: usize> {
 }
 
 impl<const D: usize> StarkProofTarget<D> {
+    /// Recover the length of the trace from a STARK proof and a STARK config.
     pub(crate) fn recover_degree_bits(&self, config: &StarkConfig) -> usize {
         let initial_merkle_proof = &self.opening_proof.query_round_proofs[0]
             .initial_trees_proof

starky/src/prover.rs

@@ -18,7 +18,7 @@ use rayon::prelude::*;
 
 use crate::config::StarkConfig;
 use crate::constraint_consumer::ConstraintConsumer;
-use crate::permutation::PermutationCheckData;
+use crate::permutation::PermutationCheckVars;
 use crate::permutation::{
     compute_permutation_z_polys, get_n_permutation_challenge_sets, PermutationChallengeSet,
 };
@@ -93,26 +93,23 @@ where
         let permutation_z_polys = compute_permutation_z_polys::<F, C, S, D>(
             &stark,
             config,
-            &mut challenger,
             &trace_poly_values,
             &permutation_challenge_sets,
         );
 
-        timed!(
+        let permutation_zs_commitment = timed!(
             timing,
             "compute permutation Z commitments",
-            (
+            PolynomialBatch::from_values(
-                PolynomialBatch::from_values(
+                permutation_z_polys,
-                    permutation_z_polys,
+                rate_bits,
-                    rate_bits,
+                false,
-                    false,
+                config.fri_config.cap_height,
-                    config.fri_config.cap_height,
+                timing,
-                    timing,
+                None,
-                    None,
-                ),
-                permutation_challenge_sets
             )
-        )
+        );
+        (permutation_zs_commitment, permutation_challenge_sets)
     });
     let permutation_zs_commitment = permutation_zs_commitment_challenges
         .as_ref()
@@ -251,6 +248,8 @@ where
     // Retrieve the LDE values at index `i`.
     let get_at_index =
         |comm: &'a PolynomialBatch<F, C, D>, i: usize| -> &'a [F] { comm.get_lde_values(i * step) };
+    let get_trace_at_index = |i| get_at_index(trace_commitment, i).try_into().unwrap();
+
     // Last element of the subgroup.
     let last = F::primitive_root_of_unity(degree_bits).inverse();
     let size = degree << quotient_degree_bits;
@@ -271,21 +270,20 @@ where
                 lagrange_last.values[i],
             );
             let vars = StarkEvaluationVars::<F, F, { S::COLUMNS }, { S::PUBLIC_INPUTS }> {
-                local_values: &get_at_index(trace_commitment, i).try_into().unwrap(),
+                local_values: &get_trace_at_index(i),
-                next_values: &get_at_index(trace_commitment, (i + next_step) % size)
+                next_values: &get_trace_at_index((i + next_step) % size),
-                    .try_into()
-                    .unwrap(),
                 public_inputs: &public_inputs,
             };
             let permutation_check_data = permutation_zs_commitment_challenges.as_ref().map(
-                |(permutation_zs_commitment, permutation_challenge_sets)| PermutationCheckData {
+                |(permutation_zs_commitment, permutation_challenge_sets)| PermutationCheckVars {
                     local_zs: get_at_index(permutation_zs_commitment, i).to_vec(),
                     next_zs: get_at_index(permutation_zs_commitment, (i + next_step) % size)
                         .to_vec(),
                     permutation_challenge_sets: permutation_challenge_sets.to_vec(),
                 },
             );
-            eval_vanishing_poly::<F, F, C, S, D, 1>(
+            // TODO: Use packed field for F.
+            eval_vanishing_poly::<F, F, F, C, S, D, 1>(
                 stark,
                 config,
                 vars,

starky/src/recursive_verifier.rs

@@ -1,5 +1,6 @@
 use std::iter::once;
 
+use anyhow::{ensure, Result};
 use itertools::Itertools;
 use plonky2::field::extension_field::Extendable;
 use plonky2::field::field_types::Field;
@@ -69,6 +70,7 @@ fn recursively_verify_stark_proof_with_challenges<
     [(); S::COLUMNS]:,
     [(); S::PUBLIC_INPUTS]:,
 {
+    check_permutation_options(&stark, &proof_with_pis, &challenges).unwrap();
     let one = builder.one_extension();
 
     let StarkProofWithPublicInputsTarget {
@@ -202,18 +204,14 @@ pub fn add_virtual_stark_proof<F: RichField + Extendable<D>, S: Stark<F, D>, con
     let fri_params = config.fri_params(degree_bits);
     let cap_height = fri_params.config.cap_height;
 
-    let num_leaves_per_oracle = if stark.uses_permutation_args() {
+    let num_leaves_per_oracle = once(S::COLUMNS)
-        vec![
+        .chain(
-            S::COLUMNS,
+            stark
-            stark.num_permutation_batches(config),
+                .uses_permutation_args()
-            stark.quotient_degree_factor() * config.num_challenges,
+                .then(|| stark.num_permutation_batches(config)),
-        ]
+        )
-    } else {
+        .chain(once(stark.quotient_degree_factor() * config.num_challenges))
-        vec![
+        .collect_vec();
-            S::COLUMNS,
-            stark.quotient_degree_factor() * config.num_challenges,
-        ]
-    };
 
     let permutation_zs_cap = stark
         .uses_permutation_args()
@@ -299,3 +297,25 @@ pub fn set_stark_proof_target<F, C: GenericConfig<D, F = F>, W, const D: usize>(
 
     set_fri_proof_target(witness, &proof_target.opening_proof, &proof.opening_proof);
 }
+
+/// Utility function to check that all permutation data wrapped in `Option`s are `Some` iff
+/// the Stark uses a permutation argument.
+fn check_permutation_options<F: RichField + Extendable<D>, S: Stark<F, D>, const D: usize>(
+    stark: &S,
+    proof_with_pis: &StarkProofWithPublicInputsTarget<D>,
+    challenges: &StarkProofChallengesTarget<D>,
+) -> Result<()> {
+    let options_is_some = [
+        proof_with_pis.proof.permutation_zs_cap.is_some(),
+        proof_with_pis.proof.openings.permutation_zs.is_some(),
+        proof_with_pis.proof.openings.permutation_zs_right.is_some(),
+        challenges.permutation_challenge_sets.is_some(),
+    ];
+    ensure!(
+        options_is_some
+            .into_iter()
+            .all(|b| b == stark.uses_permutation_args()),
+        "Permutation data doesn't match with Stark configuration."
+    );
+    Ok(())
+}

starky/src/vanishing_poly.rs

@@ -1,4 +1,5 @@
 use plonky2::field::extension_field::{Extendable, FieldExtension};
+use plonky2::field::packed_field::PackedField;
 use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::config::GenericConfig;
@@ -6,21 +7,22 @@ use plonky2::plonk::config::GenericConfig;
 use crate::config::StarkConfig;
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::permutation::{
-    eval_permutation_checks, eval_permutation_checks_recursively, PermutationCheckData,
+    eval_permutation_checks, eval_permutation_checks_recursively, PermutationCheckDataTarget,
-    PermutationCheckDataTarget,
+    PermutationCheckVars,
 };
 use crate::stark::Stark;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
-pub(crate) fn eval_vanishing_poly<F, FE, C, S, const D: usize, const D2: usize>(
+pub(crate) fn eval_vanishing_poly<F, FE, P, C, S, const D: usize, const D2: usize>(
     stark: &S,
     config: &StarkConfig,
     vars: StarkEvaluationVars<FE, FE, { S::COLUMNS }, { S::PUBLIC_INPUTS }>,
-    permutation_data: Option<PermutationCheckData<F, FE, D2>>,
+    permutation_data: Option<PermutationCheckVars<F, FE, D2>>,
     consumer: &mut ConstraintConsumer<FE>,
 ) where
     F: RichField + Extendable<D>,
     FE: FieldExtension<D2, BaseField = F>,
+    P: PackedField<Scalar = FE>,
     C: GenericConfig<D, F = F>,
     S: Stark<F, D>,
     [(); S::COLUMNS]:,
@@ -28,7 +30,7 @@ pub(crate) fn eval_vanishing_poly<F, FE, C, S, const D: usize, const D2: usize>(
 {
     stark.eval_packed_generic(vars, consumer);
     if let Some(permutation_data) = permutation_data {
-        eval_permutation_checks::<F, FE, C, S, D, D2>(
+        eval_permutation_checks::<F, FE, P, C, S, D, D2>(
             stark,
             config,
             vars,

starky/src/verifier.rs

@@ -11,7 +11,7 @@ use plonky2::plonk::plonk_common::reduce_with_powers;
 
 use crate::config::StarkConfig;
 use crate::constraint_consumer::ConstraintConsumer;
-use crate::permutation::PermutationCheckData;
+use crate::permutation::PermutationCheckVars;
 use crate::proof::{StarkOpeningSet, StarkProofChallenges, StarkProofWithPublicInputs};
 use crate::stark::Stark;
 use crate::vanishing_poly::eval_vanishing_poly;
@@ -55,6 +55,7 @@ where
     [(); S::PUBLIC_INPUTS]:,
     [(); C::Hasher::HASH_SIZE]:,
 {
+    check_permutation_options(&stark, &proof_with_pis, &challenges)?;
     let StarkProofWithPublicInputs {
         proof,
         public_inputs,
@@ -90,12 +91,12 @@ where
         l_1,
         l_last,
     );
-    let permutation_data = stark.uses_permutation_args().then(|| PermutationCheckData {
+    let permutation_data = stark.uses_permutation_args().then(|| PermutationCheckVars {
         local_zs: permutation_zs.as_ref().unwrap().clone(),
         next_zs: permutation_zs_right.as_ref().unwrap().clone(),
         permutation_challenge_sets: challenges.permutation_challenge_sets.unwrap(),
     });
-    eval_vanishing_poly::<F, F::Extension, C, S, D, D>(
+    eval_vanishing_poly::<F, F::Extension, F::Extension, C, S, D, D>(
         &stark,
         config,
         vars,
@@ -153,7 +154,32 @@ fn eval_l_1_and_l_last<F: Field>(log_n: usize, x: F) -> (F, F) {
     (z_x * invs[0], z_x * invs[1])
 }
 
-/// Recover the length of the trace from a STARK proof and a STARK config.
+/// Utility function to check that all permutation data wrapped in `Option`s are `Some` iff
+/// the Stark uses a permutation argument.
+fn check_permutation_options<
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    const D: usize,
+>(
+    stark: &S,
+    proof_with_pis: &StarkProofWithPublicInputs<F, C, D>,
+    challenges: &StarkProofChallenges<F, D>,
+) -> Result<()> {
+    let options_is_some = [
+        proof_with_pis.proof.permutation_zs_cap.is_some(),
+        proof_with_pis.proof.openings.permutation_zs.is_some(),
+        proof_with_pis.proof.openings.permutation_zs_right.is_some(),
+        challenges.permutation_challenge_sets.is_some(),
+    ];
+    ensure!(
+        options_is_some
+            .into_iter()
+            .all(|b| b == stark.uses_permutation_args()),
+        "Permutation data doesn't match with Stark configuration."
+    );
+    Ok(())
+}
 
 #[cfg(test)]
 mod tests {