shockwave-plus/shockwave_plus/src/polynomial/ml_poly.rs
shockwave_plus/src/polynomial/ml_poly.rsRaw
1use tensor_pcs::EqPoly;
2
3use crate::FieldExt;
4
5#[derive(Clone, Debug)]
6pub struct MlPoly<F> {
7 pub evals: Vec<F>,
8 pub num_vars: usize,
9}
10
11impl<F: FieldExt> MlPoly<F> {
12 #[allow(dead_code)]
13 pub fn new(evals: Vec<F>) -> Self {
14 assert!(evals.len().is_power_of_two());
15 let num_vars = (evals.len() as f64).log2() as usize;
16 Self { evals, num_vars }
17 }
18
19 #[allow(dead_code)]
20 fn dot_prod(x: &[F], y: &[F]) -> F {
21 assert_eq!(x.len(), y.len());
22 let mut result = F::ZERO;
23 for i in 0..x.len() {
24 result += x[i] * y[i];
25 }
26 result
27 }
28
29 // Evaluate the multilinear extension of the polynomial `a`, at point `t`.
30 // `a` is in evaluation form.
31 // `t` should be in big-endian form.
32 #[allow(dead_code)]
33 pub fn eval(&self, t: &[F]) -> F {
34 let n = self.evals.len();
35 debug_assert_eq!((n as f64).log2() as usize, t.len());
36
37 let eq_evals = EqPoly::new(t.to_vec()).evals();
38
39 Self::dot_prod(&self.evals, &eq_evals)
40 }
41}
42
43#[cfg(test)]
44mod tests {
45 use super::*;
46 type F = halo2curves::secp256k1::Fp;
47 use halo2curves::ff::Field;
48
49 #[test]
50 fn test_ml_poly_eval() {
51 let num_vars = 4;
52 let num_evals = 2usize.pow(num_vars as u32);
53 let evals = (0..num_evals)
54 .map(|x| F::from(x as u64))
55 .collect::<Vec<F>>();
56
57 let ml_poly = MlPoly::new(evals.clone());
58 let eval_last = ml_poly.eval(&[F::ONE, F::ONE, F::ONE, F::ONE]);
59 assert_eq!(
60 eval_last,
61 evals[evals.len() - 1],
62 "The last evaluation is not correct"
63 );
64
65 let eval_first = ml_poly.eval(&[F::ZERO, F::ZERO, F::ZERO, F::ZERO]);
66 assert_eq!(eval_first, evals[0], "The first evaluation is not correct");
67
68 let eval_second = ml_poly.eval(&[F::ZERO, F::ZERO, F::ZERO, F::ONE]);
69 assert_eq!(
70 eval_second, evals[1],
71 "The second evaluation is not correct"
72 );
73 }
74}