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