1// Copyright (c) 2023 Espresso Systems (espressosys.com)
  2// This file is part of the HyperPlonk library.
  3
  4// You should have received a copy of the MIT License
  5// along with the HyperPlonk library. If not, see <https://mit-license.org/>.
  6
  7use crate::{util::get_batched_nv, ArithErrors};
  8use ark_ff::{Field, PrimeField};
  9use ark_poly::MultilinearExtension;
 10use ark_std::{end_timer, rand::RngCore, start_timer};
 11#[cfg(feature = "parallel")]
 12use rayon::prelude::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
 13use std::sync::Arc;
 14
 15pub use ark_poly::DenseMultilinearExtension;
 16
 17/// Sample a random list of multilinear polynomials.
 18/// Returns
 19/// - the list of polynomials,
 20/// - its sum of polynomial evaluations over the boolean hypercube.
 21pub fn random_mle_list<F: PrimeField, R: RngCore>(
 22    nv: usize,
 23    degree: usize,
 24    rng: &mut R,
 25) -> (Vec<Arc<DenseMultilinearExtension<F>>>, F) {
 26    let start = start_timer!(|| "sample random mle list");
 27    let mut multiplicands = Vec::with_capacity(degree);
 28    for _ in 0..degree {
 29        multiplicands.push(Vec::with_capacity(1 << nv))
 30    }
 31    let mut sum = F::zero();
 32
 33    for _ in 0..(1 << nv) {
 34        let mut product = F::one();
 35
 36        for e in multiplicands.iter_mut() {
 37            let val = F::rand(rng);
 38            e.push(val);
 39            product *= val;
 40        }
 41        sum += product;
 42    }
 43
 44    let list = multiplicands
 45        .into_iter()
 46        .map(|x| Arc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
 47        .collect();
 48
 49    end_timer!(start);
 50    (list, sum)
 51}
 52
 53// Build a randomize list of mle-s whose sum is zero.
 54pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
 55    nv: usize,
 56    degree: usize,
 57    rng: &mut R,
 58) -> Vec<Arc<DenseMultilinearExtension<F>>> {
 59    let start = start_timer!(|| "sample random zero mle list");
 60
 61    let mut multiplicands = Vec::with_capacity(degree);
 62    for _ in 0..degree {
 63        multiplicands.push(Vec::with_capacity(1 << nv))
 64    }
 65    for _ in 0..(1 << nv) {
 66        multiplicands[0].push(F::zero());
 67        for e in multiplicands.iter_mut().skip(1) {
 68            e.push(F::rand(rng));
 69        }
 70    }
 71
 72    let list = multiplicands
 73        .into_iter()
 74        .map(|x| Arc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
 75        .collect();
 76
 77    end_timer!(start);
 78    list
 79}
 80
 81pub fn identity_permutation<F: PrimeField>(num_vars: usize, num_chunks: usize) -> Vec<F> {
 82    let len = (num_chunks as u64) * (1u64 << num_vars);
 83    (0..len).map(F::from).collect()
 84}
 85
 86/// A list of MLEs that represents an identity permutation
 87pub fn identity_permutation_mles<F: PrimeField>(
 88    num_vars: usize,
 89    num_chunks: usize,
 90) -> Vec<Arc<DenseMultilinearExtension<F>>> {
 91    let mut res = vec![];
 92    for i in 0..num_chunks {
 93        let shift = (i * (1 << num_vars)) as u64;
 94        let s_id_vec = (shift..shift + (1u64 << num_vars)).map(F::from).collect();
 95        res.push(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
 96            num_vars, s_id_vec,
 97        )));
 98    }
 99    res
