diff --git a/src/core/air/evaluation.rs b/src/core/air/evaluation.rs index f617cc7cb..48f6a4bc1 100644 --- a/src/core/air/evaluation.rs +++ b/src/core/air/evaluation.rs @@ -3,8 +3,10 @@ //! Given a random alpha, the combined polynomial is defined as //! f(p) = sum_i alpha^{N-1-i} u_i (P). +use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::QM31; use crate::core::fields::Field; +use crate::core::poly::circle::{CircleDomain, CircleEvaluation, CirclePoly}; /// Accumulates evaluations of u_i(P0) at a single point. /// Computes f(P0), the combined polynomial at that point. @@ -54,14 +56,117 @@ impl PointEvaluationAccumulator { } } -pub struct DomainEvaluationAccumulator; +type Column = Vec; + +/// Accumulates evaluations of u_i(P), each at an evaluation domain of the size of that polynomial. +/// Computes the coefficients of f(P). +pub struct DomainEvaluationAccumulator { + random_coeff: QM31, + // Accumulated evaluations for each log_size. + // Each `sub_accumulation` holds `sum_{i=0}^{n-1} evaluation_i * alpha^(n-1-i)`, + // where `n` is the number of accumulated evaluations for this log_size. + sub_accumulations: Vec, + // Number of accumulated evaluations for each log_size. + n_cols_per_size: Vec, +} +impl DomainEvaluationAccumulator { + /// Creates a new accumulator. + /// `random_coeff` should be a secure random field element, drawn from the channel. + /// `max_log_size` is the maximum log_size of the accumulated evaluations. + pub fn new(random_coeff: QM31, max_log_size: u32) -> Self { + let max_log_size = max_log_size as usize; + Self { + random_coeff, + sub_accumulations: (0..(max_log_size + 1)) + .map(|n| vec![QM31::default(); 1 << n]) + .collect(), + n_cols_per_size: vec![0; max_log_size + 1], + } + } + + /// Gets accumulators for some sizes. + /// `n_cols_per_size` is an array of pairs (log_size, n_cols). + /// For each entry, a [ColumnAccumulator] is returned, expecting to accumulate `n_cols` + /// evaluations of size `log_size`. + /// The array size, `N`, is the number of different sizes. + pub fn columns( + &mut self, + n_cols_per_size: [(u32, usize); N], + ) -> [ColumnAccumulator<'_>; N] { + n_cols_per_size.iter().for_each(|(log_size, n_col)| { + self.n_cols_per_size[*log_size as usize] += n_col; + }); + self.sub_accumulations + .get_many_mut(n_cols_per_size.map(|(log_size, _)| log_size as usize)) + .unwrap_or_else(|e| panic!("invalid log_sizes: {}", e)) + .map(|c| ColumnAccumulator { + random_coeff: self.random_coeff, + col: c, + }) + } + + /// Computes f(P) as coefficients. + pub fn finalize(self) -> CirclePoly { + let mut res_coeffs = vec![QM31::default(); 1 << self.log_size()]; + let res_log_size = self.log_size(); + for (coeffs, n_cols) in self + .sub_accumulations + .into_iter() + .enumerate() + .map(|(log_size, values)| { + if log_size == 0 { + return values; + } + CircleEvaluation::new( + CircleDomain::constraint_evaluation_domain(log_size as u32), + values, + ) + .interpolate() + .extend(res_log_size) + .coeffs + }) + .zip(self.n_cols_per_size.iter()) + { + // Add poly.coeffs into coeffs, elementwise, inplace. + let multiplier = self.random_coeff.pow(*n_cols as u128); + res_coeffs + .iter_mut() + .zip(coeffs.iter()) + .for_each(|(res_coeff, current_coeff)| { + *res_coeff = *res_coeff * multiplier + *current_coeff + }); + } + + CirclePoly::new(res_coeffs) + } + + /// Returns the log_size of the resulting polynomial. + pub fn log_size(&self) -> u32 { + (self.sub_accumulations.len() - 1) as u32 + } +} + +/// An domain accumulator for polynomials of a single size. +pub struct ColumnAccumulator<'a> { + random_coeff: QM31, + col: &'a mut Column, +} +impl<'a> ColumnAccumulator<'a> { + pub fn accumulate(&mut self, index: usize, evaluation: BaseField) { + let accum = &mut self.col[index]; + *accum = *accum * self.random_coeff + evaluation; + } +} #[cfg(test)] mod tests { + use std::array; + use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use super::*; + use crate::core::circle::CirclePoint; use crate::core::fields::m31::{M31, P}; use crate::qm31; @@ -100,4 +205,69 @@ mod tests { assert_eq!(accumulator_res, res); } + + #[test] + fn test_domain_evaluation_accumulator() { + // Generate a vector of random sizes with a constant seed. + let rng = &mut StdRng::seed_from_u64(0); + const LOG_SIZE_MIN: u32 = 4; + const LOG_SIZE_BOUND: u32 = 10; + const MASK: u32 = P; + let log_sizes = (0..100) + .map(|_| rng.gen_range(LOG_SIZE_MIN..LOG_SIZE_BOUND)) + .collect::>(); + + // Generate random evaluations. + let evaluations = log_sizes + .iter() + .map(|log_size| { + (0..(1 << *log_size)) + .map(|_| M31::from_u32_unchecked(rng.gen::() & MASK)) + .collect::>() + }) + .collect::>(); + let alpha = qm31!(2, 3, 4, 5); + + // Use accumulator. + let mut accumulator = DomainEvaluationAccumulator::new(alpha, LOG_SIZE_BOUND); + let n_cols_per_size: [(u32, usize); (LOG_SIZE_BOUND - LOG_SIZE_MIN) as usize] = + array::from_fn(|i| { + let current_log_size = LOG_SIZE_MIN + i as u32; + let n_cols = log_sizes + .iter() + .copied() + .filter(|&log_size| log_size == current_log_size) + .count(); + (current_log_size, n_cols) + }); + let mut cols = accumulator.columns(n_cols_per_size); + for (log_size, evaluation) in log_sizes.iter().zip(evaluations.iter()) { + for (index, evaluation) in evaluation.iter().enumerate() { + cols[(log_size - LOG_SIZE_MIN) as usize].accumulate(index, *evaluation); + } + } + let accumulator_poly = accumulator.finalize(); + + // Pick an arbitrary sample point. + let point = CirclePoint::::get_point(98989892); + let accumulator_res = accumulator_poly.eval_at_point(point); + + // Sort evaluations by log_size. + let mut pairs = log_sizes.into_iter().zip(evaluations).collect::>(); + pairs.sort_by_key(|(log_size, _)| *log_size); + + // Use direct computation. + let mut res = QM31::default(); + for (log_size, values) in pairs.into_iter() { + res = res * alpha + + CircleEvaluation::new( + CircleDomain::constraint_evaluation_domain(log_size), + values, + ) + .interpolate() + .eval_at_point(point); + } + + assert_eq!(accumulator_res, res); + } } diff --git a/src/lib.rs b/src/lib.rs index f935ba523..8929b091c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,7 +4,8 @@ is_sorted, new_uninit, slice_group_by, - stdsimd + stdsimd, + get_many_mut )] pub mod commitment_scheme; pub mod core;