hyperplonk/arithmetic/src/virtual_polynomial.rs
arithmetic/src/virtual_polynomial.rsRaw
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
7//! This module defines our main mathematical object `VirtualPolynomial`; and
8//! various functions associated with it.
9
10use crate::{errors::ArithErrors, multilinear_polynomial::random_zero_mle_list, random_mle_list};
11use ark_ff::PrimeField;
12use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
13use ark_serialize::{CanonicalSerialize, SerializationError, Write};
14use ark_std::{
15 end_timer,
16 rand::{Rng, RngCore},
17 start_timer,
18};
19use rayon::prelude::*;
20use std::{cmp::max, collections::HashMap, marker::PhantomData, ops::Add, sync::Arc};
21
22#[rustfmt::skip]
23/// A virtual polynomial is a sum of products of multilinear polynomials;
24/// where the multilinear polynomials are stored via their multilinear
25/// extensions: `(coefficient, DenseMultilinearExtension)`
26///
27/// * Number of products n = `polynomial.products.len()`,
28/// * Number of multiplicands of ith product m_i =
29/// `polynomial.products[i].1.len()`,
30/// * Coefficient of ith product c_i = `polynomial.products[i].0`
31///
32/// The resulting polynomial is
33///
34/// $$ \sum_{i=0}^{n} c_i \cdot \prod_{j=0}^{m_i} P_{ij} $$
35///
36/// Example:
37/// f = c0 * f0 * f1 * f2 + c1 * f3 * f4
38/// where f0 ... f4 are multilinear polynomials
39///
40/// - flattened_ml_extensions stores the multilinear extension representation of
41/// f0, f1, f2, f3 and f4
42/// - products is
43/// \[
44/// (c0, \[0, 1, 2\]),
45/// (c1, \[3, 4\])
46/// \]
47/// - raw_pointers_lookup_table maps fi to i
48///
49#[derive(Clone, Debug, Default, PartialEq)]
50pub struct VirtualPolynomial<F: PrimeField> {
51 /// Aux information about the multilinear polynomial
52 pub aux_info: VPAuxInfo<F>,
53 /// list of reference to products (as usize) of multilinear extension
54 pub products: Vec<(F, Vec<usize>)>,
55 /// Stores multilinear extensions in which product multiplicand can refer
56 /// to.
57 pub flattened_ml_extensions: Vec<Arc<DenseMultilinearExtension<F>>>,
58 /// Pointers to the above poly extensions
59 raw_pointers_lookup_table: HashMap<*const DenseMultilinearExtension<F>, usize>,
60}
61
62#[derive(Clone, Debug, Default, PartialEq, Eq, CanonicalSerialize)]
63/// Auxiliary information about the multilinear polynomial
64pub struct VPAuxInfo<F: PrimeField> {
65 /// max number of multiplicands in each product
66 pub max_degree: usize,
67 /// number of variables of the polynomial
68 pub num_variables: usize,
69 /// Associated field
70 #[doc(hidden)]
71 pub phantom: PhantomData<F>,
72}
73
74impl<F: PrimeField> Add for &VirtualPolynomial<F> {
75 type Output = VirtualPolynomial<F>;
76 fn add(self, other: &VirtualPolynomial<F>) -> Self::Output {
77 let start = start_timer!(|| "virtual poly add");
78 let mut res = self.clone();
79 for products in other.products.iter() {
80 let cur: Vec<Arc<DenseMultilinearExtension<F>>> = products
81 .1
82 .iter()
83 .map(|&x| other.flattened_ml_extensions[x].clone())
84 .collect();
85
86 res.add_mle_list(cur, products.0)
87 .expect("add product failed");
88 }
89 end_timer!(start);
90 res
91 }
92}
93
94// TODO: convert this into a trait
95impl<F: PrimeField> VirtualPolynomial<F> {
96 /// Creates an empty virtual polynomial with `num_variables`.
97 pub fn new(num_variables: usize) -> Self {
98 VirtualPolynomial {
99 aux_info: VPAuxInfo {
100 max_degree: 0,
101 num_variables,
102 phantom: PhantomData::default(),
103 },
104 products: Vec::new(),
105 flattened_ml_extensions: Vec::new(),
106 raw_pointers_lookup_table: HashMap::new(),
107 }
108 }
109
110 /// Creates an new virtual polynomial from a MLE and its coefficient.
111 pub fn new_from_mle(mle: &Arc<DenseMultilinearExtension<F>>, coefficient: F) -> Self {
112 let mle_ptr: *const DenseMultilinearExtension<F> = Arc::as_ptr(mle);
113 let mut hm = HashMap::new();
114 hm.insert(mle_ptr, 0);
115
116 VirtualPolynomial {
117 aux_info: VPAuxInfo {
118 // The max degree is the max degree of any individual variable
119 max_degree: 1,
120 num_variables: mle.num_vars,
121 phantom: PhantomData::default(),
122 },
123 // here `0` points to the first polynomial of `flattened_ml_extensions`
124 products: vec![(coefficient, vec![0])],
125 flattened_ml_extensions: vec![mle.clone()],
126 raw_pointers_lookup_table: hm,
127 }
128 }
129
130 /// Add a product of list of multilinear extensions to self
131 /// Returns an error if the list is empty, or the MLE has a different
132 /// `num_vars` from self.
133 ///
134 /// The MLEs will be multiplied together, and then multiplied by the scalar
135 /// `coefficient`.
136 pub fn add_mle_list(
137 &mut self,
138 mle_list: impl IntoIterator<Item = Arc<DenseMultilinearExtension<F>>>,
139 coefficient: F,
140 ) -> Result<(), ArithErrors> {
141 let mle_list: Vec<Arc<DenseMultilinearExtension<F>>> = mle_list.into_iter().collect();
142 let mut indexed_product = Vec::with_capacity(mle_list.len());
143
144 if mle_list.is_empty() {
145 return Err(ArithErrors::InvalidParameters(
146 "input mle_list is empty".to_string(),
147 ));
148 }
149
150 self.aux_info.max_degree = max(self.aux_info.max_degree, mle_list.len());
151
152 for mle in mle_list {
153 if mle.num_vars != self.aux_info.num_variables {
154 return Err(ArithErrors::InvalidParameters(format!(
155 "product has a multiplicand with wrong number of variables {} vs {}",
156 mle.num_vars, self.aux_info.num_variables
157 )));
158 }
159
160 let mle_ptr: *const DenseMultilinearExtension<F> = Arc::as_ptr(&mle);
161 if let Some(index) = self.raw_pointers_lookup_table.get(&mle_ptr) {
162 indexed_product.push(*index)
163 } else {
164 let curr_index = self.flattened_ml_extensions.len();
165 self.flattened_ml_extensions.push(mle.clone());
166 self.raw_pointers_lookup_table.insert(mle_ptr, curr_index);
167 indexed_product.push(curr_index);
168 }
169 }
170 self.products.push((coefficient, indexed_product));
171 Ok(())
172 }
173
174 /// Multiple the current VirtualPolynomial by an MLE:
175 /// - add the MLE to the MLE list;
176 /// - multiple each product by MLE and its coefficient.
177 /// Returns an error if the MLE has a different `num_vars` from self.
178 pub fn mul_by_mle(
179 &mut self,
180 mle: Arc<DenseMultilinearExtension<F>>,
181 coefficient: F,
182 ) -> Result<(), ArithErrors> {
183 let start = start_timer!(|| "mul by mle");
184
185 if mle.num_vars != self.aux_info.num_variables {
186 return Err(ArithErrors::InvalidParameters(format!(
187 "product has a multiplicand with wrong number of variables {} vs {}",
188 mle.num_vars, self.aux_info.num_variables
189 )));
190 }
191
192 let mle_ptr: *const DenseMultilinearExtension<F> = Arc::as_ptr(&mle);
193
194 // check if this mle already exists in the virtual polynomial
195 let mle_index = match self.raw_pointers_lookup_table.get(&mle_ptr) {
196 Some(&p) => p,
197 None => {
198 self.raw_pointers_lookup_table
199 .insert(mle_ptr, self.flattened_ml_extensions.len());
200 self.flattened_ml_extensions.push(mle);
201 self.flattened_ml_extensions.len() - 1
202 },
203 };
204
205 for (prod_coef, indices) in self.products.iter_mut() {
206 // - add the MLE to the MLE list;
207 // - multiple each product by MLE and its coefficient.
208 indices.push(mle_index);
209 *prod_coef *= coefficient;
210 }
211
212 // increase the max degree by one as the MLE has degree 1.
213 self.aux_info.max_degree += 1;
214 end_timer!(start);
215 Ok(())
216 }
217
218 /// Evaluate the virtual polynomial at point `point`.
219 /// Returns an error is point.len() does not match `num_variables`.
220 pub fn evaluate(&self, point: &[F]) -> Result<F, ArithErrors> {
221 let start = start_timer!(|| "evaluation");
222
223 if self.aux_info.num_variables != point.len() {
224 return Err(ArithErrors::InvalidParameters(format!(
225 "wrong number of variables {} vs {}",
226 self.aux_info.num_variables,
227 point.len()
228 )));
229 }
230
231 let evals: Vec<F> = self
232 .flattened_ml_extensions
233 .iter()
234 .map(|x| {
235 x.evaluate(point).unwrap() // safe unwrap here since we have
236 // already checked that num_var
237 // matches
238 })
239 .collect();
240
241 let res = self
242 .products
243 .iter()
244 .map(|(c, p)| *c * p.iter().map(|&i| evals[i]).product::<F>())
245 .sum();
246
247 end_timer!(start);
248 Ok(res)
249 }
250
251 /// Sample a random virtual polynomial, return the polynomial and its sum.
252 pub fn rand<R: RngCore>(
253 nv: usize,
254 num_multiplicands_range: (usize, usize),
255 num_products: usize,
256 rng: &mut R,
257 ) -> Result<(Self, F), ArithErrors> {
258 let start = start_timer!(|| "sample random virtual polynomial");
259
260 let mut sum = F::zero();
261 let mut poly = VirtualPolynomial::new(nv);
262 for _ in 0..num_products {
263 let num_multiplicands =
264 rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
265 let (product, product_sum) = random_mle_list(nv, num_multiplicands, rng);
266 let coefficient = F::rand(rng);
267 poly.add_mle_list(product.into_iter(), coefficient)?;
268 sum += product_sum * coefficient;
269 }
270
271 end_timer!(start);
272 Ok((poly, sum))
273 }
274
275 /// Sample a random virtual polynomial that evaluates to zero everywhere
276 /// over the boolean hypercube.
277 pub fn rand_zero<R: RngCore>(
278 nv: usize,
279 num_multiplicands_range: (usize, usize),
280 num_products: usize,
281 rng: &mut R,
282 ) -> Result<Self, ArithErrors> {
283 let mut poly = VirtualPolynomial::new(nv);
284 for _ in 0..num_products {
285 let num_multiplicands =
286 rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
287 let product = random_zero_mle_list(nv, num_multiplicands, rng);
288 let coefficient = F::rand(rng);
289 poly.add_mle_list(product.into_iter(), coefficient)?;
290 }
291
292 Ok(poly)
293 }
294
295 // Input poly f(x) and a random vector r, output
296 // \hat f(x) = \sum_{x_i \in eval_x} f(x_i) eq(x, r)
297 // where
298 // eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
299 //
300 // This function is used in ZeroCheck.
301 pub fn build_f_hat(&self, r: &[F]) -> Result<Self, ArithErrors> {
302 let start = start_timer!(|| "zero check build hat f");
303
304 if self.aux_info.num_variables != r.len() {
305 return Err(ArithErrors::InvalidParameters(format!(
306 "r.len() is different from number of variables: {} vs {}",
307 r.len(),
308 self.aux_info.num_variables
309 )));
310 }
311
312 let eq_x_r = build_eq_x_r(r)?;
313 let mut res = self.clone();
314 res.mul_by_mle(eq_x_r, F::one())?;
315
316 end_timer!(start);
317 Ok(res)
318 }
319
320 /// Print out the evaluation map for testing. Panic if the num_vars > 5.
321 pub fn print_evals(&self) {
322 if self.aux_info.num_variables > 5 {
323 panic!("this function is used for testing only. cannot print more than 5 num_vars")
324 }
325 for i in 0..1 << self.aux_info.num_variables {
326 let point = bit_decompose(i, self.aux_info.num_variables);
327 let point_fr: Vec<F> = point.iter().map(|&x| F::from(x)).collect();
328 println!("{} {}", i, self.evaluate(point_fr.as_ref()).unwrap())
329 }
330 println!()
331 }
332}
333
334/// Evaluate eq polynomial.
335pub fn eq_eval<F: PrimeField>(x: &[F], y: &[F]) -> Result<F, ArithErrors> {
336 if x.len() != y.len() {
337 return Err(ArithErrors::InvalidParameters(
338 "x and y have different length".to_string(),
339 ));
340 }
341 let start = start_timer!(|| "eq_eval");
342 let mut res = F::one();
343 for (&xi, &yi) in x.iter().zip(y.iter()) {
344 let xi_yi = xi * yi;
345 res *= xi_yi + xi_yi - xi - yi + F::one();
346 }
347 end_timer!(start);
348 Ok(res)
349}
350
351/// This function build the eq(x, r) polynomial for any given r.
352///
353/// Evaluate
354/// eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
355/// over r, which is
356/// eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
357pub fn build_eq_x_r<F: PrimeField>(
358 r: &[F],
359) -> Result<Arc<DenseMultilinearExtension<F>>, ArithErrors> {
360 let evals = build_eq_x_r_vec(r)?;
361 let mle = DenseMultilinearExtension::from_evaluations_vec(r.len(), evals);
362
363 Ok(Arc::new(mle))
364}
365/// This function build the eq(x, r) polynomial for any given r, and output the
366/// evaluation of eq(x, r) in its vector form.
367///
368/// Evaluate
369/// eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
370/// over r, which is
371/// eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
372pub fn build_eq_x_r_vec<F: PrimeField>(r: &[F]) -> Result<Vec<F>, ArithErrors> {
373 // we build eq(x,r) from its evaluations
374 // we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
375 // for example, with num_vars = 4, x is a binary vector of 4, then
376 // 0 0 0 0 -> (1-r0) * (1-r1) * (1-r2) * (1-r3)
377 // 1 0 0 0 -> r0 * (1-r1) * (1-r2) * (1-r3)
378 // 0 1 0 0 -> (1-r0) * r1 * (1-r2) * (1-r3)
379 // 1 1 0 0 -> r0 * r1 * (1-r2) * (1-r3)
380 // ....
381 // 1 1 1 1 -> r0 * r1 * r2 * r3
382 // we will need 2^num_var evaluations
383
384 let mut eval = Vec::new();
385 build_eq_x_r_helper(r, &mut eval)?;
386
387 Ok(eval)
388}
389
390/// A helper function to build eq(x, r) recursively.
391/// This function takes `r.len()` steps, and for each step it requires a maximum
392/// `r.len()-1` multiplications.
393fn build_eq_x_r_helper<F: PrimeField>(r: &[F], buf: &mut Vec<F>) -> Result<(), ArithErrors> {
394 if r.is_empty() {
395 return Err(ArithErrors::InvalidParameters("r length is 0".to_string()));
396 } else if r.len() == 1 {
397 // initializing the buffer with [1-r_0, r_0]
398 buf.push(F::one() - r[0]);
399 buf.push(r[0]);
400 } else {
401 build_eq_x_r_helper(&r[1..], buf)?;
402
403 // suppose at the previous step we received [b_1, ..., b_k]
404 // for the current step we will need
405 // if x_0 = 0: (1-r0) * [b_1, ..., b_k]
406 // if x_0 = 1: r0 * [b_1, ..., b_k]
407 // let mut res = vec![];
408 // for &b_i in buf.iter() {
409 // let tmp = r[0] * b_i;
410 // res.push(b_i - tmp);
411 // res.push(tmp);
412 // }
413 // *buf = res;
414
415 let mut res = vec![F::zero(); buf.len() << 1];
416 res.par_iter_mut().enumerate().for_each(|(i, val)| {
417 let bi = buf[i >> 1];
418 let tmp = r[0] * bi;
419 if i & 1 == 0 {
420 *val = bi - tmp;
421 } else {
422 *val = tmp;
423 }
424 });
425 *buf = res;
426 }
427
428 Ok(())
429}
430
431/// Decompose an integer into a binary vector in little endian.
432pub fn bit_decompose(input: u64, num_var: usize) -> Vec<bool> {
433 let mut res = Vec::with_capacity(num_var);
434 let mut i = input;
435 for _ in 0..num_var {
436 res.push(i & 1 == 1);
437 i >>= 1;
438 }
439 res
440}
441
442#[cfg(test)]
443mod test {
444 use super::*;
445 use ark_bls12_381::Fr;
446 use ark_ff::UniformRand;
447 use ark_std::test_rng;
448
449 #[test]
450 fn test_virtual_polynomial_additions() -> Result<(), ArithErrors> {
451 let mut rng = test_rng();
452 for nv in 2..5 {
453 for num_products in 2..5 {
454 let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
455
456 let (a, _a_sum) =
457 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
458 let (b, _b_sum) =
459 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
460 let c = &a + &b;
461
462 assert_eq!(
463 a.evaluate(base.as_ref())? + b.evaluate(base.as_ref())?,
464 c.evaluate(base.as_ref())?
465 );
466 }
467 }
468
469 Ok(())
470 }
471
472 #[test]
473 fn test_virtual_polynomial_mul_by_mle() -> Result<(), ArithErrors> {
474 let mut rng = test_rng();
475 for nv in 2..5 {
476 for num_products in 2..5 {
477 let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
478
479 let (a, _a_sum) =
480 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
481 let (b, _b_sum) = random_mle_list(nv, 1, &mut rng);
482 let b_mle = b[0].clone();
483 let coeff = Fr::rand(&mut rng);
484 let b_vp = VirtualPolynomial::new_from_mle(&b_mle, coeff);
485
486 let mut c = a.clone();
487
488 c.mul_by_mle(b_mle, coeff)?;
489
490 assert_eq!(
491 a.evaluate(base.as_ref())? * b_vp.evaluate(base.as_ref())?,
492 c.evaluate(base.as_ref())?
493 );
494 }
495 }
496
497 Ok(())
498 }
499
500 #[test]
501 fn test_eq_xr() {
502 let mut rng = test_rng();
503 for nv in 4..10 {
504 let r: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
505 let eq_x_r = build_eq_x_r(r.as_ref()).unwrap();
506 let eq_x_r2 = build_eq_x_r_for_test(r.as_ref());
507 assert_eq!(eq_x_r, eq_x_r2);
508 }
509 }
510
511 /// Naive method to build eq(x, r).
512 /// Only used for testing purpose.
513 // Evaluate
514 // eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
515 // over r, which is
516 // eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
517 fn build_eq_x_r_for_test<F: PrimeField>(r: &[F]) -> Arc<DenseMultilinearExtension<F>> {
518 // we build eq(x,r) from its evaluations
519 // we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
520 // for example, with num_vars = 4, x is a binary vector of 4, then
521 // 0 0 0 0 -> (1-r0) * (1-r1) * (1-r2) * (1-r3)
522 // 1 0 0 0 -> r0 * (1-r1) * (1-r2) * (1-r3)
523 // 0 1 0 0 -> (1-r0) * r1 * (1-r2) * (1-r3)
524 // 1 1 0 0 -> r0 * r1 * (1-r2) * (1-r3)
525 // ....
526 // 1 1 1 1 -> r0 * r1 * r2 * r3
527 // we will need 2^num_var evaluations
528
529 // First, we build array for {1 - r_i}
530 let one_minus_r: Vec<F> = r.iter().map(|ri| F::one() - ri).collect();
531
532 let num_var = r.len();
533 let mut eval = vec![];
534
535 for i in 0..1 << num_var {
536 let mut current_eval = F::one();
537 let bit_sequence = bit_decompose(i, num_var);
538
539 for (&bit, (ri, one_minus_ri)) in
540 bit_sequence.iter().zip(r.iter().zip(one_minus_r.iter()))
541 {
542 current_eval *= if bit { *ri } else { *one_minus_ri };
543 }
544 eval.push(current_eval);
545 }
546
547 let mle = DenseMultilinearExtension::from_evaluations_vec(num_var, eval);
548
549 let res = Arc::new(mle);
550 res
551 }
552}