Precompute eq(b, r)s before evaluating column MLEs in ideal check

piop/src/ideal_check/combined_poly_builder.rs

@@ -11,14 +11,18 @@ 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,
     degree_counter::{count_constraint_degrees, count_max_degree},
     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,55 +230,52 @@ 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<DynamicPolynomialF<F>>,
-                          d: usize,
+                          max_num_coeffs: usize,
                           shift: usize|
-     -> Result<F, EvaluationError> {
-        let coeff_evals: Vec<F::Inner> = (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<F> {
+        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] += &prod;
+                    }
+
+                    acc
+                },
+            );
+        res.trim();
+        res
     };
 
     // Evaluate up (all columns, shift=0).
     let up_evals: Vec<DynamicPolynomialF<F>> = cfg_iter!(trace_matrix)
-        .map(|col| {
-            let coeffs: Vec<F> = (0..max_num_coeffs)
-                .map(|d| eval_coeff_mle(col, d, 0))
-                .collect::<Result<_, _>>()?;
-            Ok(DynamicPolynomialF::new_trimmed(coeffs))
-        })
+        .map(|col| Ok(eval_coeff_mle(col, max_num_coeffs, 0)))
         .collect::<Result<Vec<_>, EvaluationError>>()?;
 
     // Evaluate down (only shifted columns, per-spec shift amount).
     let sorted_shifts = uair_sig.shifts();
     let down_evals: Vec<DynamicPolynomialF<F>> = cfg_iter!(sorted_shifts)
         .map(|spec| {
             let col = &trace_matrix[spec.source_col()];
-            let coeffs: Vec<F> = (0..max_num_coeffs)
-                .map(|d| eval_coeff_mle(col, d, spec.shift_amount()))
-                .collect::<Result<_, _>>()?;
-            Ok(DynamicPolynomialF::new_trimmed(coeffs))
+            Ok(eval_coeff_mle(col, max_num_coeffs, spec.shift_amount()))
         })
         .collect::<Result<Vec<_>, EvaluationError>>()?;
 

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<usize> },
+    #[error("Inner product error: {0}")]
+    InnerProductError(InnerProductError),
+}
+
+impl From<InnerProductError> for EvaluationError {
+    fn from(inner_product_error: InnerProductError) -> Self {
+        Self::InnerProductError(inner_product_error)
+    }
 }

poly/src/univariate/binary_ref.rs

@@ -258,7 +258,7 @@ impl<R: Clone + Zero + One + CheckedAdd + CheckedMul, const DEGREE_PLUS_ONE: usi
             .iter()
             .try_fold(
                 (self.0.coeffs[0].widen::<R>(), 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() {

poly/src/utils.rs

@@ -195,6 +195,36 @@ where
     Ok(())
 }
 
+pub fn precompute_eq_r_b_inner<F>(point: &[F], field_cfg: &F::Config) -> Vec<F::Inner>
+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<R: Semiring>(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<Value = Vec<F>> {
+        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);
+    }
+    }
 }

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<F>(
+        lhs: &[F::Inner],
+        rhs: &[F::Inner],
+        field_cfg: &F::Config,
+    ) -> Result<F, InnerProductError>
+    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).