hyperplonk/poly-iop/src/virtual_poly.rs
poly-iop/src/virtual_poly.rsRaw
1use crate::{errors::PolyIOPErrors, structs::DomainInfo};
2use ark_ff::PrimeField;
3use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
4use ark_std::{
5 end_timer,
6 rand::{Rng, RngCore},
7 start_timer,
8};
9use std::{cmp::max, collections::HashMap, marker::PhantomData, ops::Add, rc::Rc};
10
11#[rustfmt::skip]
12/// A virtual polynomial is a sum of products of multilinear polynomials;
13/// where the multilinear polynomials are stored via their multilinear
14/// extensions: `(coefficient, DenseMultilinearExtension)`
15///
16/// * Number of products n = `polynomial.products.len()`,
17/// * Number of multiplicands of ith product m_i =
18/// `polynomial.products[i].1.len()`,
19/// * Coefficient of ith product c_i = `polynomial.products[i].0`
20///
21/// The resulting polynomial is
22///
23/// $$\sum_{i=0}^{n} c_i \cdot \prod_{j=0}^{m_i} P_{ij} $$
24///
25/// Example:
26/// f = c0 * f0 * f1 * f2 + c1 * f3 * f4
27/// where f0 ... f4 are multilinear polynomials
28///
29/// - flattened_ml_extensions stores the multilinear extension representation of
30/// f0, f1, f2, f3 and f4
31/// - products is
32/// \[
33/// (c0, \[0, 1, 2\]),
34/// (c1, \[3, 4\])
35/// \]
36/// - raw_pointers_lookup_table maps fi to i
37///
38#[derive(Clone, Debug, Default, PartialEq)]
39pub struct VirtualPolynomial<F: PrimeField> {
40 /// Aux information about the multilinear polynomial
41 pub domain_info: DomainInfo<F>,
42 /// list of reference to products (as usize) of multilinear extension
43 pub products: Vec<(F, Vec<usize>)>,
44 /// Stores multilinear extensions in which product multiplicand can refer
45 /// to.
46 pub flattened_ml_extensions: Vec<Rc<DenseMultilinearExtension<F>>>,
47 /// Pointers to the above poly extensions
48 raw_pointers_lookup_table: HashMap<*const DenseMultilinearExtension<F>, usize>,
49}
50
51impl<F: PrimeField> Add for &VirtualPolynomial<F> {
52 type Output = VirtualPolynomial<F>;
53 fn add(self, other: &VirtualPolynomial<F>) -> Self::Output {
54 let start = start_timer!(|| "virtual poly add");
55 let mut res = self.clone();
56 for products in other.products.iter() {
57 let cur: Vec<Rc<DenseMultilinearExtension<F>>> = products
58 .1
59 .iter()
60 .map(|&x| other.flattened_ml_extensions[x].clone())
61 .collect();
62
63 res.add_mle_list(cur, products.0)
64 .expect("add product failed");
65 }
66 end_timer!(start);
67 res
68 }
69}
70
71impl<F: PrimeField> VirtualPolynomial<F> {
72 /// Returns an empty polynomial
73 pub fn new(num_variables: usize) -> Self {
74 VirtualPolynomial {
75 domain_info: DomainInfo {
76 max_degree: 0,
77 num_variables,
78 phantom: PhantomData::default(),
79 },
80 products: Vec::new(),
81 flattened_ml_extensions: Vec::new(),
82 raw_pointers_lookup_table: HashMap::new(),
83 }
84 }
85
86 /// Returns an new virtual polynomial from a MLE
87 pub fn new_from_mle(mle: Rc<DenseMultilinearExtension<F>>, coefficient: F) -> Self {
88 let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
89 let mut hm = HashMap::new();
90 hm.insert(mle_ptr, 0);
91
92 VirtualPolynomial {
93 domain_info: DomainInfo {
94 // The max degree is the max degree of any individual variable
95 max_degree: 1,
96 num_variables: mle.num_vars,
97 phantom: PhantomData::default(),
98 },
99 products: vec![(coefficient, vec![0])],
100 flattened_ml_extensions: vec![mle],
101 raw_pointers_lookup_table: hm,
102 }
103 }
104
105 /// Add a list of multilinear extensions that is meant to be multiplied
106 /// together. The resulting polynomial will be multiplied by the scalar
107 /// `coefficient`.
108 pub fn add_mle_list(
109 &mut self,
110 mle_list: impl IntoIterator<Item = Rc<DenseMultilinearExtension<F>>>,
111 coefficient: F,
112 ) -> Result<(), PolyIOPErrors> {
113 let mle_list: Vec<Rc<DenseMultilinearExtension<F>>> = mle_list.into_iter().collect();
114 let mut indexed_product = Vec::with_capacity(mle_list.len());
115 assert!(!mle_list.is_empty());
116 self.domain_info.max_degree = max(self.domain_info.max_degree, mle_list.len());
117 for mle in mle_list {
118 if mle.num_vars != self.domain_info.num_variables {
119 return Err(PolyIOPErrors::InvalidParameters(format!(
120 "product has a multiplicand with wrong number of variables {} vs {}",
121 mle.num_vars, self.domain_info.num_variables
122 )));
123 }
124
125 let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
126 if let Some(index) = self.raw_pointers_lookup_table.get(&mle_ptr) {
127 indexed_product.push(*index)
128 } else {
129 let curr_index = self.flattened_ml_extensions.len();
130 self.flattened_ml_extensions.push(mle.clone());
131 self.raw_pointers_lookup_table.insert(mle_ptr, curr_index);
132 indexed_product.push(curr_index);
133 }
134 }
135 self.products.push((coefficient, indexed_product));
136 Ok(())
137 }
138
139 /// Multiple the current VirtualPolynomial by an MLE:
140 /// - add the MLE to the MLE list
141 /// - multiple each product by MLE and its coefficient
142 pub fn mul_by_mle(
143 &mut self,
144 mle: Rc<DenseMultilinearExtension<F>>,
145 coefficient: F,
146 ) -> Result<(), PolyIOPErrors> {
147 let start = start_timer!(|| "mul by mle");
148 let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
149
150 let mle_index = match self.raw_pointers_lookup_table.get(&mle_ptr) {
151 Some(&p) => p,
152 None => {
153 self.raw_pointers_lookup_table
154 .insert(mle_ptr, self.flattened_ml_extensions.len());
155 self.flattened_ml_extensions.push(mle);
156 self.flattened_ml_extensions.len() - 1
157 },
158 };
159
160 for (prod_coef, indices) in self.products.iter_mut() {
161 indices.push(mle_index);
162 *prod_coef *= coefficient;
163 }
164 self.domain_info.max_degree += 1;
165 end_timer!(start);
166 Ok(())
167 }
168
169 /// Evaluate the polynomial at point `point`
170 pub fn evaluate(&self, point: &[F]) -> Result<F, PolyIOPErrors> {
171 let start = start_timer!(|| "evaluation");
172
173 if self.domain_info.num_variables != point.len() {
174 return Err(PolyIOPErrors::InvalidParameters(format!(
175 "wrong number of variables {} vs {}",
176 self.domain_info.num_variables,
177 point.len()
178 )));
179 }
180
181 let evals: Vec<F> = self
182 .flattened_ml_extensions
183 .iter()
184 .map(|x| {
185 x.evaluate(point).unwrap() // safe unwrap here since we have
186 // already checked that num_var
187 // matches
188 })
189 .collect();
190
191 let res = self
192 .products
193 .iter()
194 .map(|(c, p)| *c * p.iter().map(|&i| evals[i]).product::<F>())
195 .sum();
196
197 end_timer!(start);
198 Ok(res)
199 }
200
201 /// Sample a random virtual polynomial, return the polynomial and its sum.
202 pub fn rand<R: RngCore>(
203 nv: usize,
204 num_multiplicands_range: (usize, usize),
205 num_products: usize,
206 rng: &mut R,
207 ) -> Result<(Self, F), PolyIOPErrors> {
208 let start = start_timer!(|| "sample random virtual polynomial");
209
210 let mut sum = F::zero();
211 let mut poly = VirtualPolynomial::new(nv);
212 for _ in 0..num_products {
213 let num_multiplicands =
214 rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
215 let (product, product_sum) = random_mle_list(nv, num_multiplicands, rng);
216 let coefficient = F::rand(rng);
217 poly.add_mle_list(product.into_iter(), coefficient)?;
218 sum += product_sum * coefficient;
219 }
220
221 end_timer!(start);
222 Ok((poly, sum))
223 }
224
225 /// Sample a random virtual polynomial that evaluates to zero everywhere on
226 /// the boolean hypercube.
227 pub fn rand_zero<R: RngCore>(
228 nv: usize,
229 num_multiplicands_range: (usize, usize),
230 num_products: usize,
231 rng: &mut R,
232 ) -> Result<Self, PolyIOPErrors> {
233 let mut poly = VirtualPolynomial::new(nv);
234 for _ in 0..num_products {
235 let num_multiplicands =
236 rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
237 let product = random_zero_mle_list(nv, num_multiplicands, rng);
238 let coefficient = F::rand(rng);
239 poly.add_mle_list(product.into_iter(), coefficient).unwrap();
240 }
241
242 Ok(poly)
243 }
244}
245
246/// Sample a random list of multilinear polynomials.
247/// Returns
248/// - the list of polynomials,
249/// - its sum of polynomial evaluations over the boolean hypercube.
250fn random_mle_list<F: PrimeField, R: RngCore>(
251 nv: usize,
252 degree: usize,
253 rng: &mut R,
254) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
255 let start = start_timer!(|| "sample random mle list");
256 let mut multiplicands = Vec::with_capacity(degree);
257 for _ in 0..degree {
258 multiplicands.push(Vec::with_capacity(1 << nv))
259 }
260 let mut sum = F::zero();
261
262 for _ in 0..(1 << nv) {
263 let mut product = F::one();
264
265 for e in multiplicands.iter_mut() {
266 let val = F::rand(rng);
267 e.push(val);
268 product *= val;
269 }
270 sum += product;
271 }
272
273 let list = multiplicands
274 .into_iter()
275 .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
276 .collect();
277
278 end_timer!(start);
279 (list, sum)
280}
281
282// Build a randomize list of mle-s whose sum is zero.
283pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
284 nv: usize,
285 degree: usize,
286 rng: &mut R,
287) -> Vec<Rc<DenseMultilinearExtension<F>>> {
288 let start = start_timer!(|| "sample random zero mle list");
289
290 let mut multiplicands = Vec::with_capacity(degree);
291 for _ in 0..degree {
292 multiplicands.push(Vec::with_capacity(1 << nv))
293 }
294 for _ in 0..(1 << nv) {
295 multiplicands[0].push(F::zero());
296 for e in multiplicands.iter_mut().skip(1) {
297 e.push(F::rand(rng));
298 }
299 }
300
301 let list = multiplicands
302 .into_iter()
303 .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
304 .collect();
305
306 end_timer!(start);
307 list
308}
309
310#[cfg(test)]
311pub(crate) mod test {
312 use super::*;
313 use ark_bls12_381::Fr;
314 use ark_ff::UniformRand;
315 use ark_std::test_rng;
316
317 #[test]
318 fn test_virtual_polynomial_additions() -> Result<(), PolyIOPErrors> {
319 let mut rng = test_rng();
320 for nv in 2..5 {
321 for num_products in 2..5 {
322 let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
323
324 let (a, _a_sum) =
325 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
326 let (b, _b_sum) =
327 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
328 let c = &a + &b;
329
330 assert_eq!(
331 a.evaluate(base.as_ref())? + b.evaluate(base.as_ref())?,
332 c.evaluate(base.as_ref())?
333 );
334 }
335 }
336
337 Ok(())
338 }
339
340 #[test]
341 fn test_virtual_polynomial_mul_by_mle() -> Result<(), PolyIOPErrors> {
342 let mut rng = test_rng();
343 for nv in 2..5 {
344 for num_products in 2..5 {
345 let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
346
347 let (a, _a_sum) =
348 VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
349 let (b, _b_sum) = random_mle_list(nv, 1, &mut rng);
350 let b_mle = b[0].clone();
351 let coeff = Fr::rand(&mut rng);
352 let b_vp = VirtualPolynomial::new_from_mle(b_mle.clone(), coeff);
353
354 let mut c = a.clone();
355
356 c.mul_by_mle(b_mle, coeff)?;
357
358 assert_eq!(
359 a.evaluate(base.as_ref())? * b_vp.evaluate(base.as_ref())?,
360 c.evaluate(base.as_ref())?
361 );
362 }
363 }
364
365 Ok(())
366 }
367}