diff --git a/piop/src/ideal_check/combined_poly_builder.rs b/piop/src/ideal_check/combined_poly_builder.rs index 8ba8da89..af49f59d 100644 --- a/piop/src/ideal_check/combined_poly_builder.rs +++ b/piop/src/ideal_check/combined_poly_builder.rs @@ -11,6 +11,7 @@ use zinc_poly::{ DenseMultilinearExtension, MultilinearExtensionWithConfig, dense::CollectDenseMleWithZero, }, univariate::dynamic::over_field::DynamicPolynomialF, + utils::precompute_eq_r_b_inner, }; use zinc_uair::{ ColumnLayout, ConstraintBuilder, TraceRow, Uair, @@ -18,7 +19,10 @@ use zinc_uair::{ ideal::ImpossibleIdeal, }; use zinc_utils::{ - cfg_into_iter, cfg_iter, from_ref::FromRef, inner_transparent_field::InnerTransparentField, + cfg_into_iter, cfg_iter, + from_ref::FromRef, + inner_product::{InnerProductError, MBSInnerProduct}, + inner_transparent_field::InnerTransparentField, }; /// Given a UAIR `U` and a trace `trace` this function @@ -226,44 +230,44 @@ where let uair_sig = U::signature(); let down_layout = uair_sig.down_cols().as_column_layout(); + let precomputed_eqs = precompute_eq_r_b_inner(evaluation_point, field_cfg); + // Helper: evaluate one column's coefficient-d MLE at `evaluation_point`, // reading row `i + shift` (zero-padded beyond trace length). let eval_coeff_mle = |col: &DenseMultilinearExtension>, - d: usize, + max_num_coeffs: usize, shift: usize| - -> Result { - let coeff_evals: Vec = (0..num_rows) - .map(|i| { - // Two conditions needed: - // 1. i < num_rows - 1: zero out the last row for all columns (both up and down) - // to match the combined poly builder's explicit zero-padding at row N-1. - // 2. i + shift < num_rows: prevent OOB access for shifts > 0. - if i < num_rows - 1 && i + shift < num_rows { - col.evaluations[i + shift] - .coeffs - .get(d) - .map(|c| c.inner().clone()) - .unwrap_or_else(|| zero_inner.clone()) - } else { - zero_inner.clone() - } + -> DynamicPolynomialF { + let mut res = precomputed_eqs + .iter() + .zip(if shift > 0 { + col[shift..].iter() + } else { + col[..num_rows - 1].iter() }) - .collect(); - let coeff_mle = DenseMultilinearExtension { - evaluations: coeff_evals, - num_vars, - }; - coeff_mle.evaluate_with_config(evaluation_point, field_cfg) + .fold( + DynamicPolynomialF { + coeffs: vec![F::zero_with_cfg(field_cfg); max_num_coeffs], + }, + |mut acc, (eq, poly)| { + let mut prod = F::zero_with_cfg(field_cfg); + + for (i, poly_coeff) in poly.coeffs.iter().enumerate() { + *prod.inner_mut() = poly_coeff.inner().clone(); + prod.mul_assign_by_inner(eq); + acc.coeffs[i] += ∏ + } + + acc + }, + ); + res.trim(); + res }; // Evaluate up (all columns, shift=0). let up_evals: Vec> = cfg_iter!(trace_matrix) - .map(|col| { - let coeffs: Vec = (0..max_num_coeffs) - .map(|d| eval_coeff_mle(col, d, 0)) - .collect::>()?; - Ok(DynamicPolynomialF::new_trimmed(coeffs)) - }) + .map(|col| Ok(eval_coeff_mle(col, max_num_coeffs, 0))) .collect::, EvaluationError>>()?; // Evaluate down (only shifted columns, per-spec shift amount). @@ -271,10 +275,7 @@ where let down_evals: Vec> = cfg_iter!(sorted_shifts) .map(|spec| { let col = &trace_matrix[spec.source_col()]; - let coeffs: Vec = (0..max_num_coeffs) - .map(|d| eval_coeff_mle(col, d, spec.shift_amount())) - .collect::>()?; - Ok(DynamicPolynomialF::new_trimmed(coeffs)) + Ok(eval_coeff_mle(col, max_num_coeffs, spec.shift_amount())) }) .collect::, EvaluationError>>()?; diff --git a/poly/src/lib.rs b/poly/src/lib.rs index a49bd77f..0a642d28 100644 --- a/poly/src/lib.rs +++ b/poly/src/lib.rs @@ -4,6 +4,7 @@ pub mod utils; pub mod zero_degree; use thiserror::Error; +use zinc_utils::inner_product::InnerProductError; /// Polynomial with coefficients of type `C` and degree bounded by /// `DEGREE_BOUND`. @@ -35,4 +36,12 @@ pub enum EvaluationError { EmptyPolynomial, #[error("Unsupported constraint degrees: {degrees:?}")] UnsupportedConstraintDegrees { degrees: Vec }, + #[error("Inner product error: {0}")] + InnerProductError(InnerProductError), +} + +impl From for EvaluationError { + fn from(inner_product_error: InnerProductError) -> Self { + Self::InnerProductError(inner_product_error) + } } diff --git a/poly/src/univariate/binary_ref.rs b/poly/src/univariate/binary_ref.rs index 1370a5cc..62ae60b1 100644 --- a/poly/src/univariate/binary_ref.rs +++ b/poly/src/univariate/binary_ref.rs @@ -258,7 +258,7 @@ impl(), R::one()), - |(mut acc, mut pow), coeff| { + |(mut acc, mut pow), coeff| -> Result<(R, R), EvaluationError> { pow = pow.checked_mul(point).ok_or(EvaluationError::Overflow)?; if coeff.inner() { diff --git a/poly/src/utils.rs b/poly/src/utils.rs index 03696c32..78cd5901 100644 --- a/poly/src/utils.rs +++ b/poly/src/utils.rs @@ -195,6 +195,36 @@ where Ok(()) } +pub fn precompute_eq_r_b_inner(point: &[F], field_cfg: &F::Config) -> Vec +where + F: InnerTransparentField, +{ + if point.is_empty() { + return vec![]; + } + + let one = F::one_with_cfg(point[0].cfg()); + let mut res = vec![F::zero_with_cfg(field_cfg).into_inner(); 1 << point.len()]; + + res[0] = one.inner().clone(); + + let mut a = one.clone(); + let mut b = one.clone(); + + for (i, r) in point.iter().enumerate() { + for j in (0..1 << i).rev() { + *a.inner_mut() = r.inner().clone(); + a.mul_assign_by_inner(&res[j]); + *b.inner_mut() = F::sub_inner(&res[j], a.inner(), field_cfg); + + res[j] = b.inner().clone(); + res[j | (1 << i)] = a.inner().clone(); + } + } + + res +} + /// Build the shift selector MLE `next_c_mle(r, *)` with the first `num_vars` /// variables fixed to `r`. /// @@ -364,8 +394,10 @@ pub fn next_mle_eval(u: &[R], v: &[R], zero: R, one: R) -> R { mod tests { use crypto_bigint::{U128, const_monty_params}; use crypto_primitives::{IntoWithConfig, crypto_bigint_const_monty::ConstMontyField}; + use itertools::Itertools; use num_traits::One; use proptest::{prelude::*, proptest}; + use zinc_utils::inner_product::MBSInnerProduct; use crate::mle::MultilinearExtensionWithConfig; @@ -428,6 +460,10 @@ mod tests { prop::collection::vec(any_f(()), n) } + fn mle_evals_n_vars(n: usize) -> impl Strategy> { + prop::collection::vec(any_f(()), n) + } + #[test] fn next_mle_eval_coincides_with_next_mle_evaluated_at_successors() { let next_mle = next_mle_inner(NUM_VARS, F::zero(), F::one()).unwrap(); @@ -555,4 +591,31 @@ mod tests { } } } + + proptest! { + #[test] + fn prop_precompute_eq_r_b_inner_correct(r in point_n(4)) { + let precomputed_eq_r_bs = precompute_eq_r_b_inner(&r, &()); + + for (b, eq_b) in precomputed_eq_r_bs.iter().enumerate() { + let point_from_b = (0..4).map(|i| if b & (1 << i) == 0 { F::zero() } else { F::one() }).collect_vec(); + let eq_b_built_at_r = build_eq_x_r(&point_from_b, &()).unwrap().evaluate(&r, F::zero()).unwrap(); + prop_assert_eq!(eq_b, eq_b_built_at_r.inner()); + } + } + } + + proptest! { + #[test] + fn prop_precompute_eq_r_b_inner_compare_with_mle_eval(r in point_n(4), mle_evals in mle_evals_n_vars(4)) { + let precomputed_eq_r_bs = precompute_eq_r_b_inner(&r, &()).into_iter().map(F::new_unchecked).collect_vec(); + + let mle = DenseMultilinearExtension::from_evaluations_vec(4, mle_evals, F::zero()); + + let mle_val_expected = mle.evaluate(&r, F::zero()).unwrap(); + let mle_val_computed = MBSInnerProduct::inner_product_field(&mle, &precomputed_eq_r_bs, F::zero()).unwrap(); + + prop_assert_eq!(mle_val_computed, mle_val_expected); + } + } } diff --git a/utils/src/inner_product.rs b/utils/src/inner_product.rs index 53effe36..b2addf59 100644 --- a/utils/src/inner_product.rs +++ b/utils/src/inner_product.rs @@ -1,6 +1,8 @@ -use crate::{from_ref::FromRef, mul_by_scalar::MulByScalar}; +use crate::{ + from_ref::FromRef, inner_transparent_field::InnerTransparentField, mul_by_scalar::MulByScalar, +}; use crypto_primitives::{FromWithConfig, PrimeField, boolean::Boolean}; -use num_traits::CheckedAdd; +use num_traits::{CheckedAdd, Zero}; use thiserror::Error; /// A trait for inner product algorithms implementations. @@ -84,6 +86,33 @@ impl MBSInnerProduct { acc + product })) } + + pub fn inner_product_inner_field( + lhs: &[F::Inner], + rhs: &[F::Inner], + field_cfg: &F::Config, + ) -> Result + where + F: InnerTransparentField, + { + if lhs.len() != rhs.len() { + return Err(InnerProductError::LengthMismatch { + lhs: lhs.len(), + rhs: rhs.len(), + }); + } + + Ok(lhs + .iter() + .zip(rhs) + .fold(F::zero_with_cfg(field_cfg), |acc, (a, r)| { + let mut product = F::new_unchecked_with_cfg(a.clone(), field_cfg); + + product.mul_assign_by_inner(r); + + acc + product + })) + } } /// The inner product for vectors of length 1 (a.k.a. scalars).