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