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}