100}
101
102pub fn random_permutation<F: PrimeField, R: RngCore>(
103    num_vars: usize,
104    num_chunks: usize,
105    rng: &mut R,
106) -> Vec<F> {
107    let len = (num_chunks as u64) * (1u64 << num_vars);
108    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
109    let mut s_perm_vec = vec![];
110    for _ in 0..len {
111        let index = rng.next_u64() as usize % s_id_vec.len();
112        s_perm_vec.push(s_id_vec.remove(index));
113    }
114    s_perm_vec
115}
116
117/// A list of MLEs that represent a random permutation
118pub fn random_permutation_mles<F: PrimeField, R: RngCore>(
119    num_vars: usize,
120    num_chunks: usize,
121    rng: &mut R,
122) -> Vec<Arc<DenseMultilinearExtension<F>>> {
123    let s_perm_vec = random_permutation(num_vars, num_chunks, rng);
124    let mut res = vec![];
125    let n = 1 << num_vars;
126    for i in 0..num_chunks {
127        res.push(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
128            num_vars,
129            s_perm_vec[i * n..i * n + n].to_vec(),
130        )));
131    }
132    res
133}
134
135pub fn evaluate_opt<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
136    assert_eq!(poly.num_vars, point.len());
137    fix_variables(poly, point).evaluations[0]
138}
139
140pub fn fix_variables<F: Field>(
141    poly: &DenseMultilinearExtension<F>,
142    partial_point: &[F],
143) -> DenseMultilinearExtension<F> {
144    assert!(
145        partial_point.len() <= poly.num_vars,
146        "invalid size of partial point"
147    );
148    let nv = poly.num_vars;
149    let mut poly = poly.evaluations.to_vec();
150    let dim = partial_point.len();
151    // evaluate single variable of partial point from left to right
152    for (i, point) in partial_point.iter().enumerate().take(dim) {
153        poly = fix_one_variable_helper(&poly, nv - i, point);
154    }
155
156    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
157}
158
159fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
160    let mut res = vec![F::zero(); 1 << (nv - 1)];
161
162    // evaluate single variable of partial point from left to right
163    #[cfg(not(feature = "parallel"))]
164    for i in 0..(1 << (nv - 1)) {
165        res[i] = data[i] + (data[(i << 1) + 1] - data[i << 1]) * point;
166    }
167
168    #[cfg(feature = "parallel")]
169    res.par_iter_mut().enumerate().for_each(|(i, x)| {
170        *x = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
171    });
172
173    res
174}
175
176pub fn evaluate_no_par<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
177    assert_eq!(poly.num_vars, point.len());
178    fix_variables_no_par(poly, point).evaluations[0]
179}
180
181fn fix_variables_no_par<F: Field>(
182    poly: &DenseMultilinearExtension<F>,
183    partial_point: &[F],
184) -> DenseMultilinearExtension<F> {
185    assert!(
186        partial_point.len() <= poly.num_vars,
187        "invalid size of partial point"
188    );
189    let nv = poly.num_vars;
190    let mut poly = poly.evaluations.to_vec();
191    let dim = partial_point.len();
192    // evaluate single variable of partial point from left to right
193    for i in 1..dim + 1 {
194        let r = partial_point[i - 1];
195        for b in 0..(1 << (nv - i)) {
196            poly[b] = poly[b << 1] + (poly[(b << 1) + 1] - poly[b << 1]) * r;
197        }
198    }
199    DenseMultilinearExtension::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
200}
201
202/// merge a set of polynomials. Returns an error if the
203/// polynomials do not share a same number of nvs.
204pub fn merge_polynomials<F: PrimeField>(
205    polynomials: &[Arc<DenseMultilinearExtension<F>>],
206) -> Result<Arc<DenseMultilinearExtension<F>>, ArithErrors> {
207    let nv = polynomials[0].num_vars();
208    for poly in polynomials.iter() {
209        if nv != poly.num_vars() {
210            return Err(ArithErrors::InvalidParameters(
211                "num_vars do not match for polynomials".to_string(),
212            ));
213        }
214    }
215
216    let merged_nv = get_batched_nv(nv, polynomials.len());
217    let mut scalars = vec![];
218    for poly in polynomials.iter() {
219        scalars.extend_from_slice(poly.to_evaluations().as_slice());
220    }
221    scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
222    Ok(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
223        merged_nv, scalars,
224    )))
225}
226
227pub fn fix_last_variables_no_par<F: PrimeField>(
228    poly: &DenseMultilinearExtension<F>,
229    partial_point: &[F],
230) -> DenseMultilinearExtension<F> {
231    let mut res = fix_last_variable_no_par(poly, partial_point.last().unwrap());
232    for p in partial_point.iter().rev().skip(1) {
233        res = fix_last_variable_no_par(&res, p);
234    }
235    res
236}
237
238fn fix_last_variable_no_par<F: PrimeField>(
239    poly: &DenseMultilinearExtension<F>,
240    partial_point: &F,
241) -> DenseMultilinearExtension<F> {
242    let nv = poly.num_vars();
243    let half_len = 1 << (nv - 1);
244    let mut res = vec![F::zero(); half_len];
245    for (i, e) in res.iter_mut().enumerate().take(half_len) {
246        *e = poly.evaluations[i]
247            + *partial_point * (poly.evaluations[i + half_len] - poly.evaluations[i]);
248    }
249    DenseMultilinearExtension::from_evaluations_vec(nv - 1, res)
250}
251pub fn fix_last_variables<F: PrimeField>(
252    poly: &DenseMultilinearExtension<F>,
253    partial_point: &[F],
254) -> DenseMultilinearExtension<F> {
255    assert!(
256        partial_point.len() <= poly.num_vars,
257        "invalid size of partial point"
258    );
259    let nv = poly.num_vars;
260    let mut poly = poly.evaluations.to_vec();
261    let dim = partial_point.len();
262    // evaluate single variable of partial point from left to right
263    for (i, point) in partial_point.iter().rev().enumerate().take(dim) {
264        poly = fix_last_variable_helper(&poly, nv - i, point);
265    }
266
267    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
268}
269
270fn fix_last_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
271    let half_len = 1 << (nv - 1);
272    let mut res = vec![F::zero(); half_len];
273
274    // evaluate single variable of partial point from left to right
275    #[cfg(not(feature = "parallel"))]
276    for b in 0..half_len {
277        res[b] = data[b] + (data[b + half_len] - data[b]) * point;
278    }
279
280    #[cfg(feature = "parallel")]
281    res.par_iter_mut().enumerate().for_each(|(i, x)| {
282        *x = data[i] + (data[i + half_len] - data[i]) * point;
283    });
284
285    res
286}