p521: Improve field arithmetic (#948)

- sqn: Has an extra copy
- sqrt: Can have simplified math

Signed-off-by: Arvind Mukund <[email protected]>

p521/src/arithmetic/field.rs

@@ -201,8 +201,8 @@ impl FieldElement {
 
     /// Returns self^(2^n) mod p
     const fn sqn(&self, n: usize) -> Self {
-        let mut x = *self;
-        let mut i = 0;
+        let mut x = self.square();
+        let mut i = 1;
         while i < n {
             x = x.square();
             i += 1;
@@ -276,22 +276,18 @@ impl FieldElement {
 
     /// Returns the square root of self mod p, or `None` if no square root
     /// exists.
+    ///
+    /// # Implementation details
+    /// If _x_ has a sqrt, then due to Euler's criterion this implies x<sup>(p - 1)/2</sup> = 1.
+    /// 1. x<sup>(p + 1)/2</sup> = x.
+    /// 2. There's a special property due to _p ≡ 3 (mod 4)_ which implies _(p + 1)/4_ is an integer.
+    /// 3. We can rewrite `1.` as x<sup>((p+1)/4)<sup>2</sup></sup>
+    /// 4. x<sup>(p+1)/4</sup> is the square root.
+    /// 5. This is simplified as (2<sup>251</sup> - 1 + 1) /4 = 2<sup>519</sup>
+    /// 6. Hence, x<sup>2<sup>519</sup></sup> is the square root iff _result.square() == self_
     pub fn sqrt(&self) -> CtOption<Self> {
-        // Tonelli-Shank's algorithm for q mod 4 = 3 (i.e. Shank's algorithm)
-        // https://eprint.iacr.org/2012/685.pdf
-        let w = self.pow_vartime(&[
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000000,
-            0x0000000000000080,
-        ]);
-
-        CtOption::new(w, w.square().ct_eq(self))
+        let sqrt = self.sqn(519);
+        CtOption::new(sqrt, sqrt.square().ct_eq(self))
     }
 
     /// Relax a tight field element into a loose one.