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::Radix2Domain;
  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: Radix2Domain<F>,
 22}
 23
 24impl<F: PrimeField> EvaluationsVar<F> {
 25    /// Construct `Self` from evaluations and a domain.
 26    /// `interpolate` indicates if user wants to interpolate this polynomial
 27    /// using lagrange interpolation.
 28    pub fn from_vec_and_domain(
 29        evaluations: Vec<FpVar<F>>,
 30        domain: Radix2Domain<F>,
 31        interpolate: bool,
 32    ) -> Self {
 33        assert_eq!(
 34            evaluations.len(),
 35            1 << domain.dim,
 36            "evaluations and domain has different dimensions"
 37        );
 38
 39        let mut ev = Self {
 40            evals: evaluations,
 41            lagrange_interpolator: None,
 42            domain,
 43        };
 44        if interpolate {
 45            ev.generate_lagrange_interpolator();
 46        }
 47        ev
 48    }
 49
 50    /// Generate lagrange interpolator and mark it ready to interpolate
 51    pub fn generate_lagrange_interpolator(&mut self) {
 52        let poly_evaluations_val: Vec<_> = self.evals.iter().map(|v| v.value().unwrap()).collect();
 53        let domain = &self.domain;
 54        let lagrange_interpolator =
 55            LagrangeInterpolator::new(domain.offset, domain.gen, domain.dim, poly_evaluations_val);
 56        self.lagrange_interpolator = Some(lagrange_interpolator)
 57    }
 58
 59    fn compute_lagrange_coefficients(
 60        &self,
 61        interpolation_point: &FpVar<F>,
 62    ) -> Result<Vec<FpVar<F>>, SynthesisError> {
 63        // ref: https://github.com/alexchmit/perfect-constraints/blob/79692f2652a95a57f2c7187f5b5276345e680230/fractal/src/algebra/lagrange_interpolation.rs#L159
 64        let cs = interpolation_point.cs();
 65        let t = interpolation_point;
 66        let lagrange_interpolator = self
 67            .lagrange_interpolator
 68            .as_ref()
 69            .expect("lagrange interpolator has not been initialized. \
 70            Call `self.generate_lagrange_interpolator` first or set `interpolate` to true in constructor. ");
 71        let lagrange_coeffs =
 72            lagrange_interpolator.compute_lagrange_coefficients(t.value().unwrap());
 73        let mut lagrange_coeffs_fg = Vec::new();
 74        // Now we convert these lagrange coefficients to gadgets, and then constrain them.
 75        // The i-th lagrange coefficients constraint is:
 76        // (v_inv[i] * t - v_inv[i] * domain_elem[i]) * (coeff) = 1/Z_I(t)
 77        let vp_t = lagrange_interpolator.domain_vp.evaluate_constraints(t)?;
 78        // let inv_vp_t = vp_t.inverse()?;
 79        for i in 0..lagrange_interpolator.domain_order {
 80            let constant: F =
 81                (-lagrange_interpolator.all_domain_elems[i]) * lagrange_interpolator.v_inv_elems[i];
 82            let mut a_element: FpVar<F> =
 83                t * &FpVar::constant(lagrange_interpolator.v_inv_elems[i]);
 84            a_element += FpVar::constant(constant);
 85
 86            let lag_coeff: FpVar<F> =
 87                FpVar::new_witness(ns!(cs, "generate lagrange coefficient"), || {
 88                    Ok(lagrange_coeffs[i])
 89                })?;
 90            // Enforce the actual constraint (A_element) * (lagrange_coeff) = 1/Z_I(t)
 91            assert_eq!(
 92                (lagrange_interpolator.v_inv_elems[i] * t.value().unwrap()
 93                    - lagrange_interpolator.v_inv_elems[i]
 94                        * lagrange_interpolator.all_domain_elems[i])
 95                    * lagrange_coeffs[i],
 96                vp_t.value().unwrap()
 97            );
 98            a_element.mul_equals(&lag_coeff, &vp_t)?;
 99            lagrange_coeffs_fg.push(lag_coeff);
100        }
101        Ok(lagrange_coeffs_fg)
102    }
103
104    /// Returns constraints for Interpolating and evaluating at `interpolation_point`
105    pub fn interpolate_and_evaluate(
106        &self,
107        interpolation_point: &FpVar<F>,
108    ) -> Result<FpVar<F>, SynthesisError> {
109        let lagrange_interpolator = self
110            .lagrange_interpolator
111            .as_ref()
112            .expect("lagrange interpolator has not been initialized. ");
113        let lagrange_coeffs = self.compute_lagrange_coefficients(interpolation_point)?;
114        let mut interpolation: FpVar<F> = FpVar::zero();
115        for i in 0..lagrange_interpolator.domain_order {
116            let intermediate = &lagrange_coeffs[i] * &self.evals[i];
117            interpolation += &intermediate
118        }
119
120        Ok(interpolation)
121    }
122}
123
124impl<'a, 'b, F: PrimeField> Add<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
125    type Output = EvaluationsVar<F>;
126
127    fn add(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
128        let mut result = self.clone();
129        result += rhs;
130        result
131    }
132}
133
134impl<'a, F: PrimeField> AddAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
135    fn add_assign(&mut self, other: &'a EvaluationsVar<F>) {
136        assert_eq!(self.domain, other.domain, "domains are unequal");
137
138        self.lagrange_interpolator = None;
139        self.evals
140            .iter_mut()
141            .zip(&other.evals)
142            .for_each(|(a, b)| *a = &*a + b)
143    }
144}
145
146impl<'a, 'b, F: PrimeField> Sub<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
147    type Output = EvaluationsVar<F>;
148
149    fn sub(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
150        let mut result = self.clone();
151        result -= rhs;
152        result
153    }
154}
155
156impl<'a, F: PrimeField> SubAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
157    fn sub_assign(&mut self, other: &'a EvaluationsVar<F>) {
158        assert_eq!(self.domain, other.domain, "domains are unequal");
159
160        self.lagrange_interpolator = None;
161        self.evals
162            .iter_mut()
163            .zip(&other.evals)
164            .for_each(|(a, b)| *a = &*a - b)
165    }
166}
167
168impl<'a, 'b, F: PrimeField> Mul<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
169    type Output = EvaluationsVar<F>;
170
171    fn mul(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
172        let mut result = self.clone();
173        result *= rhs;
174        result
175    }
176}
177
178impl<'a, F: PrimeField> MulAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
179    fn mul_assign(&mut self, other: &'a EvaluationsVar<F>) {
180        assert_eq!(self.domain, other.domain, "domains are unequal");
181
182        self.lagrange_interpolator = None;
183        self.evals
184            .iter_mut()
185            .zip(&other.evals)
186            .for_each(|(a, b)| *a = &*a * b)
187    }
188}
189
190impl<'a, 'b, F: PrimeField> Div<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
191    type Output = EvaluationsVar<F>;
192
193    fn div(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
194        let mut result = self.clone();
195        result /= rhs;
196        result
197    }
198}
199
200impl<'a, F: PrimeField> DivAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
201    fn div_assign(&mut self, other: &'a EvaluationsVar<F>) {
202        assert_eq!(self.domain, other.domain, "domains are unequal");
203        let cs = self.evals[0].cs();
204        // the prover can generate result = (1 / other) * self offline
205        let mut result_val: Vec<_> = other.evals.iter().map(|x| x.value().unwrap()).collect();
206        batch_inversion(&mut result_val);
207        result_val
208            .iter_mut()
209            .zip(&self.evals)
210            .for_each(|(a, self_var)| *a *= self_var.value().unwrap());
211        let result_var: Vec<_> = result_val
212            .iter()
213            .map(|x| FpVar::new_witness(ns!(cs, "div result"), || Ok(*x)).unwrap())
214            .collect();
215        // enforce constraint
216        for i in 0..result_var.len() {
217            result_var[i]
218                .mul_equals(&other.evals[i], &self.evals[i])
219                .unwrap();
220        }
221
222        self.lagrange_interpolator = None;
223        self.evals = result_var
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use crate::alloc::AllocVar;
230    use crate::fields::fp::FpVar;
231    use crate::poly::domain::Radix2Domain;
232    use crate::poly::evaluations::univariate::EvaluationsVar;
233    use crate::R1CSVar;
234    use ark_ff::{FftField, Field, One, UniformRand};
235    use ark_poly::polynomial::univariate::DensePolynomial;
236    use ark_poly::{Polynomial, UVPolynomial};
237    use ark_relations::r1cs::ConstraintSystem;
238    use ark_std::test_rng;
239    use ark_test_curves::bls12_381::Fr;
240
241    #[test]
242    fn test_interpolate() {
243        let mut rng = test_rng();
244        let poly = DensePolynomial::rand(15, &mut rng);
245        let gen = Fr::get_root_of_unity(1 << 4).unwrap();
246        assert_eq!(gen.pow(&[1 << 4]), Fr::one());
247        let domain = Radix2Domain {
248            gen,
249            offset: Fr::multiplicative_generator(),
250            dim: 4, // 2^4 = 16
251        };
252        let mut coset_point = domain.offset;
253        let mut oracle_evals = Vec::new();
254        for _ in 0..(1 << 4) {
255            oracle_evals.push(poly.evaluate(&coset_point));
256            coset_point *= gen;
257        }
258        let cs = ConstraintSystem::new_ref();
259        let evaluations_fp: Vec<_> = oracle_evals
260            .iter()
261            .map(|x| FpVar::new_input(ns!(cs, "evaluations"), || Ok(x)).unwrap())
262            .collect();
263        let evaluations_var = EvaluationsVar::from_vec_and_domain(evaluations_fp, domain, true);
264
265        let interpolate_point = Fr::rand(&mut rng);
266        let interpolate_point_fp =
267            FpVar::new_input(ns!(cs, "interpolate point"), || Ok(interpolate_point)).unwrap();
268
269        let expected = poly.evaluate(&interpolate_point);
270
271        let actual = evaluations_var
272            .interpolate_and_evaluate(&interpolate_point_fp)
273            .unwrap()
274            .value()
275            .unwrap();
276
277        assert_eq!(actual, expected);
278        assert!(cs.is_satisfied().unwrap());
279    }
280
281    #[test]
282    fn test_division() {
283        let mut rng = test_rng();
284        let gen = Fr::get_root_of_unity(1 << 4).unwrap();
285        assert_eq!(gen.pow(&[1 << 4]), Fr::one());
286        let domain = Radix2Domain {
287            gen,
288            offset: Fr::multiplicative_generator(),
289            dim: 4, // 2^4 = 16
290        };
291
292        let cs = ConstraintSystem::new_ref();
293
294        let ev_a = EvaluationsVar::from_vec_and_domain(
295            (0..16)
296                .map(|_| FpVar::new_input(ns!(cs, "poly_a"), || Ok(Fr::rand(&mut rng))).unwrap())
297                .collect(),
298            domain.clone(),
299            false,
300        );
301        let ev_b = EvaluationsVar::from_vec_and_domain(
302            (0..16)
303                .map(|_| FpVar::new_input(ns!(cs, "poly_a"), || Ok(Fr::rand(&mut rng))).unwrap())
304                .collect(),
305            domain.clone(),
306            false,
307        );
308
309        let a_div_b = (&ev_a) / (&ev_b);
310        assert!(cs.is_satisfied().unwrap());
311        let b_div_a = (&ev_b) / (&ev_a);
312
313        let one = &a_div_b * &b_div_a;
314        for ev in one.evals.iter() {
315            assert!(Fr::is_one(&ev.value().unwrap()))
316        }
317
318        assert!(cs.is_satisfied().unwrap());
319    }
320}