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