diff --git a/starky/src/lib.rs b/starky/src/lib.rs index 1df9629e..51a73479 100644 --- a/starky/src/lib.rs +++ b/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; diff --git a/starky/src/permutation.rs b/starky/src/permutation.rs index 01cfa8bf..9306d0b2 100644 --- a/starky/src/permutation.rs +++ b/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 - // 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. - 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(); + let permutation_batches = get_permutation_batches( + &permutation_pairs, + &permutation_challenge_sets, + config.num_challenges, + stark.permutation_batch_size(), + ); permutation_batches .into_par_iter() @@ -178,3 +164,31 @@ pub(crate) fn get_n_permutation_challenge_sets>( .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], + num_challenges: usize, + batch_size: usize, +) -> Vec>> { + 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() +} diff --git a/starky/src/vanishing_poly.rs b/starky/src/vanishing_poly.rs new file mode 100644 index 00000000..6f7225b5 --- /dev/null +++ b/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( + stark: S, + config: &StarkConfig, + vars: StarkEvaluationVars, + local_zs: &[F::Extension], + next_zs: &[F::Extension], + mut consumer: ConstraintConsumer, + permutation_challenge_sets: &[PermutationChallenge], +) where + F: RichField + Extendable, + C: GenericConfig, + S: Stark, + [(); S::COLUMNS]:, + [(); S::PUBLIC_INPUTS]:, +{ + stark.eval_packed_base(vars, &mut consumer); +} + +fn eval_permutation_checks( + stark: S, + config: &StarkConfig, + vars: StarkEvaluationVars, + local_zs: &[F::Extension], + next_zs: &[F::Extension], + mut consumer: ConstraintConsumer, + permutation_challenge_sets: &[PermutationChallenge], +) where + F: RichField + Extendable, + C: GenericConfig, + S: Stark, + [(); 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() +}