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