1//! This module defines basic types related to polynomials
  2use core::ops::Index;
  3use ff::PrimeField;
  4use rayon::prelude::*;
  5
  6pub(crate) struct EqPolynomial<Scalar: PrimeField> {
  7  r: Vec<Scalar>,
  8}
  9
 10impl<Scalar: PrimeField> EqPolynomial<Scalar> {
 11  /// Creates a new polynomial from its succinct specification
 12  pub fn new(r: Vec<Scalar>) -> Self {
 13    EqPolynomial { r }
 14  }
 15
 16  /// Evaluates the polynomial at the specified point
 17  pub fn evaluate(&self, rx: &[Scalar]) -> Scalar {
 18    assert_eq!(self.r.len(), rx.len());
 19    (0..rx.len())
 20      .map(|i| rx[i] * self.r[i] + (Scalar::one() - rx[i]) * (Scalar::one() - self.r[i]))
 21      .fold(Scalar::one(), |acc, item| acc * item)
 22  }
 23
 24  pub fn evals(&self) -> Vec<Scalar> {
 25    let ell = self.r.len();
 26    let mut evals: Vec<Scalar> = vec![Scalar::zero(); (2_usize).pow(ell as u32)];
 27    let mut size = 1;
 28    evals[0] = Scalar::one();
 29
 30    for r in self.r.iter().rev() {
 31      let (evals_left, evals_right) = evals.split_at_mut(size);
 32      let (evals_right, _) = evals_right.split_at_mut(size);
 33
 34      evals_left
 35        .par_iter_mut()
 36        .zip(evals_right.par_iter_mut())
 37        .for_each(|(x, y)| {
 38          *y = *x * r;
 39          *x -= &*y;
 40        });
 41
 42      size *= 2;
 43    }
 44    evals
 45  }
 46}
 47
 48#[derive(Debug)]
 49pub(crate) struct MultilinearPolynomial<Scalar: PrimeField> {
 50  num_vars: usize, // the number of variables in the multilinear polynomial
 51  Z: Vec<Scalar>,  // evaluations of the polynomial in all the 2^num_vars Boolean inputs
 52}
 53
 54impl<Scalar: PrimeField> MultilinearPolynomial<Scalar> {
 55  pub fn new(Z: Vec<Scalar>) -> Self {
 56    assert_eq!(Z.len(), (2_usize).pow((Z.len() as f64).log2() as u32));
 57    MultilinearPolynomial {
 58      num_vars: (Z.len() as f64).log2() as usize,
 59      Z,
 60    }
 61  }
 62
 63  pub fn get_num_vars(&self) -> usize {
 64    self.num_vars
 65  }
 66
 67  pub fn len(&self) -> usize {
 68    self.Z.len()
 69  }
 70
 71  pub fn bound_poly_var_top(&mut self, r: &Scalar) {
 72    let n = self.len() / 2;
 73
 74    let (left, right) = self.Z.split_at_mut(n);
 75    let (right, _) = right.split_at(n);
 76
 77    left
 78      .par_iter_mut()
 79      .zip(right.par_iter())
 80      .for_each(|(a, b)| {
 81        *a += *r * (*b - *a);
 82      });
 83
 84    self.Z.resize(n, Scalar::zero());
 85    self.num_vars -= 1;
 86  }
 87
 88  // returns Z(r) in O(n) time
 89  pub fn evaluate(&self, r: &[Scalar]) -> Scalar {
 90    // r must have a value for each variable
 91    assert_eq!(r.len(), self.get_num_vars());
 92    let chis = EqPolynomial::new(r.to_vec()).evals();
 93    assert_eq!(chis.len(), self.Z.len());
 94
 95    (0..chis.len())
 96      .into_par_iter()
 97      .map(|i| chis[i] * self.Z[i])
 98      .reduce(Scalar::zero, |x, y| x + y)
 99  }
100}
101
102impl<Scalar: PrimeField> Index<usize> for MultilinearPolynomial<Scalar> {
103  type Output = Scalar;
104
105  #[inline(always)]
106  fn index(&self, _index: usize) -> &Scalar {
107    &(self.Z[_index])
108  }
109}
110
111pub(crate) struct SparsePolynomial<Scalar: PrimeField> {
112  num_vars: usize,
113  Z: Vec<(usize, Scalar)>,
114}
115
116impl<Scalar: PrimeField> SparsePolynomial<Scalar> {
117  pub fn new(num_vars: usize, Z: Vec<(usize, Scalar)>) -> Self {
118    SparsePolynomial { num_vars, Z }
119  }
120
121  fn compute_chi(a: &[bool], r: &[Scalar]) -> Scalar {
122    assert_eq!(a.len(), r.len());
123    let mut chi_i = Scalar::one();
124    for j in 0..r.len() {
125      if a[j] {
126        chi_i *= r[j];
127      } else {
128        chi_i *= Scalar::one() - r[j];
129      }
130    }
131    chi_i
132  }
133
134  // Takes O(n log n). TODO: do this in O(n) where n is the number of entries in Z
135  pub fn evaluate(&self, r: &[Scalar]) -> Scalar {
136    assert_eq!(self.num_vars, r.len());
137
138    let get_bits = |num: usize, num_bits: usize| -> Vec<bool> {
139      (0..num_bits)
140        .into_par_iter()
141        .map(|shift_amount| ((num & (1 << (num_bits - shift_amount - 1))) > 0))
142        .collect::<Vec<bool>>()
143    };
144
145    (0..self.Z.len())
146      .into_par_iter()
147      .map(|i| {
148        let bits = get_bits(self.Z[i].0, r.len());
149        SparsePolynomial::compute_chi(&bits, r) * self.Z[i].1
150      })
151      .reduce(Scalar::zero, |x, y| x + y)
152  }
153}