1use crate::{util::get_batched_nv, ArithErrors};
  2use ark_ff::{Field, PrimeField};
  3use ark_poly::MultilinearExtension;
  4use ark_std::{end_timer, rand::RngCore, start_timer};
  5#[cfg(feature = "parallel")]
  6use rayon::prelude::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
  7use std::rc::Rc;
  8
  9pub use ark_poly::DenseMultilinearExtension;
 10
 11/// Sample a random list of multilinear polynomials.
 12/// Returns
 13/// - the list of polynomials,
 14/// - its sum of polynomial evaluations over the boolean hypercube.
 15pub fn random_mle_list<F: PrimeField, R: RngCore>(
 16    nv: usize,
 17    degree: usize,
 18    rng: &mut R,
 19) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
 20    let start = start_timer!(|| "sample random mle list");
 21    let mut multiplicands = Vec::with_capacity(degree);
 22    for _ in 0..degree {
 23        multiplicands.push(Vec::with_capacity(1 << nv))
 24    }
 25    let mut sum = F::zero();
 26
 27    for _ in 0..(1 << nv) {
 28        let mut product = F::one();
 29
 30        for e in multiplicands.iter_mut() {
 31            let val = F::rand(rng);
 32            e.push(val);
 33            product *= val;
 34        }
 35        sum += product;
 36    }
 37
 38    let list = multiplicands
 39        .into_iter()
 40        .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
 41        .collect();
 42
 43    end_timer!(start);
 44    (list, sum)
 45}
 46
 47// Build a randomize list of mle-s whose sum is zero.
 48pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
 49    nv: usize,
 50    degree: usize,
 51    rng: &mut R,
 52) -> Vec<Rc<DenseMultilinearExtension<F>>> {
 53    let start = start_timer!(|| "sample random zero mle list");
 54
 55    let mut multiplicands = Vec::with_capacity(degree);
 56    for _ in 0..degree {
 57        multiplicands.push(Vec::with_capacity(1 << nv))
 58    }
 59    for _ in 0..(1 << nv) {
 60        multiplicands[0].push(F::zero());
 61        for e in multiplicands.iter_mut().skip(1) {
 62            e.push(F::rand(rng));
 63        }
 64    }
 65
 66    let list = multiplicands
 67        .into_iter()
 68        .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
 69        .collect();
 70
 71    end_timer!(start);
 72    list
 73}
 74
 75/// An MLE that represent an identity permutation: `f(index) \mapto index`
 76pub fn identity_permutation_mle<F: PrimeField>(
 77    num_vars: usize,
 78) -> Rc<DenseMultilinearExtension<F>> {
 79    let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
 80    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
 81        num_vars, s_id_vec,
 82    ))
 83}
 84
 85/// An MLE that represent a random permutation
 86pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
 87    num_vars: usize,
 88    rng: &mut R,
 89) -> Rc<DenseMultilinearExtension<F>> {
 90    let len = 1u64 << num_vars;
 91    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
 92    let mut s_perm_vec = vec![];
 93    for _ in 0..len {
 94        let index = rng.next_u64() as usize % s_id_vec.len();
 95        s_perm_vec.push(s_id_vec.remove(index));
 96    }
 97    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
 98        num_vars, s_perm_vec,
 99    ))
100}
101
102pub fn evaluate_opt<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
103    assert_eq!(poly.num_vars, point.len());
104    fix_variables(poly, point).evaluations[0]
105}
106
107pub fn fix_variables<F: Field>(
108    poly: &DenseMultilinearExtension<F>,
109    partial_point: &[F],
110) -> DenseMultilinearExtension<F> {
111    assert!(
112        partial_point.len() <= poly.num_vars,
113        "invalid size of partial point"
114    );
115    let nv = poly.num_vars;
116    let mut poly = poly.evaluations.to_vec();
117    let dim = partial_point.len();
118    // evaluate single variable of partial point from left to right
119    for (i, point) in partial_point.iter().enumerate().take(dim) {
120        poly = fix_one_variable_helper(&poly, nv - i, point);
121    }
122
123    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
124}
125
126pub fn fix_first_variable<F: Field>(
127    poly: &DenseMultilinearExtension<F>,
128    partial_point: &F,
129) -> DenseMultilinearExtension<F> {
130    assert!(poly.num_vars != 0, "invalid size of partial point");
131
132    let nv = poly.num_vars;
133    let res = fix_one_variable_helper(&poly.evaluations, nv, partial_point);
134    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - 1, &res)
135}
136
137fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
138    let mut res = vec![F::zero(); 1 << (nv - 1)];
139    let one_minus_p = F::one() - point;
140
141    // evaluate single variable of partial point from left to right
142    #[cfg(not(feature = "parallel"))]
143    for b in 0..(1 << (nv - 1)) {
144        res[b] = data[b << 1] * one_minus_p + data[(b << 1) + 1] * point;
145    }
146
147    #[cfg(feature = "parallel")]
148    if nv >= 13 {
149        // on my computer we parallelization doesn't help till nv >= 13
150        res.par_iter_mut().enumerate().for_each(|(i, x)| {
151            *x = data[i << 1] * one_minus_p + data[(i << 1) + 1] * point;
152        });
153    } else {
154        for b in 0..(1 << (nv - 1)) {
155            res[b] = data[b << 1] * one_minus_p + data[(b << 1) + 1] * point;
156        }
157    }
158
159    res
160}
161
162pub fn evaluate_no_par<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
163    assert_eq!(poly.num_vars, point.len());
164    fix_variables_no_par(poly, point).evaluations[0]
165}
166
167fn fix_variables_no_par<F: Field>(
168    poly: &DenseMultilinearExtension<F>,
169    partial_point: &[F],
170) -> DenseMultilinearExtension<F> {
171    assert!(
172        partial_point.len() <= poly.num_vars,
173        "invalid size of partial point"
174    );
175    let nv = poly.num_vars;
176    let mut poly = poly.evaluations.to_vec();
177    let dim = partial_point.len();
178    // evaluate single variable of partial point from left to right
179    for i in 1..dim + 1 {
180        let r = partial_point[i - 1];
181        let one_minus_r = F::one() - r;
182        for b in 0..(1 << (nv - i)) {
183            poly[b] = poly[b << 1] * one_minus_r + poly[(b << 1) + 1] * r;
184        }
185    }
186    DenseMultilinearExtension::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
187}
188
189/// merge a set of polynomials. Returns an error if the
190/// polynomials do not share a same number of nvs.
191pub fn merge_polynomials<F: PrimeField>(
192    polynomials: &[Rc<DenseMultilinearExtension<F>>],
193) -> Result<Rc<DenseMultilinearExtension<F>>, ArithErrors> {
194    let nv = polynomials[0].num_vars();
195    for poly in polynomials.iter() {
196        if nv != poly.num_vars() {
197            return Err(ArithErrors::InvalidParameters(
198                "num_vars do not match for polynomials".to_string(),
199            ));
200        }
201    }
202
203    let merged_nv = get_batched_nv(nv, polynomials.len());
204    let mut scalars = vec![];
205    for poly in polynomials.iter() {
206        scalars.extend_from_slice(poly.to_evaluations().as_slice());
207    }
208    scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
209    Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
210        merged_nv, scalars,
211    )))
212}