From d4a6b35436355c033c85e43468e337db6521ccf2 Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Thu, 5 May 2022 22:21:09 +0200 Subject: [PATCH] More cleaning --- starky2/src/cross_table_lookups.rs | 170 +++++++++++++++++++++++++++++ starky2/src/lib.rs | 1 + starky2/src/prover.rs | 89 +-------------- 3 files changed, 172 insertions(+), 88 deletions(-) create mode 100644 starky2/src/cross_table_lookups.rs diff --git a/starky2/src/cross_table_lookups.rs b/starky2/src/cross_table_lookups.rs new file mode 100644 index 00000000..b35733a2 --- /dev/null +++ b/starky2/src/cross_table_lookups.rs @@ -0,0 +1,170 @@ +use plonky2::field::extension_field::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; +use plonky2::plonk::config::GenericConfig; +use plonky2::plonk::plonk_common::reduce_with_powers; + +use crate::config::StarkConfig; +use crate::prover::CrossTableLookup; + +/// Lookup data for one table. +#[derive(Clone)] +pub struct LookupData { + zs_beta_gammas: Vec<(PolynomialValues, F, F)>, +} + +impl Default for LookupData { + fn default() -> Self { + Self { + zs_beta_gammas: Vec::new(), + } + } +} + +impl LookupData { + pub fn is_empty(&self) -> bool { + self.zs_beta_gammas.is_empty() + } + + pub fn z_polys(&self) -> Vec> { + self.zs_beta_gammas + .iter() + .map(|(p, _, _)| p.clone()) + .collect() + } +} + +pub fn cross_table_lookup_zs, const D: usize>( + config: &StarkConfig, + trace_poly_values: &[Vec>], + cross_table_lookups: &[CrossTableLookup], + challenger: &mut Challenger, +) -> Vec> { + cross_table_lookups.iter().fold( + vec![LookupData::default(); trace_poly_values.len()], + |mut acc, cross_table_lookup| { + let CrossTableLookup { + looking_table, + looking_columns, + looked_table, + looked_columns, + .. + } = cross_table_lookup; + + for _ in 0..config.num_challenges { + let beta = challenger.get_challenge(); + let gamma = challenger.get_challenge(); + let z_looking = partial_products( + &trace_poly_values[*looking_table as usize], + &looking_columns, + beta, + gamma, + ); + let z_looked = partial_products( + &trace_poly_values[*looked_table as usize], + &looked_columns, + beta, + gamma, + ); + + acc[*looking_table as usize] + .zs_beta_gammas + .push((z_looking, beta, gamma)); + acc[*looked_table as usize] + .zs_beta_gammas + .push((z_looked, beta, gamma)); + } + acc + }, + ) +} + +fn partial_products( + trace: &[PolynomialValues], + columns: &[usize], + beta: F, + gamma: F, +) -> PolynomialValues { + let mut partial_prod = F::ONE; + let mut res = Vec::new(); + for i in 0..trace[0].len() { + partial_prod *= + gamma + reduce_with_powers(columns.iter().map(|&j| &trace[i].values[j]), beta); + res.push(partial_prod); + } + res.into() +} + +pub struct CTLCheckVars +where + F: Field, + FE: FieldExtension, + P: PackedField, +{ + pub(crate) local_zs: Vec

, + pub(crate) next_zs: Vec

, + pub(crate) permutation_challenge_sets: Vec>, +} + +pub(crate) fn eval_permutation_checks( + stark: &S, + config: &StarkConfig, + vars: StarkEvaluationVars, + permutation_data: PermutationCheckVars, + consumer: &mut ConstraintConsumer

, +) where + F: RichField + Extendable, + FE: FieldExtension, + P: PackedField, + C: GenericConfig, + S: Stark, +{ + let PermutationCheckVars { + local_zs, + next_zs, + permutation_challenge_sets, + } = permutation_data; + + // Check that Z(1) = 1; + for &z in &local_zs { + consumer.constraint_first_row(z - FE::ONE); + } + + 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. + for (i, instances) in permutation_batches.iter().enumerate() { + // Z(gx) * down = Z x * up + let (reduced_lhs, reduced_rhs): (Vec

, Vec

) = instances + .iter() + .map(|instance| { + let PermutationInstance { + pair: PermutationPair { column_pairs }, + challenge: PermutationChallenge { beta, gamma }, + } = instance; + let mut factor = ReducingFactor::new(*beta); + let (lhs, rhs): (Vec<_>, Vec<_>) = column_pairs + .iter() + .map(|&(i, j)| (vars.local_values[i], vars.local_values[j])) + .unzip(); + ( + factor.reduce_ext(lhs.into_iter()) + FE::from_basefield(*gamma), + factor.reduce_ext(rhs.into_iter()) + FE::from_basefield(*gamma), + ) + }) + .unzip(); + let constraint = next_zs[i] * reduced_rhs.into_iter().product::

