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