1use ark_ec::mnt6::{
  2    g2::{AteAdditionCoefficients, AteDoubleCoefficients},
  3    G1Prepared, G2Prepared, MNT6Config,
  4};
  5use ark_ff::Field;
  6use ark_relations::r1cs::{Namespace, SynthesisError};
  7
  8use crate::{
  9    fields::{fp::FpVar, fp3::Fp3Var, FieldVar},
 10    groups::curves::short_weierstrass::ProjectiveVar,
 11    pairing::mnt6::PairingVar,
 12    prelude::*,
 13    Vec,
 14};
 15use core::borrow::Borrow;
 16
 17/// Represents a projective point in G1.
 18pub type G1Var<P> = ProjectiveVar<<P as MNT6Config>::G1Config, FpVar<<P as MNT6Config>::Fp>>;
 19
 20/// Represents a projective point in G2.
 21pub type G2Var<P> = ProjectiveVar<<P as MNT6Config>::G2Config, Fp3G<P>>;
 22
 23/// Represents the cached precomputation that can be performed on a G1 element
 24/// which enables speeding up pairing computation.
 25#[derive(Derivative)]
 26#[derivative(Clone(bound = "P: MNT6Config"), Debug(bound = "P: MNT6Config"))]
 27pub struct G1PreparedVar<P: MNT6Config> {
 28    #[doc(hidden)]
 29    pub x: FpVar<P::Fp>,
 30    #[doc(hidden)]
 31    pub y: FpVar<P::Fp>,
 32    #[doc(hidden)]
 33    pub x_twist: Fp3Var<P::Fp3Config>,
 34    #[doc(hidden)]
 35    pub y_twist: Fp3Var<P::Fp3Config>,
 36}
 37
 38impl<P: MNT6Config> G1PreparedVar<P> {
 39    /// Returns the value assigned to `self` in the underlying constraint
 40    /// system.
 41    pub fn value(&self) -> Result<G1Prepared<P>, SynthesisError> {
 42        let x = self.x.value()?;
 43        let y = self.y.value()?;
 44        let x_twist = self.x_twist.value()?;
 45        let y_twist = self.y_twist.value()?;
 46        Ok(G1Prepared {
 47            x,
 48            y,
 49            x_twist,
 50            y_twist,
 51        })
 52    }
 53
 54    /// Constructs `Self` from a `G1Var`.
 55    #[tracing::instrument(target = "r1cs")]
 56    pub fn from_group_var(q: &G1Var<P>) -> Result<Self, SynthesisError> {
 57        let q = q.to_affine()?;
 58        let zero = FpVar::<P::Fp>::zero();
 59        let x_twist = Fp3Var::new(q.x.clone(), zero.clone(), zero.clone()) * P::TWIST;
 60        let y_twist = Fp3Var::new(q.y.clone(), zero.clone(), zero) * P::TWIST;
 61        let result = G1PreparedVar {
 62            x: q.x,
 63            y: q.y,
 64            x_twist,
 65            y_twist,
 66        };
 67        Ok(result)
 68    }
 69}
 70
 71impl<P: MNT6Config> AllocVar<G1Prepared<P>, P::Fp> for G1PreparedVar<P> {
 72    #[tracing::instrument(target = "r1cs", skip(cs, f))]
 73    fn new_variable<T: Borrow<G1Prepared<P>>>(
 74        cs: impl Into<Namespace<P::Fp>>,
 75        f: impl FnOnce() -> Result<T, SynthesisError>,
 76        mode: AllocationMode,
 77    ) -> Result<Self, SynthesisError> {
 78        let ns = cs.into();
 79        let cs = ns.cs();
 80
 81        let g1_prep = f().map(|b| *b.borrow());
 82
 83        let x = FpVar::new_variable(ark_relations::ns!(cs, "x"), || g1_prep.map(|g| g.x), mode)?;
 84        let y = FpVar::new_variable(ark_relations::ns!(cs, "y"), || g1_prep.map(|g| g.y), mode)?;
 85        let x_twist = Fp3Var::new_variable(
 86            ark_relations::ns!(cs, "x_twist"),
 87            || g1_prep.map(|g| g.x_twist),
 88            mode,
 89        )?;
 90        let y_twist = Fp3Var::new_variable(
 91            ark_relations::ns!(cs, "y_twist"),
 92            || g1_prep.map(|g| g.y_twist),
 93            mode,
 94        )?;
 95        Ok(Self {
 96            x,
 97            y,
 98            x_twist,
 99            y_twist,
100        })
101    }
102}
103
104impl<P: MNT6Config> ToBytesGadget<P::Fp> for G1PreparedVar<P> {
105    #[inline]
106    #[tracing::instrument(target = "r1cs")]
107    fn to_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
108        let mut x = self.x.to_bytes()?;
109        let mut y = self.y.to_bytes()?;
110        let mut x_twist = self.x_twist.to_bytes()?;
111        let mut y_twist = self.y_twist.to_bytes()?;
112
113        x.append(&mut y);
114        x.append(&mut x_twist);
115        x.append(&mut y_twist);
116        Ok(x)
117    }
118
119    #[tracing::instrument(target = "r1cs")]
120    fn to_non_unique_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
121        let mut x = self.x.to_non_unique_bytes()?;
122        let mut y = self.y.to_non_unique_bytes()?;
123        let mut x_twist = self.x_twist.to_non_unique_bytes()?;
124        let mut y_twist = self.y_twist.to_non_unique_bytes()?;
125
126        x.append(&mut y);
127        x.append(&mut x_twist);
128        x.append(&mut y_twist);
129        Ok(x)
130    }
131}
132
133type Fp3G<P> = Fp3Var<<P as MNT6Config>::Fp3Config>;
134
135/// Represents the cached precomputation that can be performed on a G2 element
136/// which enables speeding up pairing computation.
137#[derive(Derivative)]
138#[derivative(Clone(bound = "P: MNT6Config"), Debug(bound = "P: MNT6Config"))]
139pub struct G2PreparedVar<P: MNT6Config> {
140    #[doc(hidden)]
141    pub x: Fp3Var<P::Fp3Config>,
142    #[doc(hidden)]
143    pub y: Fp3Var<P::Fp3Config>,
144    #[doc(hidden)]
145    pub x_over_twist: Fp3Var<P::Fp3Config>,
146    #[doc(hidden)]
147    pub y_over_twist: Fp3Var<P::Fp3Config>,
148    #[doc(hidden)]
149    pub double_coefficients: Vec<AteDoubleCoefficientsVar<P>>,
150    #[doc(hidden)]
151    pub addition_coefficients: Vec<AteAdditionCoefficientsVar<P>>,
152}
153
154impl<P: MNT6Config> AllocVar<G2Prepared<P>, P::Fp> for G2PreparedVar<P> {
155    #[tracing::instrument(target = "r1cs", skip(cs, f))]
156    fn new_variable<T: Borrow<G2Prepared<P>>>(
157        cs: impl Into<Namespace<P::Fp>>,
158        f: impl FnOnce() -> Result<T, SynthesisError>,
159        mode: AllocationMode,
160    ) -> Result<Self, SynthesisError> {
161        let ns = cs.into();
162        let cs = ns.cs();
163
164        let g2_prep = f().map(|b| b.borrow().clone());
165        let g2 = g2_prep.as_ref().map_err(|e| *e);
166
167        let x = Fp3Var::new_variable(ark_relations::ns!(cs, "x"), || g2.map(|g| g.x), mode)?;
168        let y = Fp3Var::new_variable(ark_relations::ns!(cs, "y"), || g2.map(|g| g.y), mode)?;
169        let x_over_twist = Fp3Var::new_variable(
170            ark_relations::ns!(cs, "x_over_twist"),
171            || g2.map(|g| g.x_over_twist),
172            mode,
173        )?;
174        let y_over_twist = Fp3Var::new_variable(
175            ark_relations::ns!(cs, "y_over_twist"),
176            || g2.map(|g| g.y_over_twist),
177            mode,
178        )?;
179        let double_coefficients = Vec::new_variable(
180            ark_relations::ns!(cs, "double coeffs"),
181            || g2.map(|g| g.double_coefficients.clone()),
182            mode,
183        )?;
184        let addition_coefficients = Vec::new_variable(
185            ark_relations::ns!(cs, "add coeffs"),
186            || g2.map(|g| g.addition_coefficients.clone()),
187            mode,
188        )?;
189        Ok(Self {
190            x,
191            y,
192            x_over_twist,
193            y_over_twist,
194            double_coefficients,
195            addition_coefficients,
196        })
197    }
198}
199
200impl<P: MNT6Config> ToBytesGadget<P::Fp> for G2PreparedVar<P> {
201    #[inline]
202    #[tracing::instrument(target = "r1cs")]
203    fn to_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
204        let mut x = self.x.to_bytes()?;
205        let mut y = self.y.to_bytes()?;
206        let mut x_over_twist = self.x_over_twist.to_bytes()?;
207        let mut y_over_twist = self.y_over_twist.to_bytes()?;
208
209        x.append(&mut y);
210        x.append(&mut x_over_twist);
211        x.append(&mut y_over_twist);
212
213        for coeff in self.double_coefficients.iter() {
214            x.extend_from_slice(&coeff.to_bytes()?);
215        }
216        for coeff in self.addition_coefficients.iter() {
217            x.extend_from_slice(&coeff.to_bytes()?);
218        }
219        Ok(x)
220    }
221
222    #[tracing::instrument(target = "r1cs")]
223    fn to_non_unique_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
224        let mut x = self.x.to_non_unique_bytes()?;
225        let mut y = self.y.to_non_unique_bytes()?;
226        let mut x_over_twist = self.x_over_twist.to_non_unique_bytes()?;
227        let mut y_over_twist = self.y_over_twist.to_non_unique_bytes()?;
228
229        x.append(&mut y);
230        x.append(&mut x_over_twist);
231        x.append(&mut y_over_twist);
232
233        for coeff in self.double_coefficients.iter() {
234            x.extend_from_slice(&coeff.to_non_unique_bytes()?);
235        }
236        for coeff in self.addition_coefficients.iter() {
237            x.extend_from_slice(&coeff.to_non_unique_bytes()?);
238        }
239        Ok(x)
240    }
241}
242
243impl<P: MNT6Config> G2PreparedVar<P> {
244    /// Returns the value assigned to `self` in the underlying constraint
245    /// system.
246    pub fn value(&self) -> Result<G2Prepared<P>, SynthesisError> {
247        let x = self.x.value()?;
248        let y = self.y.value()?;
249        let x_over_twist = self.x_over_twist.value()?;
250        let y_over_twist = self.y_over_twist.value()?;
251        let double_coefficients = self
252            .double_coefficients
253            .iter()
254            .map(|coeff| coeff.value())
255            .collect::<Result<Vec<_>, SynthesisError>>()?;
256        let addition_coefficients = self
257            .addition_coefficients
258            .iter()
259            .map(|coeff| coeff.value())
260            .collect::<Result<Vec<_>, SynthesisError>>()?;
261        Ok(G2Prepared {
262            x,
263            y,
264            x_over_twist,
265            y_over_twist,
266            double_coefficients,
267            addition_coefficients,
268        })
269    }
270
271    /// Constructs `Self` from a `G2Var`.
272    #[tracing::instrument(target = "r1cs")]
273    pub fn from_group_var(q: &G2Var<P>) -> Result<Self, SynthesisError> {
274        let q = q.to_affine()?;
275        let twist_inv = P::TWIST.inverse().unwrap();
276
277        let mut g2p = G2PreparedVar {
278            x: q.x.clone(),
279            y: q.y.clone(),
280            x_over_twist: &q.x * twist_inv,
281            y_over_twist: &q.y * twist_inv,
282            double_coefficients: vec![],
283            addition_coefficients: vec![],
284        };
285
286        let mut r = G2ProjectiveExtendedVar {
287            x: q.x.clone(),
288            y: q.y.clone(),
289            z: Fp3G::<P>::one(),
290            t: Fp3G::<P>::one(),
291        };
292
293        for bit in P::ATE_LOOP_COUNT.iter().skip(1) {
294            let (r2, coeff) = PairingVar::<P>::doubling_step_for_flipped_miller_loop(&r)?;
295            g2p.double_coefficients.push(coeff);
296            r = r2;
297
298            let add_coeff;
299            let r_temp;
300            match bit {
301                1 => {
302                    (r_temp, add_coeff) =
303                        PairingVar::<P>::mixed_addition_step_for_flipped_miller_loop(
304                            &q.x, &q.y, &r,
305                        )?;
306                },
307                -1 => {
308                    (r_temp, add_coeff) =
309                        PairingVar::<P>::mixed_addition_step_for_flipped_miller_loop(
310                            &q.x,
311                            &q.y.negate()?,
312                            &r,
313                        )?;
314                },
315                _ => continue,
316            }
317            g2p.addition_coefficients.push(add_coeff);
318            r = r_temp;
319        }
320
321        if P::ATE_IS_LOOP_COUNT_NEG {
322            let rz_inv = r.z.inverse()?;
323            let rz2_inv = rz_inv.square()?;
324            let rz3_inv = &rz_inv * &rz2_inv;
325
326            let minus_r_affine_x = &r.x * &rz2_inv;
327            let minus_r_affine_y = r.y.negate()? * &rz3_inv;
328
329            let add_result = PairingVar::<P>::mixed_addition_step_for_flipped_miller_loop(
330                &minus_r_affine_x,
331                &minus_r_affine_y,
332                &r,
333            )?;
334            g2p.addition_coefficients.push(add_result.1);
335        }
336
337        Ok(g2p)
338    }
339}
340
341#[doc(hidden)]
342#[derive(Derivative)]
343#[derivative(Clone(bound = "P: MNT6Config"), Debug(bound = "P: MNT6Config"))]
344pub struct AteDoubleCoefficientsVar<P: MNT6Config> {
345    pub c_h: Fp3Var<P::Fp3Config>,
346    pub c_4c: Fp3Var<P::Fp3Config>,
347    pub c_j: Fp3Var<P::Fp3Config>,
348    pub c_l: Fp3Var<P::Fp3Config>,
349}
350
351impl<P: MNT6Config> AllocVar<AteDoubleCoefficients<P>, P::Fp> for AteDoubleCoefficientsVar<P> {
352    #[tracing::instrument(target = "r1cs", skip(cs, f))]
353    fn new_variable<T: Borrow<AteDoubleCoefficients<P>>>(
354        cs: impl Into<Namespace<P::Fp>>,
355        f: impl FnOnce() -> Result<T, SynthesisError>,
356        mode: AllocationMode,
357    ) -> Result<Self, SynthesisError> {
358        let ns = cs.into();
359        let cs = ns.cs();
360
361        let c_prep = f().map(|c| c.borrow().clone());
362        let c = c_prep.as_ref().map_err(|e| *e);
363
364        let c_h = Fp3Var::new_variable(ark_relations::ns!(cs, "c_h"), || c.map(|c| c.c_h), mode)?;
365        let c_4c =
366            Fp3Var::new_variable(ark_relations::ns!(cs, "c_4c"), || c.map(|c| c.c_4c), mode)?;
367        let c_j = Fp3Var::new_variable(ark_relations::ns!(cs, "c_j"), || c.map(|c| c.c_j), mode)?;
368        let c_l = Fp3Var::new_variable(ark_relations::ns!(cs, "c_l"), || c.map(|c| c.c_l), mode)?;
369        Ok(Self {
370            c_h,
371            c_4c,
372            c_j,
373            c_l,
374        })
375    }
376}
377
378impl<P: MNT6Config> ToBytesGadget<P::Fp> for AteDoubleCoefficientsVar<P> {
379    #[inline]
380    #[tracing::instrument(target = "r1cs")]
381    fn to_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
382        let mut c_h = self.c_h.to_bytes()?;
383        let mut c_4c = self.c_4c.to_bytes()?;
384        let mut c_j = self.c_j.to_bytes()?;
385        let mut c_l = self.c_l.to_bytes()?;
386
387        c_h.append(&mut c_4c);
388        c_h.append(&mut c_j);
389        c_h.append(&mut c_l);
390        Ok(c_h)
391    }
392
393    #[tracing::instrument(target = "r1cs")]
394    fn to_non_unique_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
395        let mut c_h = self.c_h.to_non_unique_bytes()?;
396        let mut c_4c = self.c_4c.to_non_unique_bytes()?;
397        let mut c_j = self.c_j.to_non_unique_bytes()?;
398        let mut c_l = self.c_l.to_non_unique_bytes()?;
399
400        c_h.append(&mut c_4c);
401        c_h.append(&mut c_j);
402        c_h.append(&mut c_l);
403        Ok(c_h)
404    }
405}
406
407impl<P: MNT6Config> AteDoubleCoefficientsVar<P> {
408    /// Returns the value assigned to `self` in the underlying constraint
409    /// system.
410    pub fn value(&self) -> Result<AteDoubleCoefficients<P>, SynthesisError> {
411        let c_h = self.c_h.value()?;
412        let c_4c = self.c_4c.value()?;
413        let c_j = self.c_j.value()?;
414        let c_l = self.c_l.value()?;
415        Ok(AteDoubleCoefficients {
416            c_h,
417            c_4c,
418            c_j,
419            c_l,
420        })
421    }
422}
423
424#[doc(hidden)]
425#[derive(Derivative)]
426#[derivative(Clone(bound = "P: MNT6Config"), Debug(bound = "P: MNT6Config"))]
427pub struct AteAdditionCoefficientsVar<P: MNT6Config> {
428    pub c_l1: Fp3Var<P::Fp3Config>,
429    pub c_rz: Fp3Var<P::Fp3Config>,
430}
431
432impl<P: MNT6Config> AllocVar<AteAdditionCoefficients<P>, P::Fp> for AteAdditionCoefficientsVar<P> {
433    #[tracing::instrument(target = "r1cs", skip(cs, f))]
434    fn new_variable<T: Borrow<AteAdditionCoefficients<P>>>(
435        cs: impl Into<Namespace<P::Fp>>,
436        f: impl FnOnce() -> Result<T, SynthesisError>,
437        mode: AllocationMode,
438    ) -> Result<Self, SynthesisError> {
439        let ns = cs.into();
440        let cs = ns.cs();
441
442        let c_prep = f().map(|c| c.borrow().clone());
443        let c = c_prep.as_ref().map_err(|e| *e);
444
445        let c_l1 =
446            Fp3Var::new_variable(ark_relations::ns!(cs, "c_l1"), || c.map(|c| c.c_l1), mode)?;
447        let c_rz =
448            Fp3Var::new_variable(ark_relations::ns!(cs, "c_rz"), || c.map(|c| c.c_rz), mode)?;
449        Ok(Self { c_l1, c_rz })
450    }
451}
452
453impl<P: MNT6Config> ToBytesGadget<P::Fp> for AteAdditionCoefficientsVar<P> {
454    #[inline]
455    #[tracing::instrument(target = "r1cs")]
456    fn to_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
457        let mut c_l1 = self.c_l1.to_bytes()?;
458        let mut c_rz = self.c_rz.to_bytes()?;
459
460        c_l1.append(&mut c_rz);
461        Ok(c_l1)
462    }
463
464    #[tracing::instrument(target = "r1cs")]
465    fn to_non_unique_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
466        let mut c_l1 = self.c_l1.to_non_unique_bytes()?;
467        let mut c_rz = self.c_rz.to_non_unique_bytes()?;
468
469        c_l1.append(&mut c_rz);
470        Ok(c_l1)
471    }
472}
473
474impl<P: MNT6Config> AteAdditionCoefficientsVar<P> {
475    /// Returns the value assigned to `self` in the underlying constraint
476    /// system.
477    pub fn value(&self) -> Result<AteAdditionCoefficients<P>, SynthesisError> {
478        let c_l1 = self.c_l1.value()?;
479        let c_rz = self.c_rz.value()?;
480        Ok(AteAdditionCoefficients { c_l1, c_rz })
481    }
482}
483
484#[doc(hidden)]
485pub struct G2ProjectiveExtendedVar<P: MNT6Config> {
486    pub x: Fp3Var<P::Fp3Config>,
487    pub y: Fp3Var<P::Fp3Config>,
488    pub z: Fp3Var<P::Fp3Config>,
489    pub t: Fp3Var<P::Fp3Config>,
490}