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::sync::Arc;
  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<Arc<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| Arc::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<Arc<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| Arc::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<Arc<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(Arc::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<Arc<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(Arc::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    res.par_iter_mut().enumerate().for_each(|(i, x)| {
164        *x = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
165    });
166
167    res
168}
169
170pub fn evaluate_no_par<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
171    assert_eq!(poly.num_vars, point.len());
172    fix_variables_no_par(poly, point).evaluations[0]
173}
174
175fn fix_variables_no_par<F: Field>(
176    poly: &DenseMultilinearExtension<F>,
177    partial_point: &[F],
178) -> DenseMultilinearExtension<F> {
179    assert!(
180        partial_point.len() <= poly.num_vars,
181        "invalid size of partial point"
182    );
183    let nv = poly.num_vars;
184    let mut poly = poly.evaluations.to_vec();
185    let dim = partial_point.len();
186    // evaluate single variable of partial point from left to right
187    for i in 1..dim + 1 {
188        let r = partial_point[i - 1];
189        for b in 0..(1 << (nv - i)) {
190            poly[b] = poly[b << 1] + (poly[(b << 1) + 1] - poly[b << 1]) * r;
191        }
192    }
193    DenseMultilinearExtension::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
194}
195
196/// merge a set of polynomials. Returns an error if the
197/// polynomials do not share a same number of nvs.
198pub fn merge_polynomials<F: PrimeField>(
199    polynomials: &[Arc<DenseMultilinearExtension<F>>],
200) -> Result<Arc<DenseMultilinearExtension<F>>, ArithErrors> {
201    let nv = polynomials[0].num_vars();
202    for poly in polynomials.iter() {
203        if nv != poly.num_vars() {
204            return Err(ArithErrors::InvalidParameters(
205                "num_vars do not match for polynomials".to_string(),
206            ));
207        }
208    }
209
210    let merged_nv = get_batched_nv(nv, polynomials.len());
211    let mut scalars = vec![];
212    for poly in polynomials.iter() {
213        scalars.extend_from_slice(poly.to_evaluations().as_slice());
214    }
215    scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
216    Ok(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
217        merged_nv, scalars,
218    )))
219}
220
221pub fn fix_last_variables_no_par<F: PrimeField>(
222    poly: &DenseMultilinearExtension<F>,
223    partial_point: &[F],
224) -> DenseMultilinearExtension<F> {
225    let mut res = fix_last_variable_no_par(poly, partial_point.last().unwrap());
226    for p in partial_point.iter().rev().skip(1) {
227        res = fix_last_variable_no_par(&res, p);
228    }
229    res
230}
231
232fn fix_last_variable_no_par<F: PrimeField>(
233    poly: &DenseMultilinearExtension<F>,
234    partial_point: &F,
235) -> DenseMultilinearExtension<F> {
236    let nv = poly.num_vars();
237    let half_len = 1 << (nv - 1);
238    let mut res = vec![F::zero(); half_len];
239    for (i, e) in res.iter_mut().enumerate().take(half_len) {
240        *e = poly.evaluations[i]
241            + *partial_point * (poly.evaluations[i + half_len] - poly.evaluations[i]);
242    }
243    DenseMultilinearExtension::from_evaluations_vec(nv - 1, res)
244}
245pub fn fix_last_variables<F: PrimeField>(
246    poly: &DenseMultilinearExtension<F>,
247    partial_point: &[F],
248) -> DenseMultilinearExtension<F> {
249    assert!(
250        partial_point.len() <= poly.num_vars,
251        "invalid size of partial point"
252    );
253    let nv = poly.num_vars;
254    let mut poly = poly.evaluations.to_vec();
255    let dim = partial_point.len();
256    // evaluate single variable of partial point from left to right
257    for (i, point) in partial_point.iter().rev().enumerate().take(dim) {
258        poly = fix_last_variable_helper(&poly, nv - i, point);
259    }
260
261    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
262}
263
264fn fix_last_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
265    let half_len = 1 << (nv - 1);
266    let mut res = vec![F::zero(); half_len];
267
268    // evaluate single variable of partial point from left to right
269    #[cfg(not(feature = "parallel"))]
270    for b in 0..half_len {
271        res[b] = data[b] + (data[b + half_len] - data[b]) * point;
272    }
273
274    #[cfg(feature = "parallel")]
275    res.par_iter_mut().enumerate().for_each(|(i, x)| {
276        *x = data[i] + (data[i + half_len] - data[i]) * point;
277    });
278
279    res
280}