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