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