1use ark_ec::{
  2    bls12::{Bls12Config, G1Prepared, G2Prepared, TwistType},
  3    short_weierstrass::Affine as GroupAffine,
  4};
  5use ark_ff::{BitIteratorBE, Field, One};
  6use ark_relations::r1cs::{Namespace, SynthesisError};
  7
  8use crate::{
  9    fields::{fp::FpVar, fp2::Fp2Var, FieldVar},
 10    groups::curves::short_weierstrass::*,
 11    Vec,
 12};
 13use core::fmt::Debug;
 14
 15/// Represents a projective point in G1.
 16pub type G1Var<P> = ProjectiveVar<<P as Bls12Config>::G1Config, FpVar<<P as Bls12Config>::Fp>>;
 17
 18/// Represents an affine point on G1. Should be used only for comparison and
 19/// when a canonical representation of a point is required, and not for
 20/// arithmetic.
 21pub type G1AffineVar<P> = AffineVar<<P as Bls12Config>::G1Config, FpVar<<P as Bls12Config>::Fp>>;
 22
 23/// Represents a projective point in G2.
 24pub type G2Var<P> = ProjectiveVar<<P as Bls12Config>::G2Config, Fp2G<P>>;
 25/// Represents an affine point on G2. Should be used only for comparison and
 26/// when a canonical representation of a point is required, and not for
 27/// arithmetic.
 28pub type G2AffineVar<P> = AffineVar<<P as Bls12Config>::G2Config, Fp2G<P>>;
 29
 30/// Represents the cached precomputation that can be performed on a G1 element
 31/// which enables speeding up pairing computation.
 32#[derive(Derivative)]
 33#[derivative(Clone(bound = "G1Var<P>: Clone"), Debug(bound = "G1Var<P>: Debug"))]
 34pub struct G1PreparedVar<P: Bls12Config>(pub AffineVar<P::G1Config, FpVar<P::Fp>>);
 35
 36impl<P: Bls12Config> G1PreparedVar<P> {
 37    /// Returns the value assigned to `self` in the underlying constraint
 38    /// system.
 39    pub fn value(&self) -> Result<G1Prepared<P>, SynthesisError> {
 40        let x = self.0.x.value()?;
 41        let y = self.0.y.value()?;
 42        let infinity = self.0.infinity.value()?;
 43        let g = infinity
 44            .then_some(GroupAffine::identity())
 45            .unwrap_or(GroupAffine::new(x, y))
 46            .into();
 47        Ok(g)
 48    }
 49
 50    /// Constructs `Self` from a `G1Var`.
 51    pub fn from_group_var(q: &G1Var<P>) -> Result<Self, SynthesisError> {
 52        let g = q.to_affine()?;
 53        Ok(Self(g))
 54    }
 55}
 56
 57impl<P: Bls12Config> AllocVar<G1Prepared<P>, P::Fp> for G1PreparedVar<P> {
 58    fn new_variable<T: Borrow<G1Prepared<P>>>(
 59        cs: impl Into<Namespace<P::Fp>>,
 60        f: impl FnOnce() -> Result<T, SynthesisError>,
 61        mode: AllocationMode,
 62    ) -> Result<Self, SynthesisError> {
 63        let ns = cs.into();
 64        let cs = ns.cs();
 65        let g1_prep = f().map(|b| b.borrow().0);
 66
 67        let x = FpVar::new_variable(ark_relations::ns!(cs, "x"), || g1_prep.map(|g| g.x), mode)?;
 68        let y = FpVar::new_variable(ark_relations::ns!(cs, "y"), || g1_prep.map(|g| g.y), mode)?;
 69        let infinity = Boolean::new_variable(
 70            ark_relations::ns!(cs, "inf"),
 71            || g1_prep.map(|g| g.infinity),
 72            mode,
 73        )?;
 74        let g = AffineVar::new(x, y, infinity);
 75        Ok(Self(g))
 76    }
 77}
 78
 79impl<P: Bls12Config> ToBytesGadget<P::Fp> for G1PreparedVar<P> {
 80    #[inline]
 81    #[tracing::instrument(target = "r1cs")]
 82    fn to_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
 83        let mut bytes = self.0.x.to_bytes()?;
 84        let y_bytes = self.0.y.to_bytes()?;
 85        let inf_bytes = self.0.infinity.to_bytes()?;
 86        bytes.extend_from_slice(&y_bytes);
 87        bytes.extend_from_slice(&inf_bytes);
 88        Ok(bytes)
 89    }
 90
 91    #[tracing::instrument(target = "r1cs")]
 92    fn to_non_unique_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
 93        let mut bytes = self.0.x.to_non_unique_bytes()?;
 94        let y_bytes = self.0.y.to_non_unique_bytes()?;
 95        let inf_bytes = self.0.infinity.to_non_unique_bytes()?;
 96        bytes.extend_from_slice(&y_bytes);
 97        bytes.extend_from_slice(&inf_bytes);
 98        Ok(bytes)
 99    }