() + - local_zs[i] * reduced_lhs.into_iter().product::

(); + consumer.constraint(constraint); + } +} diff --git a/starky2/src/lib.rs b/starky2/src/lib.rs index 0924600d..03ce7076 100644 --- a/starky2/src/lib.rs +++ b/starky2/src/lib.rs @@ -5,6 +5,7 @@ pub mod config; pub mod constraint_consumer; +pub mod cross_table_lookups; mod get_challenges; pub mod mock_stark; pub mod permutation; diff --git a/starky2/src/prover.rs b/starky2/src/prover.rs index cc19edbc..d998b074 100644 --- a/starky2/src/prover.rs +++ b/starky2/src/prover.rs @@ -22,6 +22,7 @@ use rayon::prelude::*; use crate::config::StarkConfig; use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; +use crate::cross_table_lookups::{cross_table_lookup_zs, LookupData}; use crate::permutation::PermutationCheckVars; use crate::permutation::{ compute_permutation_z_polys, get_n_permutation_challenge_sets, PermutationChallengeSet, @@ -381,94 +382,6 @@ where }) } -/// Lookup data for one table. -#[derive(Clone)] -struct LookupData { - zs_beta_gammas: Vec<(PolynomialValues, F, F)>, -} - -impl Default for LookupData { - fn default() -> Self { - Self { - zs_beta_gammas: Vec::new(), - } - } -} - -impl LookupData { - fn is_empty(&self) -> bool { - self.zs_beta_gammas.is_empty() - } - - fn z_polys(&self) -> Vec> { - self.zs_beta_gammas - .iter() - .map(|(p, _, _)| p.clone()) - .collect() - } -} - -fn cross_table_lookup_zs, const D: usize>( - config: &StarkConfig, - trace_poly_values: &[Vec>], - cross_table_lookups: &[CrossTableLookup], - challenger: &mut Challenger, -) -> Vec> { - cross_table_lookups.iter().fold( - vec![LookupData::default(); trace_poly_values.len()], - |mut acc, cross_table_lookup| { - let CrossTableLookup { - looking_table, - looking_columns, - looked_table, - looked_columns, - .. - } = cross_table_lookup; - - for _ in 0..config.num_challenges { - let beta = challenger.get_challenge(); - let gamma = challenger.get_challenge(); - let z_looking = partial_products( - &trace_poly_values[*looking_table as usize], - &looking_columns, - beta, - gamma, - ); - let z_looked = partial_products( - &trace_poly_values[*looked_table as usize], - &looked_columns, - beta, - gamma, - ); - - acc[*looking_table as usize] - .zs_beta_gammas - .push((z_looking, beta, gamma)); - acc[*looked_table as usize] - .zs_beta_gammas - .push((z_looked, beta, gamma)); - } - acc - }, - ) -} - -fn partial_products( - trace: &[PolynomialValues], - columns: &[usize], - beta: F, - gamma: F, -) -> PolynomialValues { - let mut partial_prod = F::ONE; - let mut res = Vec::new(); - for i in 0..trace[0].len() { - partial_prod *= - gamma + reduce_with_powers(columns.iter().map(|&j| &trace[i].values[j]), beta); - res.push(partial_prod); - } - res.into() -} - /// Computes the quotient polynomials `(sum alpha^i C_i(x)) / Z_H(x)` for `alpha` in `alphas`, /// where the `C_i`s are the Stark constraints. fn compute_quotient_polys<'a, F, P, C, S, const D: usize>(