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