100}
101
102type Fp2G<P> = Fp2Var<<P as Bls12Config>::Fp2Config>;
103type LCoeff<P> = (Fp2G<P>, Fp2G<P>);
104/// Represents the cached precomputation that can be performed on a G2 element
105/// which enables speeding up pairing computation.
106#[derive(Derivative)]
107#[derivative(
108    Clone(bound = "Fp2Var<P::Fp2Config>: Clone"),
109    Debug(bound = "Fp2Var<P::Fp2Config>: Debug")
110)]
111pub struct G2PreparedVar<P: Bls12Config> {
112    #[doc(hidden)]
113    pub ell_coeffs: Vec<LCoeff<P>>,
114}
115
116impl<P: Bls12Config> AllocVar<G2Prepared<P>, P::Fp> for G2PreparedVar<P> {
117    #[tracing::instrument(target = "r1cs", skip(cs, f, mode))]
118    fn new_variable<T: Borrow<G2Prepared<P>>>(
119        cs: impl Into<Namespace<P::Fp>>,
120        f: impl FnOnce() -> Result<T, SynthesisError>,
121        mode: AllocationMode,
122    ) -> Result<Self, SynthesisError> {
123        let ns = cs.into();
124        let cs = ns.cs();
125        let g2_prep = f().map(|b| {
126            let projective_coeffs = &b.borrow().ell_coeffs;
127            match P::TWIST_TYPE {
128                TwistType::M => {
129                    let mut z_s = projective_coeffs
130                        .iter()
131                        .map(|(_, _, z)| *z)
132                        .collect::<Vec<_>>();
133                    ark_ff::fields::batch_inversion(&mut z_s);
134                    projective_coeffs
135                        .iter()
136                        .zip(z_s)
137                        .map(|((x, y, _), z_inv)| (*x * &z_inv, *y * &z_inv))
138                        .collect::<Vec<_>>()
139                },
140                TwistType::D => {
141                    let mut z_s = projective_coeffs
142                        .iter()
143                        .map(|(z, ..)| *z)
144                        .collect::<Vec<_>>();
145                    ark_ff::fields::batch_inversion(&mut z_s);
146                    projective_coeffs
147                        .iter()
148                        .zip(z_s)
149                        .map(|((_, x, y), z_inv)| (*x * &z_inv, *y * &z_inv))
150                        .collect::<Vec<_>>()
151                },
152            }
153        });
154
155        let l = Vec::new_variable(
156            ark_relations::ns!(cs, "l"),
157            || {
158                g2_prep
159                    .clone()
160                    .map(|c| c.iter().map(|(l, _)| *l).collect::<Vec<_>>())
161            },
162            mode,
163        )?;
164        let r = Vec::new_variable(
165            ark_relations::ns!(cs, "r"),
166            || g2_prep.map(|c| c.iter().map(|(_, r)| *r).collect::<Vec<_>>()),
167            mode,
168        )?;
169        let ell_coeffs = l.into_iter().zip(r).collect();
170        Ok(Self { ell_coeffs })
171    }
172}
173
174impl<P: Bls12Config> ToBytesGadget<P::Fp> for G2PreparedVar<P> {
175    #[inline]
176    #[tracing::instrument(target = "r1cs")]
177    fn to_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
178        let mut bytes = Vec::new();
179        for coeffs in &self.ell_coeffs {
180            bytes.extend_from_slice(&coeffs.0.to_bytes()?);
181            bytes.extend_from_slice(&coeffs.1.to_bytes()?);
182        }
183        Ok(bytes)
184    }
185
186    #[tracing::instrument(target = "r1cs")]
187    fn to_non_unique_bytes(&self) -> Result<Vec<UInt8<P::Fp>>, SynthesisError> {
188        let mut bytes = Vec::new();
189        for coeffs in &self.ell_coeffs {
190            bytes.extend_from_slice(&coeffs.0.to_non_unique_bytes()?);
191            bytes.extend_from_slice(&coeffs.1.to_non_unique_bytes()?);
192        }
193        Ok(bytes)
194    }
195}
196
197impl<P: Bls12Config> G2PreparedVar<P> {
198    /// Constructs `Self` from a `G2Var`.
199    #[tracing::instrument(target = "r1cs")]
200    pub fn from_group_var(q: &G2Var<P>) -> Result<Self, SynthesisError> {
201        let q = q.to_affine()?;
202        let two_inv = P::Fp::one().double().inverse().unwrap();
203        // Enforce that `q` is not the point at infinity.
204        q.infinity.enforce_not_equal(&Boolean::TRUE)?;
205        let mut ell_coeffs = vec![];
206        let mut r = q.clone();
207
208        for i in BitIteratorBE::new(P::X).skip(1) {
209            ell_coeffs.push(Self::double(&mut r, &two_inv)?);
210
211            if i {
212                ell_coeffs.push(Self::add(&mut r, &q)?);
213            }
214        }
215
216        Ok(Self { ell_coeffs })
217    }
218
219    #[tracing::instrument(target = "r1cs")]
220    fn double(r: &mut G2AffineVar<P>, two_inv: &P::Fp) -> Result<LCoeff<P>, SynthesisError> {
221        let a = r.y.inverse()?;
222        let mut b = r.x.square()?;
223        let b_tmp = b.clone();
224        b.mul_assign_by_base_field_constant(*two_inv);
225        b += &b_tmp;
226
227        let c = &a * &b;
228        let d = r.x.double()?;
229        let x3 = c.square()? - &d;
230        let e = &c * &r.x - &r.y;
231        let c_x3 = &c * &x3;
232        let y3 = &e - &c_x3;
233        let mut f = c;
234        f.negate_in_place()?;
235        r.x = x3;
236        r.y = y3;
237        match P::TWIST_TYPE {
238            TwistType::M => Ok((e, f)),
239            TwistType::D => Ok((f, e)),
240        }
241    }
242
243    #[tracing::instrument(target = "r1cs")]
244    fn add(r: &mut G2AffineVar<P>, q: &G2AffineVar<P>) -> Result<LCoeff<P>, SynthesisError> {
245        let a = (&q.x - &r.x).inverse()?;
246        let b = &q.y - &r.y;
247        let c = &a * &b;
248        let d = &r.x + &q.x;
249        let x3 = c.square()? - &d;
250
251        let e = (&r.x - &x3) * &c;
252        let y3 = e - &r.y;
253        let g = &c * &r.x - &r.y;
254        let mut f = c;
255        f.negate_in_place()?;
256        r.x = x3;
257        r.y = y3;
258        match P::TWIST_TYPE {
259            TwistType::M => Ok((g, f)),
260            TwistType::D => Ok((f, g)),
261        }
262    }
263}