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
126fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
127    let mut res = vec![F::zero(); 1 << (nv - 1)];
128
129    // evaluate single variable of partial point from left to right
130    #[cfg(not(feature = "parallel"))]
131    for i in 0..(1 << (nv - 1)) {
132        res[i] = data[i] + (data[(i << 1) + 1] - data[i << 1]) * point;
133    }
134
135    #[cfg(feature = "parallel")]
136    if nv >= 13 {
137        // on my computer we parallelization doesn't help till nv >= 13
138        res.par_iter_mut().enumerate().for_each(|(i, x)| {
139            *x = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
140        });
141    } else {
142        for i in 0..(1 << (nv - 1)) {
143            res[i] = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
144        }
145    }
146
147    res
148}
149
150pub fn evaluate_no_par<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
151    assert_eq!(poly.num_vars, point.len());
152    fix_variables_no_par(poly, point).evaluations[0]
153}
154
155fn fix_variables_no_par<F: Field>(
156    poly: &DenseMultilinearExtension<F>,
157    partial_point: &[F],
158) -> DenseMultilinearExtension<F> {
159    assert!(
160        partial_point.len() <= poly.num_vars,
161        "invalid size of partial point"
162    );
163    let nv = poly.num_vars;
164    let mut poly = poly.evaluations.to_vec();
165    let dim = partial_point.len();
166    // evaluate single variable of partial point from left to right
167    for i in 1..dim + 1 {
168        let r = partial_point[i - 1];
169        for b in 0..(1 << (nv - i)) {
170            poly[b] = poly[b << 1] + (poly[(b << 1) + 1] - poly[b << 1]) * r;
171        }
172    }
173    DenseMultilinearExtension::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
174}
175
176/// merge a set of polynomials. Returns an error if the
177/// polynomials do not share a same number of nvs.
178pub fn merge_polynomials<F: PrimeField>(
179    polynomials: &[Rc<DenseMultilinearExtension<F>>],
180) -> Result<Rc<DenseMultilinearExtension<F>>, ArithErrors> {
181    let nv = polynomials[0].num_vars();
182    for poly in polynomials.iter() {
183        if nv != poly.num_vars() {
184            return Err(ArithErrors::InvalidParameters(
185                "num_vars do not match for polynomials".to_string(),
186            ));
187        }
188    }
189
190    let merged_nv = get_batched_nv(nv, polynomials.len());
191    let mut scalars = vec![];
192    for poly in polynomials.iter() {
193        scalars.extend_from_slice(poly.to_evaluations().as_slice());
194    }
195    scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
196    Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
197        merged_nv, scalars,
198    )))
199}
200
201pub fn fix_last_variables_no_par<F: PrimeField>(
202    poly: &DenseMultilinearExtension<F>,
203    partial_point: &[F],
204) -> DenseMultilinearExtension<F> {
205    let mut res = fix_last_variable_no_par(poly, partial_point.last().unwrap());
206    for p in partial_point.iter().rev().skip(1) {
207        res = fix_last_variable_no_par(&res, p);
208    }
209    res
210}
211
212fn fix_last_variable_no_par<F: PrimeField>(
213    poly: &DenseMultilinearExtension<F>,
214    partial_point: &F,
215) -> DenseMultilinearExtension<F> {
216    let nv = poly.num_vars();
217    let half_len = 1 << (nv - 1);
218    let mut res = vec![F::zero(); half_len];
219    for (i, e) in res.iter_mut().enumerate().take(half_len) {
220        *e = poly.evaluations[i]
221            + *partial_point * (poly.evaluations[i + half_len] - poly.evaluations[i]);
222    }
223    DenseMultilinearExtension::from_evaluations_vec(nv - 1, res)
224}
225pub fn fix_last_variables<F: PrimeField>(
226    poly: &DenseMultilinearExtension<F>,
227    partial_point: &[F],
228) -> DenseMultilinearExtension<F> {
229    assert!(
230        partial_point.len() <= poly.num_vars,
231        "invalid size of partial point"
232    );
233    let nv = poly.num_vars;
234    let mut poly = poly.evaluations.to_vec();
235    let dim = partial_point.len();
236    // evaluate single variable of partial point from left to right
237    for (i, point) in partial_point.iter().rev().enumerate().take(dim) {
238        poly = fix_last_variable_helper(&poly, nv - i, point);
239    }
240
241    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
242}
243
244fn fix_last_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
245    let half_len = 1 << (nv - 1);
246    let mut res = vec![F::zero(); half_len];
247
248    // evaluate single variable of partial point from left to right
249    #[cfg(not(feature = "parallel"))]
250    for b in 0..half_len {
251        res[b] = data[b] + (data[b + half_len] - data[b]) * point;
252    }
253
254    #[cfg(feature = "parallel")]
255    if nv >= 13 {
256        // on my computer we parallelization doesn't help till nv >= 13
257        res.par_iter_mut().enumerate().for_each(|(i, x)| {
258            *x = data[i] + (data[i + half_len] - data[i]) * point;
259        });
260    } else {
261        for b in 0..(1 << (nv - 1)) {
262            res[b] = data[b] + (data[b + half_len] - data[b]) * point;
263        }
264    }
265
266    res
267}