sonobe/src/utils/espresso/virtual_polynomial.rs
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))
328fn build_eq_x_r<F: PrimeField>(r: &[F]) -> 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/// This function build the eq(x, r) polynomial for any given r, and output the
335/// evaluation of eq(x, r) in its vector form.
336///
337/// Evaluate
338/// eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
339/// over r, which is
340/// eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
341fn build_eq_x_r_vec<F: PrimeField>(r: &[F]) -> Result<Vec<F>, ArithErrors> {
342 // we build eq(x,r) from its evaluations
343 // we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
344 // for example, with num_vars = 4, x is a binary vector of 4, then
345 // 0 0 0 0 -> (1-r0) * (1-r1) * (1-r2) * (1-r3)
346 // 1 0 0 0 -> r0 * (1-r1) * (1-r2) * (1-r3)
347 // 0 1 0 0 -> (1-r0) * r1 * (1-r2) * (1-r3)
348 // 1 1 0 0 -> r0 * r1 * (1-r2) * (1-r3)
349 // ....
350 // 1 1 1 1 -> r0 * r1 * r2 * r3
351 // we will need 2^num_var evaluations
352
353 let mut eval = Vec::new();
354 build_eq_x_r_helper(r, &mut eval)?;
355
356 Ok(eval)
357}
358
359/// A helper function to build eq(x, r) recursively.
360/// This function takes `r.len()` steps, and for each step it requires a maximum
361/// `r.len()-1` multiplications.
362fn build_eq_x_r_helper<F: PrimeField>(r: &[F], buf: &mut Vec<F>) -> Result<(), ArithErrors> {
363 if r.is_empty() {
364 return Err(ArithErrors::InvalidParameters("r length is 0".to_string()));
365 } else if r.len() == 1 {
366 // initializing the buffer with [1-r_0, r_0]
367 buf.push(F::one() - r[0]);
368 buf.push(r[0]);
369 } else {
370 build_eq_x_r_helper(&r[1..], buf)?;
371
372 // suppose at the previous step we received [b_1, ..., b_k]
373 // for the current step we will need
374 // if x_0 = 0: (1-r0) * [b_1, ..., b_k]
375 // if x_0 = 1: r0 * [b_1, ..., b_k]
376 // let mut res = vec![];
377 // for &b_i in buf.iter() {
378 // let tmp = r[0] * b_i;
379 // res.push(b_i - tmp);
380 // res.push(tmp);
381 // }
382 // *buf = res;
383
384 let mut res = vec![F::zero(); buf.len() << 1];
385 res.par_iter_mut().enumerate().for_each(|(i, val)| {
386 let bi = buf[i >> 1];
387 let tmp = r[0] * bi;
388 if i & 1 == 0 {
389 *val = bi - tmp;
390 } else {
391 *val = tmp;
392 }
393 });
394 *buf = res;
395 }
396
397 Ok(())
398}
399
400/// Decompose an integer into a binary vector in little endian.
401pub fn bit_decompose(input: u64, num_var: usize) -> Vec<bool> {
402 let mut res = Vec::with_capacity(num_var);
403 let mut i = input;
404 for _ in 0..num_var {
405 res.push(i & 1 == 1);
406 i >>= 1;
407 }
408 res
409}
410
411#[cfg(test)]
412mod test {
413 use super::*;
414 use crate::utils::multilinear_polynomial::tests::random_mle_list;
415 use ark_bls12_377::Fr;
416 use ark_ff::UniformRand;
417 use ark_std::{
418 rand::{Rng, RngCore},
419 test_rng,
420 };
421
422 impl<F: PrimeField> VirtualPolynomial<F> {
423 /// Sample a random virtual polynomial, return the polynomial and its sum.
424 fn rand<R: RngCore>(
425 nv: usize,
426 num_multiplicands_range: (usize, usize),
427 num_products: usize,
428 rng: &mut R,
429 ) -> Result<(Self, F), ArithErrors> {
430 let start = start_timer!(|| "sample random virtual polynomial");
431
432 let mut sum = F::zero();
433 let mut poly = VirtualPolynomial::new(nv);
434 for _ in 0..num_products {
435 let num_multiplicands =
436 rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
437 let (product, product_sum) = random_mle_list(nv, num_multiplicands, rng);
438 let coefficient = F::rand(rng);
439 poly.add_mle_list(product.into_iter(), coefficient)?;
440 sum += product_sum * coefficient;
441 }
442
443 end_timer!(start);
444 Ok((poly, sum))
445 }
446 }
447
448 #[test]
449 fn test_virtual_polynomial_additions() -> Result<(), ArithErrors> {
450 let mut rng = test_rng();
451 for nv in 2..5 {
452 for num_products in 2..5 {
453 let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
454
455 let (a, _a_sum) =
456 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
457 let (b, _b_sum) =
458 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
459 let c = &a + &b;
460
461 assert_eq!(
462 a.evaluate(base.as_ref())? + b.evaluate(base.as_ref())?,
463 c.evaluate(base.as_ref())?
464 );
465 }
466 }
467
468 Ok(())
469 }
470
471 #[test]
472 fn test_virtual_polynomial_mul_by_mle() -> Result<(), ArithErrors> {
473 let mut rng = test_rng();
474 for nv in 2..5 {
475 for num_products in 2..5 {
476 let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
477
478 let (a, _a_sum) =
479 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
480 let (b, _b_sum) = random_mle_list(nv, 1, &mut rng);
481 let b_mle = b[0].clone();
482 let coeff = Fr::rand(&mut rng);
483 let b_vp = VirtualPolynomial::new_from_mle(&b_mle, coeff);
484
485 let mut c = a.clone();
486
487 c.mul_by_mle(b_mle, coeff)?;
488
489 assert_eq!(
490 a.evaluate(base.as_ref())? * b_vp.evaluate(base.as_ref())?,
491 c.evaluate(base.as_ref())?
492 );
493 }
494 }
495
496 Ok(())
497 }
498
499 #[test]
500 fn test_eq_xr() {
501 let mut rng = test_rng();
502 for nv in 4..10 {
503 let r: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
504 let eq_x_r = build_eq_x_r(r.as_ref()).unwrap();
505 let eq_x_r2 = build_eq_x_r_for_test(r.as_ref());
506 assert_eq!(eq_x_r, eq_x_r2);
507 }
508 }
509
510 /// Naive method to build eq(x, r).
511 /// Only used for testing purpose.
512 // Evaluate
513 // eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
514 // over r, which is
515 // eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
516 fn build_eq_x_r_for_test<F: PrimeField>(r: &[F]) -> Arc<DenseMultilinearExtension<F>> {
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 Arc::new(mle)
549 }
550}