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