zinc_poly/
utils.rs

1use crypto_primitives::SemiringConfig;
2use thiserror::Error;
3use zinc_utils::{add, mul, sub};
4
5use crate::mle::{DenseMultilinearExtension, dense::CollectDenseMleWithZero};
6
7/// A `enum` specifying the possible failure modes of the arithmetics.
8#[derive(Debug, Clone, Error)]
9pub enum ArithErrors {
10    #[error("Invalid parameters: {0}")]
11    InvalidParameters(String),
12}
13
14/// This function build the eq(x, r) polynomial for any given r.
15///
16/// Evaluate
17///      eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
18/// over r, which is
19///      eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
20pub fn build_eq_x_r<C: SemiringConfig>(
21    cfg: &C,
22    r: &[C::Element],
23) -> Result<DenseMultilinearExtension<C::Element>, ArithErrors> {
24    let evals = build_eq_x_r_vec(cfg, r)?;
25    let mle = DenseMultilinearExtension::from_evaluations_vec(r.len(), evals, cfg.zero());
26
27    Ok(mle)
28}
29
30/// This function builds the eq(x, r) polynomial for any given r, and outputs
31/// the evaluation of eq(x, r) in its vector form.
32///
33/// Evaluate
34///      $eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))$
35/// over r, which is
36///      $eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))$
37pub fn build_eq_x_r_vec<C: SemiringConfig>(
38    cfg: &C,
39    r: &[C::Element],
40) -> Result<Vec<C::Element>, ArithErrors> {
41    // we build eq(x,r) from its evaluations
42    // we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
43    // for example, with num_vars = 4, x is a binary vector of 4, then
44    //  0 0 0 0 -> (1-r0)   * (1-r1)    * (1-r2)    * (1-r3)
45    //  1 0 0 0 -> r0       * (1-r1)    * (1-r2)    * (1-r3)
46    //  0 1 0 0 -> (1-r0)   * r1        * (1-r2)    * (1-r3)
47    //  1 1 0 0 -> r0       * r1        * (1-r2)    * (1-r3)
48    //  ....
49    //  1 1 1 1 -> r0       * r1        * r2        * r3
50    // we will need 2^num_var evaluations
51
52    if r.is_empty() {
53        return Err(ArithErrors::InvalidParameters("r length is 0".to_owned()));
54    }
55
56    let one = cfg.one();
57    let mut eval = vec![one; 1 << r.len()];
58    let mut s = 1;
59    for ri in r {
60        for j in (0..s).rev() {
61            let prev = eval[j].clone();
62            let hi = cfg.mul(&prev, ri);
63            eval[j] = cfg.sub(&prev, &hi);
64            eval[add!(j, s)] = hi;
65        }
66        s = mul!(s, 2);
67    }
68
69    Ok(eval)
70}
71
72/// Build the shift selector MLE `next_c_mle(r, *)` with the first `num_vars`
73/// variables fixed to `r`.
74///
75/// For each `b in {0,1}^{num_vars}`:
76///   next_c_mle(b) = eq(r, b - c)   if b >= c
77///   next_c_mle(b) = 0              if b < c
78///
79/// Uses the identity `next_c_mle(r, b) = eq(r, b - c)` for `b >= c` and
80/// `0` for `b < c`.
81pub fn build_next_c_r_mle<C: SemiringConfig>(
82    cfg: &C,
83    r: &[C::Element],
84    c: usize,
85) -> Result<DenseMultilinearExtension<C::Element>, ArithErrors> {
86    let num_vars = r.len();
87    let n = 1 << num_vars;
88    assert!(c < n, "shift c={c} must be < domain size {n}");
89
90    let eq_r = build_eq_x_r(cfg, r)?;
91    if c == 0 {
92        return Ok(eq_r);
93    }
94
95    // next_c_mle(r, 0) = 0 for b < c
96    // next_c_mle(r, b - c) = eq(r, b - c) for b >= c
97    let mut evaluations = Vec::with_capacity(n);
98    evaluations.resize(c, cfg.zero());
99    evaluations.extend_from_slice(&eq_r.evaluations[..sub!(n, c)]);
100
101    Ok(DenseMultilinearExtension {
102        num_vars,
103        evaluations,
104    })
105}
106
107/// Evaluate eq polynomial.
108pub fn eq_eval<C: SemiringConfig>(
109    cfg: &C,
110    x: &[C::Element],
111    y: &[C::Element],
112) -> Result<C::Element, ArithErrors> {
113    if x.len() != y.len() {
114        return Err(ArithErrors::InvalidParameters(
115            "x and y have different length".to_string(),
116        ));
117    }
118
119    let one = cfg.one();
120    let mut res = one.clone();
121    for (xi, yi) in x.iter().zip(y.iter()) {
122        // xi * yi + (1 - xi) * (1 - yi) = 2 * xi * yi - xi - yi + 1
123        let xi_yi = cfg.mul(xi, yi);
124        let mut term = cfg.add(&xi_yi, &xi_yi);
125        cfg.sub_assign(&mut term, xi);
126        cfg.sub_assign(&mut term, yi);
127        cfg.add_assign(&mut term, &one);
128        cfg.mul_assign(&mut res, &term);
129    }
130
131    Ok(res)
132}
133
134/// Evaluate an MLE at a point using a precomputed eq table.
135///
136/// Given `evaluations[b]` and `eq_table[b] = eq(b, r)` (precomputed via
137/// [`build_eq_x_r_vec`]), returns `\sum_{b} eq_table[b] * evaluations[b]`.
138///
139/// This is equivalent to `DenseMultilinearExtension::evaluate`
140/// but avoids cloning the evaluation vector (the fix-variables algorithm is
141/// destructive). When multiple MLEs share the same evaluation point, build the
142/// eq table once and call this function for each MLE.
143pub fn mle_eval_with_eq_table<C: SemiringConfig>(
144    cfg: &C,
145    evaluations: &[C::Element],
146    eq_table: &[C::Element],
147) -> C::Element {
148    let mut acc = cfg.zero();
149    assert_eq!(
150        evaluations.len(),
151        eq_table.len(),
152        "evaluations and eq_table must have the same length"
153    );
154    for (eval, eq_val) in evaluations.iter().zip(eq_table.iter()) {
155        let term = cfg.mul(eq_val, eval);
156        cfg.add_assign(&mut acc, &term);
157    }
158    acc
159}
160
161/// Returns a multilinear polynomial in 2n variables that evaluates to 1
162/// if and only if the second n-bit vector is equal to the first vector plus one
163#[allow(clippy::arithmetic_side_effects)]
164pub fn next_mle<E: Clone>(
165    num_vars: u32,
166    zero: E,
167    one: E,
168) -> Result<DenseMultilinearExtension<E>, ArithErrors> {
169    if !num_vars.is_multiple_of(2) {
170        return Err(ArithErrors::InvalidParameters(
171            "num_vars must be even".to_string(),
172        ));
173    }
174
175    let mut mle = (0..1 << num_vars)
176        .map(|_| zero.clone())
177        .collect_dense_mle_with_zero(&zero);
178
179    let half_vars = num_vars / 2;
180
181    for i in 0usize..(1 << half_vars) - 1 {
182        let next = i + 1;
183
184        let i_concat_next = (next << half_vars) | i;
185
186        mle[i_concat_next] = one.clone();
187    }
188
189    Ok(mle)
190}
191
192/// Evaluates the next MLE in O(n), by reusing suffix equality and prefix carry
193/// products across carry positions.
194///
195/// Improved from O(n²) approach here: https://github.com/TomWambsgans/Whirlaway/blob/9e3592b/crates/air/src/utils.rs#L92
196///
197/// `next_mle(u, v) = 1` iff `Val(v) = Val(u) + 1` and `Val(u) < 2^n - 1`.
198///
199/// # Arguments
200/// - `u`: first n-bit vector (LE convention: index 0 = LSB).
201/// - `v`: second n-bit vector. Must have `v.len() == u.len()`.
202///
203/// # Algorithm
204/// Uses prefix/suffix products for O(n) evaluation:
205///   `next_mle(u, v) = sum_{j=0}^{n-1}
206///       [prod_{i<j} u_i * (1 - v_i)]      -- bits below j: were 1, flip to 0
207///     * (1 - u_j) * v_j                   -- bit j: 0 → 1
208///     * [prod_{i>j} eq(u_i, v_i)]`        -- bits above j: unchanged
209///
210/// # Panics
211/// Panics if `u.len() != v.len()`.
212#[allow(clippy::arithmetic_side_effects)]
213pub fn next_mle_eval<C: SemiringConfig>(cfg: &C, u: &[C::Element], v: &[C::Element]) -> C::Element {
214    let n = u.len();
215    assert_eq!(n, v.len(), "u and v must have the same length");
216    if n == 0 {
217        return cfg.zero();
218    }
219
220    let one = cfg.one();
221
222    // suffix_eq[j] = prod_{i=j}^{n-1} eq(u_i, v_i)
223    let mut suffix_eq = vec![one.clone(); n + 1];
224    for i in (0..n).rev() {
225        // eq(u_i, v_i) = u_i * v_i + (1 - u_i) * (1 - v_i)
226        let uv = cfg.mul(&u[i], &v[i]);
227        let eq_i = cfg.add(&uv, &cfg.mul(&cfg.sub(&one, &u[i]), &cfg.sub(&one, &v[i])));
228        suffix_eq[i] = cfg.mul(&suffix_eq[i + 1], &eq_i);
229    }
230
231    // prefix_carry accumulates prod_{i<j} u_i * (1 - v_i)
232    let mut prefix_carry = one.clone();
233    let mut result = cfg.zero();
234    for j in 0..n {
235        // prefix_carry * (1 - u_j) * v_j * suffix_eq[j + 1]
236        let mut term = cfg.mul(&prefix_carry, &cfg.sub(&one, &u[j]));
237        cfg.mul_assign(&mut term, &v[j]);
238        cfg.mul_assign(&mut term, &suffix_eq[j + 1]);
239        cfg.add_assign(&mut result, &term);
240
241        let carry = cfg.mul(&u[j], &cfg.sub(&one, &v[j]));
242        cfg.mul_assign(&mut prefix_carry, &carry);
243    }
244    result
245}
246
247#[cfg(test)]
248#[allow(
249    clippy::arithmetic_side_effects,
250    clippy::cast_possible_truncation,
251    clippy::needless_range_loop
252)]
253mod tests {
254    use crypto_bigint::{U128, const_monty_params};
255    use crypto_primitives::{FixedConfig, crypto_bigint_const_monty::ConstMontyField};
256    use num_traits::{One, Zero};
257    use proptest::{prelude::*, proptest};
258
259    use super::*;
260
261    const_monty_params!(Params, U128, "00000000b933426489189cb5b47d567f");
262
263    type F = ConstMontyField<Params, { U128::LIMBS }>;
264
265    const NUM_VARS: u32 = 8;
266
267    fn cfg() -> FixedConfig<F> {
268        FixedConfig::default()
269    }
270
271    #[test]
272    fn build_eq_x_r_vec_matches_product_formula() {
273        // For each b in {0,1}^n, eq(b, r) = prod_i (b_i * r_i + (1-b_i) * (1-r_i)).
274        // The helper uses little-endian indexing: bit i of the index is x_i.
275        let r: Vec<F> = (0..NUM_VARS).map(|i| F::from(i + 11)).collect();
276        let eq_vec = build_eq_x_r_vec(&cfg(), &r).unwrap();
277
278        let n = 1usize << NUM_VARS;
279        assert_eq!(eq_vec.len(), n);
280
281        let one = F::one();
282        for b in 0..n {
283            let mut expected = one;
284            for (i, ri) in r.iter().enumerate() {
285                let bit = (b >> i) & 1;
286                expected *= if bit == 1 { *ri } else { one - ri };
287            }
288            assert_eq!(eq_vec[b], expected, "mismatch at b={b}");
289        }
290    }
291
292    #[test]
293    fn build_eq_x_r_vec_basic() {
294        let r: [F; _] = [F::from(3_u64)];
295        let evals = build_eq_x_r_vec(&cfg(), &r).unwrap();
296        assert_eq!(evals, vec![F::one() - r[0], r[0]]);
297    }
298
299    #[test]
300    fn build_eq_x_r_vec_two_vars() {
301        let r: [F; _] = [F::from(2_u64), F::from(5_u64)];
302        let evals = build_eq_x_r_vec(&cfg(), &r).unwrap();
303        let e00 = (F::one() - r[0]) * (F::one() - r[1]);
304        let e01 = r[0] * (F::one() - r[1]);
305        let e10 = (F::one() - r[0]) * r[1];
306        let e11 = r[0] * r[1];
307        assert_eq!(evals, vec![e00, e01, e10, e11]);
308    }
309
310    #[test]
311    fn build_eq_x_r_error_on_empty() {
312        let r: [F; 0] = [];
313        let err = build_eq_x_r_vec(&cfg(), &r).unwrap_err();
314        let msg = format!("{err}");
315        assert!(msg.contains("Invalid parameters"));
316    }
317
318    #[test]
319    fn build_eq_x_r_mle_properties() {
320        let r: [F; _] = [F::from(7_u64), F::from(11_u64), F::from(13_u64)];
321        let mle = build_eq_x_r(&cfg(), &r).unwrap();
322        assert_eq!(mle.num_vars, r.len());
323        let evals = mle.evaluations;
324        let direct = build_eq_x_r_vec(&cfg(), &r).unwrap();
325        assert_eq!(evals, direct);
326    }
327
328    #[test]
329    fn next_mle_is_one_on_successors() {
330        let next_mle = next_mle(NUM_VARS, F::zero(), F::one()).unwrap();
331
332        for i in 0..(1 << ((NUM_VARS / 2) - 1)) {
333            let mut point: Vec<F> = (0..(NUM_VARS / 2))
334                .map(|j| {
335                    if i & (1 << j) == 0 {
336                        F::zero()
337                    } else {
338                        F::one()
339                    }
340                })
341                .collect();
342
343            point.extend((0..(NUM_VARS / 2)).map(|j| {
344                if (i + 1) & (1 << j) == 0 {
345                    F::zero()
346                } else {
347                    F::one()
348                }
349            }));
350
351            assert_eq!(next_mle.clone().evaluate(&cfg(), &point), Ok(F::one()));
352        }
353    }
354
355    #[test]
356    fn next_mle_is_one_only_on_successors() {
357        let next_mle = next_mle(NUM_VARS, F::zero(), F::one()).unwrap();
358
359        // The number of successors is (1 << (num_vars / 2)) - 1
360        // and we know the mle is one on them. So we need to check
361        // that it is one only on that many points.
362        assert_eq!(
363            next_mle.evaluations.iter().filter(|x| !x.is_zero()).count(),
364            (1 << (NUM_VARS / 2)) - 1
365        );
366    }
367
368    fn any_f() -> impl Strategy<Value = F> + 'static {
369        any::<u128>().prop_map(F::from)
370    }
371
372    fn point_n(n: usize) -> impl Strategy<Value = Vec<F>> {
373        prop::collection::vec(any_f(), n)
374    }
375
376    #[test]
377    fn next_mle_eval_coincides_with_next_mle_evaluated_at_successors() {
378        let next_mle = next_mle(NUM_VARS, F::zero(), F::one()).unwrap();
379
380        for i in 0..(1 << ((NUM_VARS / 2) - 1)) {
381            let mut point: Vec<F> = (0..(NUM_VARS / 2))
382                .map(|j| {
383                    if i & (1 << j) == 0 {
384                        F::zero()
385                    } else {
386                        F::one()
387                    }
388                })
389                .collect();
390
391            point.extend((0..(NUM_VARS / 2)).map(|j| {
392                if (i + 1) & (1 << j) == 0 {
393                    F::zero()
394                } else {
395                    F::one()
396                }
397            }));
398
399            let (u, v) = point.split_at(NUM_VARS as usize / 2);
400            assert_eq!(
401                next_mle.clone().evaluate(&cfg(), &point),
402                Ok(next_mle_eval(&cfg(), u, v))
403            );
404        }
405    }
406
407    proptest! {
408    #[test]
409    #[cfg_attr(miri, ignore)] // long running
410    fn prop_next_mle_eval_coincides_with_next_mle_evaluate_at_point(r in point_n(NUM_VARS as usize)) {
411        let next_mle = next_mle(NUM_VARS, F::zero(), F::one()).unwrap();
412
413        let (u, v) = r.split_at(NUM_VARS as usize / 2);
414        prop_assert_eq!(
415            next_mle.evaluate(&cfg(), &r),
416            Ok(next_mle_eval(&cfg(), u, v))
417        );
418    }
419    }
420
421    #[test]
422    #[cfg_attr(miri, ignore)] // long running
423    fn next_c_r_mle_c1_matches_shift_by_1() {
424        // c=1 should give the same result as the original build_next_r_mle
425        let num_vars: usize = 4;
426        let r: Vec<F> = (0..num_vars).map(|i| F::from((i + 3) as u32)).collect();
427
428        let next_1 = build_next_c_r_mle(&cfg(), &r, 1).unwrap();
429
430        // Manually build shift-by-1: evaluations[0] = 0, evaluations[b] = eq(r, b-1)
431        let eq_r = build_eq_x_r(&cfg(), &r).unwrap();
432        let n = 1 << num_vars;
433        let mut expected = vec![F::zero(); 1];
434        expected.extend_from_slice(&eq_r.evaluations[..n - 1]);
435
436        assert_eq!(next_1.evaluations, expected);
437    }
438
439    #[test]
440    #[cfg_attr(miri, ignore)] // long running
441    fn next_c_r_mle_c0_is_eq() {
442        // c=0 should return eq(r, b)
443        let num_vars: usize = 4;
444        let r: Vec<F> = (0..num_vars).map(|i| F::from((i + 7) as u32)).collect();
445
446        let next_0 = build_next_c_r_mle(&cfg(), &r, 0).unwrap();
447        let eq_r = build_eq_x_r(&cfg(), &r).unwrap();
448
449        assert_eq!(next_0.evaluations, eq_r.evaluations);
450    }
451
452    #[test]
453    #[cfg_attr(miri, ignore)] // long running
454    fn next_c_r_mle_has_correct_structure() {
455        // For any c, evaluations[b] should be:
456        //   0 for b < c
457        //   eq(r, b-c) for b >= c
458        let num_vars: usize = 4;
459        let n = 1 << num_vars;
460        let r: Vec<F> = (0..num_vars).map(|i| F::from((i + 5) as u32)).collect();
461
462        for c in [2, 3, 5, 7] {
463            let next_c = build_next_c_r_mle(&cfg(), &r, c).unwrap();
464            let eq_r = build_eq_x_r(&cfg(), &r).unwrap();
465
466            // First c entries should be zero
467            for b in 0..c {
468                assert!(
469                    next_c.evaluations[b].is_zero(),
470                    "c={c}, b={b}: expected zero"
471                );
472            }
473            // Remaining entries should match eq(r, b-c)
474            for b in c..n {
475                assert_eq!(
476                    next_c.evaluations[b],
477                    eq_r.evaluations[b - c],
478                    "c={c}, b={b}: mismatch"
479                );
480            }
481        }
482    }
483
484    proptest! {
485    #[test]
486    #[cfg_attr(miri, ignore)] // long running
487    fn prop_next_c_r_mle_evaluates_correctly(r in point_n(4), c in 1..15usize) {
488        // build_next_c_r_mle(r, c) evaluated at random point should equal
489        // the shift-c predicate: sum_b next_c(b) * eq(b, point)
490        let next_c = build_next_c_r_mle(&cfg(), &r, c).unwrap();
491        let eq_r = build_eq_x_r(&cfg(), &r).unwrap();
492
493        // Verify the table structure holds
494        let n = 1 << r.len();
495        for b in 0..c.min(n) {
496            prop_assert!(next_c.evaluations[b].is_zero());
497        }
498        for b in c..n {
499            prop_assert_eq!(&next_c.evaluations[b], &eq_r.evaluations[b - c]);
500        }
501    }
502    }
503}
zinc_poly/utils.rs

zinc_poly/
utils.rs
