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/// A list of MLEs that represents an identity permutation
 76pub fn identity_permutation_mles<F: PrimeField>(
 77    num_vars: usize,
 78    num_chunks: usize,
 79) -> Vec<Rc<DenseMultilinearExtension<F>>> {
 80    let mut res = vec![];
 81    for i in 0..num_chunks {
 82        let shift = (i * (1 << num_vars)) as u64;
 83        let s_id_vec = (shift..shift + (1u64 << num_vars)).map(F::from).collect();
 84        res.push(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
 85            num_vars, s_id_vec,
 86        )));
 87    }
 88    res
 89}
 90
 91/// A list of MLEs that represent a random permutation
 92pub fn random_permutation_mles<F: PrimeField, R: RngCore>(
 93    num_vars: usize,
 94    num_chunks: usize,
 95    rng: &mut R,
 96) -> Vec<Rc<DenseMultilinearExtension<F>>> {
 97    let len = (num_chunks as u64) * (1u64 << num_vars);
 98    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
 99    let mut s_perm_vec = vec![];
100    for _ in 0..len {
101        let index = rng.next_u64() as usize % s_id_vec.len();
102        s_perm_vec.push(s_id_vec.remove(index));
103    }
104    let mut res = vec![];
105    let n = 1 << num_vars;
106    for i in 0..num_chunks {
107        res.push(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
108            num_vars,
109            s_perm_vec[i * n..i * n + n].to_vec(),
110        )));
111    }
112    res
113}
114
115pub fn evaluate_opt<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
116    assert_eq!(poly.num_vars, point.len());
117    fix_variables(poly, point).evaluations[0]
118}
119
120pub fn fix_variables<F: Field>(
121    poly: &DenseMultilinearExtension<F>,
122    partial_point: &[F],
123) -> DenseMultilinearExtension<F> {
124    assert!(
125        partial_point.len() <= poly.num_vars,
126        "invalid size of partial point"
127    );
128    let nv = poly.num_vars;
129    let mut poly = poly.evaluations.to_vec();
130    let dim = partial_point.len();
131    // evaluate single variable of partial point from left to right
132    for (i, point) in partial_point.iter().enumerate().take(dim) {
133        poly = fix_one_variable_helper(&poly, nv - i, point);
134    }
135
136    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
137}
138
139fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
140    let mut res = vec![F::zero(); 1 << (nv - 1)];
141
142    // evaluate single variable of partial point from left to right
143    #[cfg(not(feature = "parallel"))]
144    for i in 0..(1 << (nv - 1)) {
145        res[i] = data[i] + (data[(i << 1) + 1] - data[i << 1]) * point;
146    }
147
148    #[cfg(feature = "parallel")]
149    if nv >= 13 {
150        // on my computer we parallelization doesn't help till nv >= 13
151        res.par_iter_mut().enumerate().for_each(|(i, x)| {
152            *x = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
153        });
154    } else {
155        for i in 0..(1 << (nv - 1)) {
156            res[i] = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
157        }
158    }
159
160    res
161}
162
163pub fn evaluate_no_par<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
164    assert_eq!(poly.num_vars, point.len());
165    fix_variables_no_par(poly, point).evaluations[0]
166}
167
168fn fix_variables_no_par<F: Field>(
169    poly: &DenseMultilinearExtension<F>,
170    partial_point: &[F],
171) -> DenseMultilinearExtension<F> {
172    assert!(
173        partial_point.len() <= poly.num_vars,
174        "invalid size of partial point"
175    );
176    let nv = poly.num_vars;
177    let mut poly = poly.evaluations.to_vec();
178    let dim = partial_point.len();
179    // evaluate single variable of partial point from left to right
180    for i in 1..dim + 1 {
181        let r = partial_point[i - 1];
182        for b in 0..(1 << (nv - i)) {
183            poly[b] = poly[b << 1] + (poly[(b << 1) + 1] - poly[b << 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}
213
214pub fn fix_last_variables_no_par<F: PrimeField>(
215    poly: &DenseMultilinearExtension<F>,
216    partial_point: &[F],
217) -> DenseMultilinearExtension<F> {
218    let mut res = fix_last_variable_no_par(poly, partial_point.last().unwrap());
219    for p in partial_point.iter().rev().skip(1) {
220        res = fix_last_variable_no_par(&res, p);
221    }
222    res
223}
224
225fn fix_last_variable_no_par<F: PrimeField>(
226    poly: &DenseMultilinearExtension<F>,
227    partial_point: &F,
228) -> DenseMultilinearExtension<F> {
229    let nv = poly.num_vars();
230    let half_len = 1 << (nv - 1);
231    let mut res = vec![F::zero(); half_len];
232    for (i, e) in res.iter_mut().enumerate().take(half_len) {
233        *e = poly.evaluations[i]
234            + *partial_point * (poly.evaluations[i + half_len] - poly.evaluations[i]);
235    }
236    DenseMultilinearExtension::from_evaluations_vec(nv - 1, res)
237}
238pub fn fix_last_variables<F: PrimeField>(
239    poly: &DenseMultilinearExtension<F>,
240    partial_point: &[F],
241) -> DenseMultilinearExtension<F> {
242    assert!(
243        partial_point.len() <= poly.num_vars,
244        "invalid size of partial point"
245    );
246    let nv = poly.num_vars;
247    let mut poly = poly.evaluations.to_vec();
248    let dim = partial_point.len();
249    // evaluate single variable of partial point from left to right
250    for (i, point) in partial_point.iter().rev().enumerate().take(dim) {
251        poly = fix_last_variable_helper(&poly, nv - i, point);
252    }
253
254    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
255}
256
257fn fix_last_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
258    let half_len = 1 << (nv - 1);
259    let mut res = vec![F::zero(); half_len];
260
261    // evaluate single variable of partial point from left to right
262    #[cfg(not(feature = "parallel"))]
263    for b in 0..half_len {
264        res[b] = data[b] + (data[b + half_len] - data[b]) * point;
265    }
266
267    #[cfg(feature = "parallel")]
268    if nv >= 13 {
269        // on my computer we parallelization doesn't help till nv >= 13
270        res.par_iter_mut().enumerate().for_each(|(i, x)| {
271            *x = data[i] + (data[i + half_len] - data[i]) * point;
272        });
273    } else {
274        for b in 0..(1 << (nv - 1)) {
275            res[b] = data[b] + (data[b + half_len] - data[b]) * point;
276        }
277    }
278
279    res
280}