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