Remove legendre computation from fp2's sqrt, fix bug in fp sqrt

CHANGELOG.md

@@ -37,7 +37,7 @@
 - #122 (ark-poly) Add infrastructure for benchmarking `FFT`s.
 - #125 (ark-poly) Add parallelization to applying coset shifts within `coset_fft`.
 - #126 (ark-ec) Use `ark_ff::batch_inversion` for point normalization
-- #131 (ark-ff) Speedup `sqrt` on `PrimeField` when a square root exists. (And slows it down when doesn't exist)
+- #131, #137 (ark-ff) Speedup `sqrt` on fields when a square root exists. (And slows it down when doesn't exist)
 
 ### Bug fixes
 - #36 (ark-ec) In Short-Weierstrass curves, include an infinity bit in `ToConstraintField`.

ff/src/fields/arithmetic.rs

@@ -240,17 +240,6 @@ macro_rules! sqrt_impl {
         let mut b = x * &w;
 
         let mut v = $P::TWO_ADICITY as usize;
-        // t = self^t
-        #[cfg(debug_assertions)]
-        {
-            let mut check = b;
-            for _ in 0..(v - 1) {
-                check.square_in_place();
-            }
-            if !check.is_one() {
-                panic!("Input is not a square root, but it passed the QR test")
-            }
-        }
 
         while !b.is_one() {
             let mut k = 0usize;

ff/src/fields/models/quadratic_extension.rs

@@ -278,34 +278,40 @@ where
     fn sqrt(&self) -> Option<Self> {
         // Square root based on the complex method. See
         // https://eprint.iacr.org/2012/685.pdf (page 15, algorithm 8)
-        use crate::LegendreSymbol::*;
         if self.c1.is_zero() {
             return self.c0.sqrt().map(|c0| Self::new(c0, P::BaseField::zero()));
         }
+        // Try computing the square root
+        // Check at the end of the algorithm if it was a square root
         let alpha = self.norm();
-        // This is equivalent to self.legendre, see the comment on legendre above.
-        let legendre_symbol = alpha.legendre();
-        match legendre_symbol {
-            Zero => Some(*self),
-            QuadraticNonResidue => None,
-            QuadraticResidue => {
-                // TODO: Precompute this
-                let two_inv = P::BaseField::one()
-                    .double()
-                    .inverse()
-                    .expect("Two should always have an inverse");
-                let alpha = alpha
-                    .sqrt()
-                    .expect("We are in the QR case, the norm should have a square root");
-                let mut delta = (alpha + &self.c0) * &two_inv;
-                if delta.legendre().is_qnr() {
-                    delta -= &alpha;
-                }
-                let c0 = delta.sqrt().expect("Delta must have a square root");
-                let c0_inv = c0.inverse().expect("c0 must have an inverse");
-                Some(Self::new(c0, self.c1 * &two_inv * &c0_inv))
+        // TODO: Precompute this
+        let two_inv = P::BaseField::one()
+            .double()
+            .inverse()
+            .expect("Two should always have an inverse");
+        let alpha = alpha
+            .sqrt()
+            .expect("We are in the QR case, the norm should have a square root");
+        let mut delta = (alpha + &self.c0) * &two_inv;
+        if delta.legendre().is_qnr() {
+            delta -= &alpha;
+        }
+        let c0 = delta.sqrt().expect("Delta must have a square root");
+        let c0_inv = c0.inverse().expect("c0 must have an inverse");
+        let sqrt_cand = Self::new(c0, self.c1 * &two_inv * &c0_inv);
+        // Check if sqrt_cand is actually the square root
+        // if not, there exists no square root.
+        if sqrt_cand.square() == *self {
+            return Some(sqrt_cand);
+        }
+        #[cfg(debug_assertions)]
+        {
+            use crate::fields::LegendreSymbol::*;
+            if self.legendre() != QuadraticNonResidue {
+                panic!("Input has a square root per its legendre symbol, but it was not found")
             }
         }
+        None
     }
 
     fn sqrt_in_place(&mut self) -> Option<&mut Self> {