ark-r1cs-std/src/poly/evaluations/univariate/mod.rs
src/poly/evaluations/univariate/mod.rsRaw
1pub mod lagrange_interpolator;
2
3use crate::{
4 alloc::AllocVar,
5 eq::EqGadget,
6 fields::{fp::FpVar, FieldVar},
7 poly::{
8 domain::Radix2DomainVar,
9 evaluations::univariate::lagrange_interpolator::LagrangeInterpolator,
10 },
11 R1CSVar,
12};
13use ark_ff::{batch_inversion, PrimeField};
14use ark_relations::r1cs::SynthesisError;
15use ark_std::{
16 ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign},
17 vec::Vec,
18};
19
20#[derive(Clone)]
21/// Stores a UV polynomial in evaluation form.
22pub struct EvaluationsVar<F: PrimeField> {
23 /// Evaluations of univariate polynomial over domain
24 pub evals: Vec<FpVar<F>>,
25 /// Optional Lagrange Interpolator. Useful for lagrange interpolation.
26 pub lagrange_interpolator: Option<LagrangeInterpolator<F>>,
27 domain: Radix2DomainVar<F>,
28 /// Contains all domain elements of `domain.base_domain`.
29 ///
30 /// This is a cache for lagrange interpolation when offset is non-constant.
31 /// Will be `None` if offset is constant or `interpolate` is set to
32 /// `false`.
33 subgroup_points: Option<Vec<F>>,
34}
35
36impl<F: PrimeField> EvaluationsVar<F> {
37 /// Construct `Self` from evaluations and a domain.
38 /// `interpolate` indicates if user wants to interpolate this polynomial
39 /// using lagrange interpolation.
40 pub fn from_vec_and_domain(
41 evaluations: Vec<FpVar<F>>,
42 domain: Radix2DomainVar<F>,
43 interpolate: bool,
44 ) -> Self {
45 assert_eq!(
46 evaluations.len(),
47 1 << domain.dim,
48 "evaluations and domain has different dimensions"
49 );
50
51 let mut ev = Self {
52 evals: evaluations,
53 lagrange_interpolator: None,
54 domain,
55 subgroup_points: None,
56 };
57 if interpolate {
58 ev.generate_interpolation_cache();
59 }
60 ev
61 }
62
63 /// Precompute necessary calculation for lagrange interpolation and mark it
64 /// ready to interpolate
65 pub fn generate_interpolation_cache(&mut self) {
66 if self.domain.offset().is_constant() {
67 let poly_evaluations_val: Vec<_> =
68 self.evals.iter().map(|v| v.value().unwrap()).collect();
69 let domain = &self.domain;
70 let lagrange_interpolator = if let &FpVar::Constant(x) = domain.offset() {
71 LagrangeInterpolator::new(x, domain.gen, domain.dim, poly_evaluations_val)
72 } else {
73 panic!("Domain offset needs to be constant.")
74 };
75 self.lagrange_interpolator = Some(lagrange_interpolator)
76 } else {
77 // calculate all elements of base subgroup so that in later part we don't need
78 // to calculate the exponents again
79 let mut subgroup_points = Vec::with_capacity(self.domain.size() as usize);
80 subgroup_points.push(F::one());
81 for i in 1..self.domain.size() as usize {
82 subgroup_points.push(subgroup_points[i - 1] * self.domain.gen)
83 }
84 self.subgroup_points = Some(subgroup_points)
85 }
86 }
87
88 /// Compute lagrange coefficients for each evaluation, given
89 /// `interpolation_point`. Only valid if the domain offset is constant.
90 fn compute_lagrange_coefficients(
91 &self,
92 interpolation_point: &FpVar<F>,
93 ) -> Result<Vec<FpVar<F>>, SynthesisError> {
94 // ref: https://github.com/alexchmit/perfect-constraints/blob/79692f2652a95a57f2c7187f5b5276345e680230/fractal/src/algebra/lagrange_interpolation.rs#L159
95 let cs = interpolation_point.cs();
96 let t = interpolation_point;
97 let lagrange_interpolator = self
98 .lagrange_interpolator
99 .as_ref()
100 .expect("lagrange interpolator has not been initialized. \
101 Call `self.generate_interpolation_cache` first or set `interpolate` to true in constructor. ");
102 let lagrange_coeffs =
103 lagrange_interpolator.compute_lagrange_coefficients(t.value().unwrap());
104 let mut lagrange_coeffs_fg = Vec::new();
105 // Now we convert these lagrange coefficients to gadgets, and then constrain
106 // them. The i-th lagrange coefficients constraint is:
107 // (v_inv[i] * t - v_inv[i] * domain_elem[i]) * (coeff) = 1/Z_I(t)
108 let vp_t = lagrange_interpolator.domain_vp.evaluate_constraints(t)?;
109 // let inv_vp_t = vp_t.inverse()?;
110 for i in 0..lagrange_interpolator.domain_order {
111 let constant: F =
112 (-lagrange_interpolator.all_domain_elems[i]) * lagrange_interpolator.v_inv_elems[i];
113 let mut a_element: FpVar<F> =
114 t * &FpVar::constant(lagrange_interpolator.v_inv_elems[i]);
115 a_element += FpVar::constant(constant);
116
117 let lag_coeff: FpVar<F> =
118 FpVar::new_witness(ns!(cs, "generate lagrange coefficient"), || {
119 Ok(lagrange_coeffs[i])
120 })?;
121 // Enforce the actual constraint (A_element) * (lagrange_coeff) = 1/Z_I(t)
122 assert_eq!(
123 (lagrange_interpolator.v_inv_elems[i] * t.value().unwrap()
124 - lagrange_interpolator.v_inv_elems[i]
125 * lagrange_interpolator.all_domain_elems[i])
126 * lagrange_coeffs[i],
127 vp_t.value().unwrap()
128 );
129 a_element.mul_equals(&lag_coeff, &vp_t)?;
130 lagrange_coeffs_fg.push(lag_coeff);
131 }
132 Ok(lagrange_coeffs_fg)
133 }
134
135 /// Returns constraints for Interpolating and evaluating at
136 /// `interpolation_point`
137 pub fn interpolate_and_evaluate(
138 &self,
139 interpolation_point: &FpVar<F>,
140 ) -> Result<FpVar<F>, SynthesisError> {
141 // specialize: if domain offset is constant, we can optimize to have fewer
142 // constraints
143 if self.domain.offset().is_constant() {
144 self.lagrange_interpolate_with_constant_offset(interpolation_point)
145 } else {
146 // if domain offset is not constant, then we use standard lagrange interpolation
147 // code
148 self.lagrange_interpolate_with_non_constant_offset(interpolation_point)
149 }
150 }
151
152 fn lagrange_interpolate_with_constant_offset(
153 &self,
154 interpolation_point: &FpVar<F>,
155 ) -> Result<FpVar<F>, SynthesisError> {
156 let lagrange_interpolator = self
157 .lagrange_interpolator
158 .as_ref()
159 .expect("lagrange interpolator has not been initialized. ");
160 let lagrange_coeffs = self.compute_lagrange_coefficients(interpolation_point)?;
161
162 let interpolation = lagrange_coeffs
163 .iter()
164 .zip(&self.evals)
165 .take(lagrange_interpolator.domain_order)
166 .map(|(coeff, eval)| coeff * eval)
167 .sum::<FpVar<F>>();
168
169 Ok(interpolation)
170 }
171
172 /// Generate interpolation constraints. We assume at compile time we know
173 /// the base coset (i.e. `gen`) but not know `offset`.
174 fn lagrange_interpolate_with_non_constant_offset(
175 &self,
176 interpolation_point: &FpVar<F>, // = alpha in the following code
177 ) -> Result<FpVar<F>, SynthesisError> {
178 // first, make sure `subgroup_points` is made
179 let subgroup_points = self.subgroup_points.as_ref()
180 .expect("lagrange interpolator has not been initialized. \
181 Call `self.generate_interpolation_cache` first or set `interpolate` to true in constructor. ");
182 // Let denote interpolation_point as alpha.
183 // Lagrange polynomial for coset element `a` is
184 // \frac{1}{size * offset ^ size} * \frac{alpha^size - offset^size}{alpha *
185 // a^{-1} - 1} Notice that a = (offset * a') where a' is the
186 // corresponding element of base coset
187
188 // let `lhs` be \frac{alpha^size - offset^size}{size * offset ^ size}. This part
189 // is shared by all lagrange polynomials
190 let coset_offset_to_size = self
191 .domain
192 .offset()
193 .pow_by_constant(&[self.domain.size()])?; // offset^size
194 let alpha_to_size = interpolation_point.pow_by_constant(&[self.domain.size()])?;
195 let lhs_numerator = &alpha_to_size - &coset_offset_to_size;
196 // This enforces that `alpha` is not in the coset.
197 // This also means that the denominator is
198 lhs_numerator.enforce_not_equal(&FpVar::zero())?;
199
200 // `domain.offset()` is non-zero by construction, so `coset_offset_to_size` is
201 // also non-zero, which means `lhs_denominator` is non-zero
202 let lhs_denominator = &coset_offset_to_size * FpVar::constant(F::from(self.domain.size()));
203
204 // unchecked is okay because the denominator is non-zero.
205 let lhs = lhs_numerator.mul_by_inverse_unchecked(&lhs_denominator)?;
206
207 // `rhs` for coset element `a` is \frac{1}{alpha * a^{-1} - 1} = \frac{1}{alpha
208 // * offset^{-1} * a'^{-1} - 1} domain.offset() is non-zero by
209 // construction.
210 let alpha_coset_offset_inv =
211 interpolation_point.mul_by_inverse_unchecked(&self.domain.offset())?;
212
213 let domain_size = self.domain.size() as usize;
214
215 // `evals` stores all lagrange polynomials evaluated at alpha
216 let evals = (0..domain_size)
217 .map(|i| {
218 // a'^{-1} where a is the base coset element
219 let subgroup_point_inv = subgroup_points[(domain_size - i) % domain_size];
220 debug_assert_eq!(subgroup_points[i] * subgroup_point_inv, F::one());
221 // alpha * offset^{-1} * a'^{-1} - 1
222 let lag_denom = &alpha_coset_offset_inv * subgroup_point_inv - F::one();
223 // lag_denom cannot be zero, so we use `unchecked`.
224 //
225 // Proof: lag_denom is zero if and only if alpha * (coset_offset *
226 // subgroup_point)^{-1} == 1. This can happen only if `alpha` is
227 // itself in the coset.
228 //
229 // Earlier we asserted that `lhs_numerator` is not zero.
230 // Since `lhs_numerator` is just the vanishing polynomial for the coset
231 // evaluated at `alpha`, and since this is non-zero, `alpha` is not
232 // in the coset.
233 let lag_coeff = lhs.mul_by_inverse_unchecked(&lag_denom)?;
234
235 Ok(&self.evals[i] * lag_coeff)
236 })
237 .collect::<Result<Vec<_>, _>>()?;
238
239 let res = evals.iter().sum();
240
241 Ok(res)
242 }
243}
244
245impl<'a, 'b, F: PrimeField> Add<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
246 type Output = EvaluationsVar<F>;
247
248 fn add(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
249 let mut result = self.clone();
250 result += rhs;
251 result
252 }
253}
254
255impl<'a, F: PrimeField> AddAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
256 /// Performs the `+=` operations, assuming `domain.offset` is equal.
257 fn add_assign(&mut self, other: &'a EvaluationsVar<F>) {
258 // offset might be unknown at compile time, so we assume offset is equal
259 assert!(
260 self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
261 "domains are unequal"
262 );
263
264 self.lagrange_interpolator = None;
265 self.evals
266 .iter_mut()
267 .zip(&other.evals)
268 .for_each(|(a, b)| *a = &*a + b)
269 }
270}
271
272impl<'a, 'b, F: PrimeField> Sub<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
273 type Output = EvaluationsVar<F>;
274
275 fn sub(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
276 let mut result = self.clone();
277 result -= rhs;
278 result
279 }
280}
281
282impl<'a, F: PrimeField> SubAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
283 /// Performs the `-=` operations, assuming `domain.offset` is equal.
284 fn sub_assign(&mut self, other: &'a EvaluationsVar<F>) {
285 // offset might be unknown at compile time, so we assume offset is equal
286 assert!(
287 self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
288 "domains are unequal"
289 );
290
291 self.lagrange_interpolator = None;
292 self.evals
293 .iter_mut()
294 .zip(&other.evals)
295 .for_each(|(a, b)| *a = &*a - b)
296 }
297}
298
299impl<'a, 'b, F: PrimeField> Mul<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
300 type Output = EvaluationsVar<F>;
301
302 /// Performs the `*` operations, assuming `domain.offset` is equal.
303 fn mul(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
304 let mut result = self.clone();
305 result *= rhs;
306 result
307 }
308}
309
310impl<'a, F: PrimeField> MulAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
311 /// Performs the `*=` operations, assuming `domain.offset` is equal.
312 fn mul_assign(&mut self, other: &'a EvaluationsVar<F>) {
313 // offset might be unknown at compile time, so we assume offset is equal
314 assert!(
315 self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
316 "domains are unequal"
317 );
318
319 self.lagrange_interpolator = None;
320 self.evals
321 .iter_mut()
322 .zip(&other.evals)
323 .for_each(|(a, b)| *a = &*a * b)
324 }
325}
326
327impl<'a, 'b, F: PrimeField> Div<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
328 type Output = EvaluationsVar<F>;
329
330 fn div(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
331 let mut result = self.clone();
332 result /= rhs;
333 result
334 }
335}
336
337impl<'a, F: PrimeField> DivAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
338 /// Performs the `/=` operations, assuming `domain.offset` is equal.
339 fn div_assign(&mut self, other: &'a EvaluationsVar<F>) {
340 // offset might be unknown at compile time, so we assume offset is equal
341 assert!(
342 self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
343 "domains are unequal"
344 );
345 let cs = self.evals[0].cs();
346 // the prover can generate result = (1 / other) * self offline
347 let mut result_val: Vec<_> = other.evals.iter().map(|x| x.value().unwrap()).collect();
348 batch_inversion(&mut result_val);
349 result_val
350 .iter_mut()
351 .zip(&self.evals)
352 .for_each(|(a, self_var)| *a *= self_var.value().unwrap());
353 let result_var: Vec<_> = result_val
354 .iter()
355 .map(|x| FpVar::new_witness(ns!(cs, "div result"), || Ok(*x)).unwrap())
356 .collect();
357 // enforce constraint
358 for i in 0..result_var.len() {
359 result_var[i]
360 .mul_equals(&other.evals[i], &self.evals[i])
361 .unwrap();
362 }
363
364 self.lagrange_interpolator = None;
365 self.evals = result_var
366 }
367}
368
369#[cfg(test)]
370mod tests {
371 use crate::{
372 alloc::AllocVar,
373 fields::{fp::FpVar, FieldVar},
374 poly::{domain::Radix2DomainVar, evaluations::univariate::EvaluationsVar},
375 R1CSVar,
376 };
377 use ark_ff::{FftField, Field, One, UniformRand};
378 use ark_poly::{polynomial::univariate::DensePolynomial, DenseUVPolynomial, Polynomial};
379 use ark_relations::r1cs::ConstraintSystem;
380 use ark_std::test_rng;
381 use ark_test_curves::bls12_381::Fr;
382
383 #[test]
384 fn test_interpolate_constant_offset() {
385 for n in [11, 12, 13, 14] {
386 let mut rng = test_rng();
387
388 let poly = DensePolynomial::rand((1 << n) - 1, &mut rng);
389 let gen = Fr::get_root_of_unity(1 << n).unwrap();
390 assert_eq!(gen.pow(&[1 << n]), Fr::one());
391 let domain = Radix2DomainVar::new(gen, n, FpVar::constant(Fr::rand(&mut rng))).unwrap();
392 let mut coset_point = domain.offset().value().unwrap();
393 let mut oracle_evals = Vec::new();
394 for _ in 0..(1 << n) {
395 oracle_evals.push(poly.evaluate(&coset_point));
396 coset_point *= gen;
397 }
398 let cs = ConstraintSystem::new_ref();
399 let evaluations_fp: Vec<_> = oracle_evals
400 .iter()
401 .map(|x| FpVar::new_input(ns!(cs, "evaluations"), || Ok(x)).unwrap())
402 .collect();
403 let evaluations_var = EvaluationsVar::from_vec_and_domain(evaluations_fp, domain, true);
404
405 let interpolate_point = Fr::rand(&mut rng);
406 let interpolate_point_fp =
407 FpVar::new_input(ns!(cs, "interpolate point"), || Ok(interpolate_point)).unwrap();
408
409 let expected = poly.evaluate(&interpolate_point);
410
411 let actual = evaluations_var
412 .interpolate_and_evaluate(&interpolate_point_fp)
413 .unwrap()
414 .value()
415 .unwrap();
416
417 assert_eq!(actual, expected);
418 assert!(cs.is_satisfied().unwrap());
419 println!("number of constraints: {}", cs.num_constraints());
420 }
421 }
422
423 #[test]
424 fn test_interpolate_non_constant_offset() {
425 for n in [11, 12, 13, 14] {
426 let mut rng = test_rng();
427 let poly = DensePolynomial::rand((1 << n) - 1, &mut rng);
428 let gen = Fr::get_root_of_unity(1 << n).unwrap();
429 assert_eq!(gen.pow(&[1 << n]), Fr::one());
430 let cs = ConstraintSystem::new_ref();
431 let domain = Radix2DomainVar::new(
432 gen,
433 n,
434 FpVar::new_witness(ns!(cs, "offset"), || Ok(Fr::rand(&mut rng))).unwrap(),
435 )
436 .unwrap();
437 let mut coset_point = domain.offset().value().unwrap();
438 let mut oracle_evals = Vec::new();
439 for _ in 0..(1 << n) {
440 oracle_evals.push(poly.evaluate(&coset_point));
441 coset_point *= gen;
442 }
443
444 let evaluations_fp: Vec<_> = oracle_evals
445 .iter()
446 .map(|x| FpVar::new_input(ns!(cs, "evaluations"), || Ok(x)).unwrap())
447 .collect();
448 let evaluations_var = EvaluationsVar::from_vec_and_domain(evaluations_fp, domain, true);
449
450 let interpolate_point = Fr::rand(&mut rng);
451 let interpolate_point_fp =
452 FpVar::new_input(ns!(cs, "interpolate point"), || Ok(interpolate_point)).unwrap();
453
454 let expected = poly.evaluate(&interpolate_point);
455
456 let actual = evaluations_var
457 .interpolate_and_evaluate(&interpolate_point_fp)
458 .unwrap()
459 .value()
460 .unwrap();
461
462 assert_eq!(actual, expected);
463 assert!(cs.is_satisfied().unwrap());
464 println!("number of constraints: {}", cs.num_constraints());
465 }
466 }
467
468 #[test]
469 fn test_division() {
470 let mut rng = test_rng();
471 let gen = Fr::get_root_of_unity(1 << 4).unwrap();
472 assert_eq!(gen.pow(&[1 << 4]), Fr::one());
473 let domain = Radix2DomainVar::new(
474 gen,
475 4, // 2^4 = 16
476 FpVar::constant(Fr::GENERATOR),
477 )
478 .unwrap();
479 let cs = ConstraintSystem::new_ref();
480
481 let ev_a = EvaluationsVar::from_vec_and_domain(
482 (0..16)
483 .map(|_| FpVar::new_input(ns!(cs, "poly_a"), || Ok(Fr::rand(&mut rng))).unwrap())
484 .collect(),
485 domain.clone(),
486 false,
487 );
488 let ev_b = EvaluationsVar::from_vec_and_domain(
489 (0..16)
490 .map(|_| FpVar::new_input(ns!(cs, "poly_a"), || Ok(Fr::rand(&mut rng))).unwrap())
491 .collect(),
492 domain.clone(),
493 false,
494 );
495
496 let a_div_b = (&ev_a) / (&ev_b);
497 assert!(cs.is_satisfied().unwrap());
498 let b_div_a = (&ev_b) / (&ev_a);
499
500 let one = &a_div_b * &b_div_a;
501 for ev in one.evals.iter() {
502 assert!(Fr::is_one(&ev.value().unwrap()))
503 }
504
505 assert!(cs.is_satisfied().unwrap());
506 }
507}