From 2d3f498e751dc4fda452994f6f5c64eaec2adc94 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Michele=20Orr=C3=B9?= <maker@tumbolandia.net>
Date: Thu, 10 Aug 2017 18:18:46 +0200
Subject: [PATCH] Polish sqrt in fr.rs: use pattern matching with Legendre
 enums.

---
 src/bls12_381/fr.rs | 60 ++++++++++++++++++++++-----------------------
 1 file changed, 29 insertions(+), 31 deletions(-)

diff --git a/src/bls12_381/fr.rs b/src/bls12_381/fr.rs
index 7842de592..f91494210 100644
--- a/src/bls12_381/fr.rs
+++ b/src/bls12_381/fr.rs
@@ -555,45 +555,42 @@ impl SqrtField for Fr {
     fn sqrt(&self) -> Option<Self> {
         // Tonelli-Shank's algorithm for q mod 16 = 1
         // https://eprint.iacr.org/2012/685.pdf (page 12, algorithm 5)
+        match self.legendre() {
+            Zero => Some(*self),
+            QNonResidue => None,
+            QResidue => {
+                let mut c = Fr(ROOT_OF_UNITY);
+                // r = self^((t + 1) // 2)
+                let mut r = self.pow([0x7fff2dff80000000, 0x4d0ec02a9ded201, 0x94cebea4199cec04, 0x39f6d3a9]);
+                // t = self^t
+                let mut t = self.pow([0xfffe5bfeffffffff, 0x9a1d80553bda402, 0x299d7d483339d808, 0x73eda753]);
+                let mut m = S;
 
-        if self.is_zero() {
-            return Some(*self);
-        }
-
-        if let QNonResidue = self.legendre() {
-            None
-        } else {
-            let mut c = Fr(ROOT_OF_UNITY);
-            // r = self^((t + 1) // 2)
-            let mut r = self.pow([0x7fff2dff80000000, 0x4d0ec02a9ded201, 0x94cebea4199cec04, 0x39f6d3a9]);
-            // t = self^t
-            let mut t = self.pow([0xfffe5bfeffffffff, 0x9a1d80553bda402, 0x299d7d483339d808, 0x73eda753]);
-            let mut m = S;
-
-            while t != Self::one() {
+                while t != Self::one() {
                 let mut i = 1;
-                {
-                    let mut t2i = t;
-                    t2i.square();
-                    loop {
-                        if t2i == Self::one() {
-                            break;
-                        }
+                    {
+                        let mut t2i = t;
                         t2i.square();
-                        i += 1;
+                        loop {
+                            if t2i == Self::one() {
+                                break;
+                            }
+                            t2i.square();
+                            i += 1;
+                        }
                     }
-                }
 
-                for _ in 0..(m - i - 1) {
+                    for _ in 0..(m - i - 1) {
+                        c.square();
+                    }
+                    r.mul_assign(&c);
                     c.square();
+                    t.mul_assign(&c);
+                    m = i;
                 }
-                r.mul_assign(&c);
-                c.square();
-                t.mul_assign(&c);
-                m = i;
-            }
 
-            Some(r)
+                Some(r)
+            }
         }
     }
 }
@@ -606,6 +603,7 @@ impl LegendreField for Fr {
         else { QNonResidue }
     }
 }
+
 #[cfg(test)]
 use rand::{SeedableRng, XorShiftRng, Rand};