Add faster subgroup checks for BN254 (#808)

Co-authored-by: Marcin <[email protected]>

curves/bn254/src/curves/g1.rs

@@ -47,6 +47,12 @@ impl SWCurveConfig for Config {
         let s = Self::ScalarField::from_sign_and_limbs(true, scalar);
         GLVConfig::glv_mul_projective(*p, s)
     }
+
+    #[inline]
+    fn is_in_correct_subgroup_assuming_on_curve(_p: &G1Affine) -> bool {
+        // G1 = E(Fq) so if the point is on the curve, it is also in the subgroup.
+        true
+    }
 }
 
 impl GLVConfig for Config {

curves/bn254/src/curves/g2.rs

@@ -1,9 +1,10 @@
+use ark_ec::AffineRepr;
 use ark_ec::{
     models::{short_weierstrass::SWCurveConfig, CurveConfig},
     scalar_mul::glv::GLVConfig,
     short_weierstrass::{Affine, Projective},
 };
-use ark_ff::{AdditiveGroup, BigInt, MontFp, PrimeField, Zero};
+use ark_ff::{AdditiveGroup, BigInt, Field, MontFp, PrimeField, Zero};
 
 use crate::{Fq, Fq2, Fr};
 
@@ -49,6 +50,16 @@ impl SWCurveConfig for Config {
     fn mul_by_a(_: Self::BaseField) -> Self::BaseField {
         Self::BaseField::zero()
     }
+
+    fn is_in_correct_subgroup_assuming_on_curve(point: &G2Affine) -> bool {
+        // Subgroup check from section 4.3 of https://eprint.iacr.org/2022/352.pdf.
+        //
+        // Checks that [p]P = [6X^2]P
+
+        let x_times_point = point.mul_bigint(SIX_X_SQUARED);
+        let p_times_point = p_power_endomorphism(point);
+        x_times_point.eq(&p_times_point)
+    }
 }
 
 impl GLVConfig for Config {
@@ -102,3 +113,79 @@ pub const G2_GENERATOR_Y_C0: Fq =
 /// 4082367875863433681332203403145435568316851327593401208105741076214120093531
 pub const G2_GENERATOR_Y_C1: Fq =
     MontFp!("4082367875863433681332203403145435568316851327593401208105741076214120093531");
+
+// PSI_X = (u+9)^((p-1)/3) = TWIST_MUL_BY_Q_X
+const P_POWER_ENDOMORPHISM_COEFF_0: Fq2 = Fq2::new(
+    MontFp!("21575463638280843010398324269430826099269044274347216827212613867836435027261"),
+    MontFp!("10307601595873709700152284273816112264069230130616436755625194854815875713954"),
+);
+
+// PSI_Y = (u+9)^((p-1)/2) = TWIST_MUL_BY_Q_Y
+const P_POWER_ENDOMORPHISM_COEFF_1: Fq2 = Fq2::new(
+    MontFp!("2821565182194536844548159561693502659359617185244120367078079554186484126554"),
+    MontFp!("3505843767911556378687030309984248845540243509899259641013678093033130930403"),
+);
+
+// Integer representation of 6x^2 = t - 1
+const SIX_X_SQUARED: [u64; 2] = [17887900258952609094, 8020209761171036667];
+
+/// psi(P) is the untwist-Frobenius-twist endomorphism on E'(Fq2)
+fn p_power_endomorphism(p: &Affine<Config>) -> Affine<Config> {
+    // Maps (x,y) -> (x^p * (u+9)^((p-1)/3), y^p * (u+9)^((p-1)/2))
+
+    let mut res = *p;
+    res.x.frobenius_map_in_place(1);
+    res.y.frobenius_map_in_place(1);
+
+    res.x *= P_POWER_ENDOMORPHISM_COEFF_0;
+    res.y *= P_POWER_ENDOMORPHISM_COEFF_1;
+
+    res
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+    use crate::g2;
+    use ark_std::{rand::Rng, UniformRand};
+
+    fn sample_unchecked() -> Affine<g2::Config> {
+        let mut rng = ark_std::test_rng();
+        loop {
+            let x1 = Fq::rand(&mut rng);
+            let x2 = Fq::rand(&mut rng);
+            let greatest = rng.gen();
+            let x = Fq2::new(x1, x2);
+
+            if let Some(p) = Affine::get_point_from_x_unchecked(x, greatest) {
+                return p;
+            }
+        }
+    }
+
+    fn naive_is_in_subgroup_assuming_on_curve(p: &Affine<g2::Config>) -> bool {
+        <g2::Config as SWCurveConfig>::mul_affine(
+            p,
+            <g2::Config as CurveConfig>::ScalarField::characteristic(),
+        )
+        .is_zero()
+    }
+
+    #[test]
+    fn test_is_in_subgroup_assuming_on_curve() {
+        const SAMPLES: usize = 100;
+        for _ in 0..SAMPLES {
+            let p: Affine<g2::Config> = sample_unchecked();
+            assert!(p.is_on_curve());
+
+            assert_eq!(
+                naive_is_in_subgroup_assuming_on_curve(&p),
+                p.is_in_correct_subgroup_assuming_on_curve()
+            );
+
+            let cleared = p.clear_cofactor();
+            assert!(cleared.is_in_correct_subgroup_assuming_on_curve());
+        }
+    }
+}