arnaucube
arnaucube/shockwave-plusPublic Self-hosted repository
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}