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