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