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