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