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}