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