First pass

starky/src/lib.rs

@@ -14,6 +14,7 @@ pub mod prover;
 pub mod recursive_verifier;
 pub mod stark;
 pub mod stark_testing;
+pub mod vanishing_poly;
 pub mod vars;
 pub mod verifier;
 

starky/src/permutation.rs

@@ -62,26 +62,12 @@ where
         stark.permutation_batch_size(),
     );
 
-    // Get a list of instances of our batch-permutation argument. These are permutation arguments
+    let permutation_batches = get_permutation_batches(
-    // where the same `Z(x)` polynomial is used to check more than one permutation.
+        &permutation_pairs,
-    // Before batching, each permutation pair leads to `num_challenges` permutation arguments, so we
+        &permutation_challenge_sets,
-    // start with the cartesian product of `permutation_pairs` and `0..num_challenges`. Then we
+        config.num_challenges,
-    // chunk these arguments based on our batch size.
+        stark.permutation_batch_size(),
-    let permutation_batches = permutation_pairs
+    );
-        .iter()
-        .cartesian_product(0..config.num_challenges)
-        .chunks(stark.permutation_batch_size())
-        .into_iter()
-        .map(|batch| {
-            batch
-                .enumerate()
-                .map(|(i, (pair, chal))| {
-                    let challenge = permutation_challenge_sets[i].challenges[chal];
-                    PermutationInstance { pair, challenge }
-                })
-                .collect_vec()
-        })
-        .collect_vec();
 
     permutation_batches
         .into_par_iter()
@@ -178,3 +164,31 @@ pub(crate) fn get_n_permutation_challenge_sets<F: RichField, H: Hasher<F>>(
         .map(|_| get_permutation_challenge_set(challenger, num_challenges))
         .collect()
 }
+
+/// Get a list of instances of our batch-permutation argument. These are permutation arguments
+/// where the same `Z(x)` polynomial is used to check more than one permutation.
+/// Before batching, each permutation pair leads to `num_challenges` permutation arguments, so we
+/// start with the cartesian product of `permutation_pairs` and `0..num_challenges`. Then we
+/// chunk these arguments based on our batch size.
+pub(crate) fn get_permutation_batches<'a, F: Field>(
+    permutation_pairs: &'a [PermutationPair],
+    permutation_challenge_sets: &[PermutationChallengeSet<F>],
+    num_challenges: usize,
+    batch_size: usize,
+) -> Vec<Vec<PermutationInstance<'a, F>>> {
+    permutation_pairs
+        .iter()
+        .cartesian_product(0..num_challenges)
+        .chunks(batch_size)
+        .into_iter()
+        .map(|batch| {
+            batch
+                .enumerate()
+                .map(|(i, (pair, chal))| {
+                    let challenge = permutation_challenge_sets[i].challenges[chal];
+                    PermutationInstance { pair, challenge }
+                })
+                .collect_vec()
+        })
+        .collect()
+}

starky/src/vanishing_poly.rs

@@ -0,0 +1,60 @@
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::plonk::config::GenericConfig;
+use rayon::prelude::*;
+
+use crate::config::StarkConfig;
+use crate::constraint_consumer::ConstraintConsumer;
+use crate::permutation::{get_permutation_batches, PermutationChallenge};
+use crate::stark::Stark;
+use crate::vars::StarkEvaluationVars;
+
+pub(crate) fn eval_vanishing_poly<F, C, S, const D: usize>(
+    stark: S,
+    config: &StarkConfig,
+    vars: StarkEvaluationVars<F, F, S::COLUMNS, S::PUBLIC_INPUTS>,
+    local_zs: &[F::Extension],
+    next_zs: &[F::Extension],
+    mut consumer: ConstraintConsumer<F>,
+    permutation_challenge_sets: &[PermutationChallenge<F>],
+) where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    [(); S::COLUMNS]:,
+    [(); S::PUBLIC_INPUTS]:,
+{
+    stark.eval_packed_base(vars, &mut consumer);
+}
+
+fn eval_permutation_checks<F, C, S, const D: usize>(
+    stark: S,
+    config: &StarkConfig,
+    vars: StarkEvaluationVars<F::Extension, F::Extension, S::COLUMNS, S::PUBLIC_INPUTS>,
+    local_zs: &[F::Extension],
+    next_zs: &[F::Extension],
+    mut consumer: ConstraintConsumer<F>,
+    permutation_challenge_sets: &[PermutationChallenge<F>],
+) where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    [(); S::COLUMNS]:,
+    [(); S::PUBLIC_INPUTS]:,
+{
+    let permutation_pairs = stark.permutation_pairs();
+
+    let permutation_batches = get_permutation_batches(
+        &permutation_pairs,
+        &permutation_challenge_sets,
+        config.num_challenges,
+        stark.permutation_batch_size(),
+    );
+
+    // Each zs value corresponds to a permutation batch.
+    permutation_batches
+        .into_par_iter()
+        .map(|instances| compute_permutation_z_poly(&instances, trace_poly_values))
+        .collect()
+}