Implement sqrt_alt, a more efficient way of doing sqrt_ratio(num, one()).

Signed-off-by: Daira Hopwood <daira@jacaranda.org>

src/arithmetic/fields.rs

@@ -54,6 +54,11 @@ pub trait FieldExt:
     /// The choice of root from sqrt is unspecified.
     fn sqrt_ratio(num: &Self, div: &Self) -> (Choice, Self);
 
+    /// Equivalent to sqrt_ratio(self, one()).
+    fn sqrt_alt(&self) -> (Choice, Self) {
+        Self::sqrt_ratio(self, &Self::one())
+    }
+
     /// This computes a random element of the field using system randomness.
     fn rand() -> Self {
         Self::random(rand::rngs::OsRng)
@@ -259,13 +264,35 @@ impl<F: FieldExt> SqrtTables<F> {
         // uv = u * v
         let uv = w * num;
 
-        self.sqrt_common(num, div, &uv, &v)
+        let res = self.sqrt_common(&uv, &v);
+
+        let sqdiv = res.square() * div;
+        let is_square = (sqdiv - num).ct_is_zero();
+        let is_nonsquare = (sqdiv - F::ROOT_OF_UNITY * num).ct_is_zero();
+        assert!(bool::from(
+            num.ct_is_zero() | div.ct_is_zero() | (is_square ^ is_nonsquare)
+        ));
+
+        (is_square, res)
     }
 
-    /// Same as sqrt_ratio but given num, div, v = u^((T-1)/2), and uv = u * v as input.
-    ///
-    /// The choice of root from sqrt is unspecified.
-    fn sqrt_common(&self, num: &F, div: &F, uv: &F, v: &F) -> (Choice, F) {
+    /// Same as sqrt_ratio(u, one()) but more efficient.
+    pub fn sqrt_alt(&self, u: &F) -> (Choice, F) {
+        let v = u.pow_by_t_minus1_over2();
+        let uv = *u * v;
+
+        let res = self.sqrt_common(&uv, &v);
+
+        let sq = res.square();
+        let is_square = (sq - u).ct_is_zero();
+        let is_nonsquare = (sq - F::ROOT_OF_UNITY * u).ct_is_zero();
+        assert!(bool::from(u.ct_is_zero() | (is_square ^ is_nonsquare)));
+
+        (is_square, res)
+    }
+
+    /// Common part of sqrt_ratio and sqrt_alt: return res given v = u^((T-1)/2) and uv = u * v.
+    fn sqrt_common(&self, uv: &F, v: &F) -> F {
         let sqr = |x: F, i: u32| (0..i).fold(x, |x, _| x.square());
         let inv = |x: F| self.inv[self.hasher.hash(&x)] as usize;
 
@@ -296,20 +323,11 @@ impl<F: FieldExt> SqrtTables<F> {
                                 // 1 == x3 * ROOT_OF_UNITY^t_
         t_ = (t_ + 1) >> 1;
         assert!(t_ <= 0x80000000);
-        let res = *uv
-            * self.g0[t_ & 0xFF]
+
+        *uv * self.g0[t_ & 0xFF]
             * self.g1[(t_ >> 8) & 0xFF]
             * self.g2[(t_ >> 16) & 0xFF]
-            * self.g3[t_ >> 24];
-
-        let sqdiv = res.square() * div;
-        let is_square = (sqdiv - num).ct_is_zero();
-        let is_nonsquare = (sqdiv - F::ROOT_OF_UNITY * num).ct_is_zero();
-        assert!(bool::from(
-            num.ct_is_zero() | div.ct_is_zero() | (is_square ^ is_nonsquare)
-        ));
-
-        (is_square, res)
+            * self.g3[t_ >> 24]
     }
 }
 

src/pasta/fields/fp.rs

@@ -679,6 +679,10 @@ impl FieldExt for Fp {
         FP_TABLES.sqrt_ratio(num, div)
     }
 
+    fn sqrt_alt(&self) -> (Choice, Self) {
+        FP_TABLES.sqrt_alt(self)
+    }
+
     fn ct_is_zero(&self) -> Choice {
         self.ct_eq(&Self::zero())
     }
@@ -834,15 +838,20 @@ fn test_sqrt() {
 }
 
 #[test]
-fn test_sqrt_ratio() {
+fn test_sqrt_ratio_and_alt() {
     // (true, sqrt(num/div)), if num and div are nonzero and num/div is a square in the field
     let num = (Fp::TWO_INV).square();
     let div = Fp::from_u64(25);
+    let div_inverse = div.invert().unwrap();
     let expected = Fp::TWO_INV * Fp::from_u64(5).invert().unwrap();
     let (is_square, v) = Fp::sqrt_ratio(&num, &div);
     assert!(bool::from(is_square));
     assert!(v == expected || (-v) == expected);
 
+    let (is_square_alt, v_alt) = Fp::sqrt_alt(&(num * div_inverse));
+    assert!(bool::from(is_square_alt));
+    assert!(v_alt == v);
+
     // (false, sqrt(ROOT_OF_UNITY * num/div)), if num and div are nonzero and num/div is a nonsquare in the field
     let num = num * Fp::ROOT_OF_UNITY;
     let expected = Fp::TWO_INV * Fp::ROOT_OF_UNITY * Fp::from_u64(5).invert().unwrap();
@@ -850,6 +859,10 @@ fn test_sqrt_ratio() {
     assert!(!bool::from(is_square));
     assert!(v == expected || (-v) == expected);
 
+    let (is_square_alt, v_alt) = Fp::sqrt_alt(&(num * div_inverse));
+    assert!(!bool::from(is_square_alt));
+    assert!(v_alt == v);
+
     // (true, 0), if num is zero
     let num = Fp::zero();
     let expected = Fp::zero();
@@ -857,6 +870,10 @@ fn test_sqrt_ratio() {
     assert!(bool::from(is_square));
     assert!(v == expected);
 
+    let (is_square_alt, v_alt) = Fp::sqrt_alt(&(num * div_inverse));
+    assert!(bool::from(is_square_alt));
+    assert!(v_alt == v);
+
     // (false, 0), if num is nonzero and div is zero
     let num = (Fp::TWO_INV).square();
     let div = Fp::zero();

src/pasta/fields/fq.rs

@@ -679,6 +679,10 @@ impl FieldExt for Fq {
         FQ_TABLES.sqrt_ratio(num, div)
     }
 
+    fn sqrt_alt(&self) -> (Choice, Self) {
+        FQ_TABLES.sqrt_alt(self)
+    }
+
     fn ct_is_zero(&self) -> Choice {
         self.ct_eq(&Self::zero())
     }
@@ -834,15 +838,20 @@ fn test_sqrt() {
 }
 
 #[test]
-fn test_sqrt_ratio() {
+fn test_sqrt_ratio_and_alt() {
     // (true, sqrt(num/div)), if num and div are nonzero and num/div is a square in the field
     let num = (Fq::TWO_INV).square();
     let div = Fq::from_u64(25);
+    let div_inverse = div.invert().unwrap();
     let expected = Fq::TWO_INV * Fq::from_u64(5).invert().unwrap();
     let (is_square, v) = Fq::sqrt_ratio(&num, &div);
     assert!(bool::from(is_square));
     assert!(v == expected || (-v) == expected);
 
+    let (is_square_alt, v_alt) = Fq::sqrt_alt(&(num * div_inverse));
+    assert!(bool::from(is_square_alt));
+    assert!(v_alt == v);
+
     // (false, sqrt(ROOT_OF_UNITY * num/div)), if num and div are nonzero and num/div is a nonsquare in the field
     let num = num * Fq::ROOT_OF_UNITY;
     let expected = Fq::TWO_INV * Fq::ROOT_OF_UNITY * Fq::from_u64(5).invert().unwrap();
@@ -850,6 +859,10 @@ fn test_sqrt_ratio() {
     assert!(!bool::from(is_square));
     assert!(v == expected || (-v) == expected);
 
+    let (is_square_alt, v_alt) = Fq::sqrt_alt(&(num * div_inverse));
+    assert!(!bool::from(is_square_alt));
+    assert!(v_alt == v);
+
     // (true, 0), if num is zero
     let num = Fq::zero();
     let expected = Fq::zero();
@@ -857,6 +870,10 @@ fn test_sqrt_ratio() {
     assert!(bool::from(is_square));
     assert!(v == expected);
 
+    let (is_square_alt, v_alt) = Fq::sqrt_alt(&(num * div_inverse));
+    assert!(bool::from(is_square_alt));
+    assert!(v_alt == v);
+
     // (false, 0), if num is nonzero and div is zero
     let num = (Fq::TWO_INV).square();
     let div = Fq::zero();