Speedup sqrt on PrimeField when a square root exists. (#131)

This comes at the expense of slowing sqrt down when it does not exist.
This came from no longer computing the legendre symbol before the
square root.

CHANGELOG.md

@@ -37,6 +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)
 
 ### Bug fixes
 - #36 (ark-ec) In Short-Weierstrass curves, include an infinity bit in `ToConstraintField`.

ff/src/fields/arithmetic.rs

@@ -229,53 +229,64 @@ macro_rules! sqrt_impl {
         // Actually this is just normal Tonelli-Shanks; since `P::Generator`
         // is a quadratic non-residue, `P::ROOT_OF_UNITY = P::GENERATOR ^ t`
         // is also a quadratic non-residue (since `t` is odd).
-        match $self.legendre() {
-            Zero => Some(*$self),
-            QuadraticNonResidue => None,
-            QuadraticResidue => {
-                let mut z = $Self::qnr_to_t();
-                let mut w = $self.pow($P::T_MINUS_ONE_DIV_TWO);
-                let mut x = w * $self;
-                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;
+        if $self.is_zero() {
+            return Some($Self::zero());
+        }
+        // Try computing the square root (x at the end of the algorithm)
+        // Check at the end of the algorithm if x was a square root
+        // Begin Tonelli-Shanks
+        let mut z = $Self::qnr_to_t();
+        let mut w = $self.pow($P::T_MINUS_ONE_DIV_TWO);
+        let mut x = w * $self;
+        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")
+            }
+        }
 
-                    let mut b2k = b;
-                    while !b2k.is_one() {
-                        // invariant: b2k = b^(2^k) after entering this loop
-                        b2k.square_in_place();
-                        k += 1;
-                    }
+        while !b.is_one() {
+            let mut k = 0usize;
 
-                    let j = v - k - 1;
-                    w = z;
-                    for _ in 0..j {
-                        w.square_in_place();
-                    }
+            let mut b2k = b;
+            while !b2k.is_one() {
+                // invariant: b2k = b^(2^k) after entering this loop
+                b2k.square_in_place();
+                k += 1;
+            }
 
-                    z = w.square();
-                    b *= &z;
-                    x *= &w;
-                    v = k;
-                }
+            let j = v - k - 1;
+            w = z;
+            for _ in 0..j {
+                w.square_in_place();
+            }
 
-                Some(x)
+            z = w.square();
+            b *= &z;
+            x *= &w;
+            v = k;
+        }
+        // Is x the square root? If so, return it.
+        if (x * x == *$self) {
+            return Some(x);
+        }
+        // Consistency check that if no square root is found,
+        // it is because none exists.
+        #[cfg(debug_assertions)]
+        {
+            if ($self.legendre() != QuadraticNonResidue) {
+                panic!("Input has a square root per its legendre symbol, but it was not found")
             }
         }
+        None
     }};
 }