Add ff::Field::{sqrt_ratio, sqrt_alt}

We also provide helper methods for implementing the square root trait
methods.

Ported from zcash/pasta_curves@db83057.

CHANGELOG.md

@@ -8,6 +8,18 @@ and this library adheres to Rust's notion of
 ## [Unreleased]
 ### Added
 - `ff::Field::pow`
+- `ff::Field::{sqrt_ratio, sqrt_alt}`
+- `ff::helpers`:
+  - `sqrt_tonelli_shanks`
+  - `sqrt_ratio_generic`
+
+### Changed
+- `ff::Field::sqrt` is now a provided method that uses the `Field::sqrt_ratio`
+  method. Implementors of the `Field` trait can choose to implement
+  `Field::sqrt_ratio` and use the provided `ff::Field::sqrt` method, especially
+  if it is more efficient in practice, or they can keep their own implementation
+  of `Field::sqrt` and implement `Field::sqrt_ratio` in terms of that
+  implementation using the `ff::helpers::sqrt_ratio_generic` helper function.
 
 ## [0.12.1] - 2022-10-28
 ### Fixed

ff_derive/src/lib.rs

@@ -257,8 +257,8 @@ fn validate_struct(ast: &syn::DeriveInput, limbs: usize) -> Option<proc_macro2::
         syn::Expr::Group(expr_group) => match &*expr_group.expr {
             syn::Expr::Lit(expr_lit) => Some(&expr_lit.lit),
             _ => None,
-        }
-        _ => None
+        },
+        _ => None,
     };
     let lit_int = match match expr_lit {
         Some(syn::Lit::Int(lit_int)) => Some(lit_int),
@@ -1295,6 +1295,10 @@ fn prime_field_impl(
                 #squaring_impl
             }
 
+            fn sqrt_ratio(num: &Self, div: &Self) -> (::ff::derive::subtle::Choice, Self) {
+                ::ff::helpers::sqrt_ratio_generic(num, div)
+            }
+
             fn sqrt(&self) -> ::ff::derive::subtle::CtOption<Self> {
                 #sqrt_impl
             }

src/helpers.rs

@@ -0,0 +1,128 @@
+//! Helper methods for implementing the `ff` traits.
+
+use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
+
+use crate::PrimeField;
+
+/// Constant-time implementation of Tonelli–Shanks' square-root algorithm for
+/// `p mod 16 = 1`.
+///
+/// `tm1d2` should be set to `(t - 1) // 2`, where `t = (modulus - 1) >> F::S`.
+///
+/// ## Implementing [`Field::sqrt`]
+///
+/// This function can be used to implement [`Field::sqrt`] for fields that both implement
+/// [`PrimeField`] and satisfy `p mod 16 = 1`.
+///
+/// [`Field::sqrt`]: crate::Field::sqrt
+pub fn sqrt_tonelli_shanks<F: PrimeField, S: AsRef<[u64]>>(f: &F, tm1d2: S) -> CtOption<F> {
+    // This is a constant-time version of https://eprint.iacr.org/2012/685.pdf (page 12,
+    // algorithm 5). Steps 2-5 of the algorithm are omitted because they are only needed
+    // to detect non-square input; it is more efficient to do that by checking at the end
+    // whether the square of the result is the input.
+
+    // w = self^((t - 1) // 2)
+    let w = f.pow_vartime(tm1d2);
+
+    let mut v = F::S;
+    let mut x = w * f;
+    let mut b = x * w;
+
+    // Initialize z as the 2^S root of unity.
+    let mut z = F::root_of_unity();
+
+    for max_v in (1..=F::S).rev() {
+        let mut k = 1;
+        let mut b2k = b.square();
+        let mut j_less_than_v: Choice = 1.into();
+
+        // This loop has three phases based on the value of k for algorithm 5:
+        // - for j <= k, we square b2k in order to calculate b^{2^k}.
+        // - for k < j <= v, we square z in order to calculate ω.
+        // - for j > v, we do nothing.
+        for j in 2..max_v {
+            let b2k_is_one = b2k.ct_eq(&F::one());
+            let squared = F::conditional_select(&b2k, &z, b2k_is_one).square();
+            b2k = F::conditional_select(&squared, &b2k, b2k_is_one);
+            let new_z = F::conditional_select(&z, &squared, b2k_is_one);
+            j_less_than_v &= !j.ct_eq(&v);
+            k = u32::conditional_select(&j, &k, b2k_is_one);
+            z = F::conditional_select(&z, &new_z, j_less_than_v);
+        }
+
+        let result = x * z;
+        x = F::conditional_select(&result, &x, b.ct_eq(&F::one()));
+        z = z.square();
+        b *= z;
+        v = k;
+    }
+
+    CtOption::new(
+        x,
+        (x * x).ct_eq(f), // Only return Some if it's the square root.
+    )
+}
+
+/// Computes:
+///
+/// - $(\textsf{true}, \sqrt{\textsf{num}/\textsf{div}})$, if $\textsf{num}$ and
+///   $\textsf{div}$ are nonzero and $\textsf{num}/\textsf{div}$ is a square in the
+///   field;
+/// - $(\textsf{true}, 0)$, if $\textsf{num}$ is zero;
+/// - $(\textsf{false}, 0)$, if $\textsf{num}$ is nonzero and $\textsf{div}$ is zero;
+/// - $(\textsf{false}, \sqrt{G_S \cdot \textsf{num}/\textsf{div}})$, if
+///   $\textsf{num}$ and $\textsf{div}$ are nonzero and $\textsf{num}/\textsf{div}$ is
+///   a nonsquare in the field;
+///
+/// where $G_S$ is a non-square.
+///
+/// For this method, $G_S$ is currently [`PrimeField::root_of_unity`], a generator of the
+/// order $2^S$ subgroup. Users of this crate should not rely on this generator being
+/// fixed; it may be changed in future crate versions to simplify the implementation of
+/// the SSWU hash-to-curve algorithm.
+///
+/// The choice of root from sqrt is unspecified.
+///
+/// ## Implementing [`Field::sqrt_ratio`]
+///
+/// This function can be used to implement [`Field::sqrt_ratio`] for fields that also
+/// implement [`PrimeField`]. If doing so, the default implementation of [`Field::sqrt`]
+/// *MUST* be overridden, or else both functions will recurse in a cycle until a stack
+/// overflow occurs.
+///
+/// [`Field::sqrt_ratio`]: crate::Field::sqrt_ratio
+/// [`Field::sqrt`]: crate::Field::sqrt
+pub fn sqrt_ratio_generic<F: PrimeField>(num: &F, div: &F) -> (Choice, F) {
+    // General implementation:
+    //
+    // a = num * inv0(div)
+    //   = {    0    if div is zero
+    //     { num/div otherwise
+    //
+    // b = G_S * a
+    //   = {      0      if div is zero
+    //     { G_S*num/div otherwise
+    //
+    // Since G_S is non-square, a and b are either both zero (and both square), or
+    // only one of them is square. We can therefore choose the square root to return
+    // based on whether a is square, but for the boolean output we need to handle the
+    // num != 0 && div == 0 case specifically.
+
+    let a = div.invert().unwrap_or_else(F::zero) * num;
+    let b = a * F::root_of_unity();
+    let sqrt_a = a.sqrt();
+    let sqrt_b = b.sqrt();
+
+    let num_is_zero = num.is_zero();
+    let div_is_zero = div.is_zero();
+    let is_square = sqrt_a.is_some();
+    let is_nonsquare = sqrt_b.is_some();
+    assert!(bool::from(
+        num_is_zero | div_is_zero | (is_square ^ is_nonsquare)
+    ));
+
+    (
+        is_square & (num_is_zero | !div_is_zero),
+        CtOption::conditional_select(&sqrt_b, &sqrt_a, is_square).unwrap(),
+    )
+}

src/lib.rs

@@ -12,6 +12,8 @@ extern crate alloc;
 mod batch;
 pub use batch::*;
 
+pub mod helpers;
+
 #[cfg(feature = "derive")]
 #[cfg_attr(docsrs, doc(cfg(feature = "derive")))]
 pub use ff_derive::PrimeField;
@@ -101,9 +103,37 @@ pub trait Field:
     /// failing if the element is zero.
     fn invert(&self) -> CtOption<Self>;
 
+    /// Computes:
+    ///
+    /// - $(\textsf{true}, \sqrt{\textsf{num}/\textsf{div}})$, if $\textsf{num}$ and
+    ///   $\textsf{div}$ are nonzero and $\textsf{num}/\textsf{div}$ is a square in the
+    ///   field;
+    /// - $(\textsf{true}, 0)$, if $\textsf{num}$ is zero;
+    /// - $(\textsf{false}, 0)$, if $\textsf{num}$ is nonzero and $\textsf{div}$ is zero;
+    /// - $(\textsf{false}, \sqrt{G_S \cdot \textsf{num}/\textsf{div}})$, if
+    ///   $\textsf{num}$ and $\textsf{div}$ are nonzero and $\textsf{num}/\textsf{div}$ is
+    ///   a nonsquare in the field;
+    ///
+    /// where $G_S$ is a non-square.
+    ///
+    /// # Warnings
+    ///
+    /// - The choice of root from `sqrt` is unspecified.
+    /// - The value of $G_S$ is unspecified, and cannot be assumed to have any specific
+    ///   value in a generic context.
+    fn sqrt_ratio(num: &Self, div: &Self) -> (Choice, Self);
+
+    /// Equivalent to `Self::sqrt_ratio(self, one())`.
+    fn sqrt_alt(&self) -> (Choice, Self) {
+        Self::sqrt_ratio(self, &Self::one())
+    }
+
     /// Returns the square root of the field element, if it is
     /// quadratic residue.
-    fn sqrt(&self) -> CtOption<Self>;
+    fn sqrt(&self) -> CtOption<Self> {
+        let (is_square, res) = self.sqrt_alt();
+        CtOption::new(res, is_square)
+    }
 
     /// Exponentiates `self` by `exp`, where `exp` is a little-endian order integer
     /// exponent.