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