Add sqrt_ratio implementation.

Co-authored-by: Jack Grigg <jack@electriccoin.co>
Signed-off-by: Daira Hopwood <daira@jacaranda.org>

Cargo.toml

@@ -43,6 +43,8 @@ metrics-macros = "=0.1.0-alpha.9"
 num_cpus = "1.13"
 rand = "0.7"
 blake2b_simd = "0.5"
+lazy_static = "1.4.0"
+static_assertions = "1.1.0"
 
 [features]
 sanity-checks = []

src/arithmetic/fields.rs

@@ -1,10 +1,16 @@
 //! This module contains the `Field` abstraction that allows us to write
 //! code that generalizes over a pair of fields.
 
+use core::mem::size_of;
+use static_assertions::const_assert;
+use std::assert;
+use std::convert::TryInto;
 use subtle::{Choice, ConstantTimeEq, CtOption};
 
 use super::Group;
 
+const_assert!(size_of::<usize>() >= 4);
+
 /// This trait is a common interface for dealing with elements of a finite
 /// field.
 pub trait FieldExt:
@@ -16,6 +22,9 @@ pub trait FieldExt:
     /// Inverse of `ROOT_OF_UNITY`
     const ROOT_OF_UNITY_INV: Self;
 
+    /// The value $(T-1)/2$ such that $2^S \cdot T = p - 1$ with $T$ odd.
+    const T_MINUS1_OVER2: [u64; 4];
+
     /// Generator of the $t-order$ multiplicative subgroup
     const DELTA: Self;
 
@@ -32,6 +41,15 @@ pub trait FieldExt:
     /// Element of multiplicative order $3$.
     const ZETA: Self;
 
+    /// XOR parameter of the perfect hash function used for SqrtTables.
+    const HASH_XOR: u32;
+
+    /// Modulus of the perfect hash function used for SqrtTables.
+    const HASH_MOD: usize;
+
+    /// Tables for square root computation.
+    fn get_tables() -> &'static SqrtTables<Self>;
+
     /// This computes a random element of the field using system randomness.
     fn rand() -> Self {
         Self::random(rand::rngs::OsRng)
@@ -58,6 +76,119 @@ pub trait FieldExt:
     /// byte representation of an integer.
     fn from_bytes_wide(bytes: &[u8; 64]) -> Self;
 
+    /// Computes:
+    ///
+    /// * (true,  sqrt(num/div)),                 if num and div are nonzero and num/div is a square in the field;
+    /// * (true,  0),                             if num is zero;
+    /// * (false, 0),                             if num is nonzero and div is zero;
+    /// * (false, sqrt(ROOT_OF_UNITY * num/div)), if num and div are nonzero and num/div is a nonsquare in the field;
+    ///
+    /// where ROOT_OF_UNITY is a generator of the order 2^n subgroup (and therefore a nonsquare).
+    ///
+    /// The choice of root from sqrt is unspecified.
+    fn sqrt_ratio(num: &Self, div: &Self) -> (Choice, Self) {
+        // Based on:
+        // * [Sarkar2020](https://eprint.iacr.org/2020/1407)
+        // * [BDLSY2012](https://cr.yp.to/papers.html#ed25519)
+        //
+        // We need to calculate uv and v, where v = u^((m-1)/2), u = num/div, and p-1 = T * 2^S.
+        // We can rewrite as follows:
+        //
+        //      v = (num/div)^((T-1)/2)
+        //        = num^((T-1)/2) * div^(p-1 - (T-1)/2)    [Fermat's Little Theorem]
+        //        =       "       * div^(T * 2^S - (T-1)/2)
+        //        =       "       * div^((2^(S+1) - 1)*(T-1)/2 + 2^S)
+        //        = (num * div^(2^(S+1) - 1))^((T-1)/2) * div^(2^S)
+        //
+        // Let  w = (num * div^(2^(S+1) - 1))^((T-1)/2) * div^(2^S - 1).
+        // Then v = w * div, and uv = num * v / div = num * w.
+        //
+        // We calculate:
+        //
+        //      s = div^(2^S - 1) using an addition chain
+        //      t = div^(2^(S+1) - 1) = s^2 * div
+        //      w = (num * t)^((T-1)/2) * s using another addition chain
+        //
+        // then u and uv as above. The addition chains are given in
+        // https://github.com/zcash/pasta/blob/master/addchain_sqrt.py .
+        // The overall cost of this part is similar to a single full-width exponentiation,
+        // regardless of S.
+
+        let sqr = |x: Self, i: u32| (0..i).fold(x, |x, _| x.square());
+
+        // s = div^(2^S - 1)
+        let s = (0..5).fold(*div, |d: Self, i| sqr(d, 1 << i) * d);
+
+        // t == div^(2^(S+1) - 1)
+        let t = s.square() * div;
+
+        // TODO: replace this with an addition chain.
+        let w = ff::Field::pow_vartime(&(t * num), &Self::T_MINUS1_OVER2) * s;
+
+        // v == u^((T-1)/2)
+        let v = w * div;
+
+        // uv = u * v
+        let uv = w * num;
+
+        Self::sqrt_common(num, div, &uv, &v)
+    }
+
+    /// 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(num: &Self, div: &Self, uv: &Self, v: &Self) -> (Choice, Self) {
+        let tab = Self::get_tables();
+        let sqr = |x: Self, i: u32| (0..i).fold(x, |x, _| x.square());
+
+        let x3 = *uv * v;
+        let x2 = sqr(x3, 8);
+        let x1 = sqr(x2, 8);
+        let x0 = sqr(x1, 8);
+
+        // i = 0, 1
+        let mut t_: usize = tab.inv[x0.hash()] as usize; // = t >> 16
+                                                         // 1 == x0 * ROOT_OF_UNITY^(t_ << 24)
+        assert!(t_ < 0x100);
+        let alpha = x1 * tab.g2[t_];
+
+        // i = 2
+        t_ += (tab.inv[alpha.hash()] as usize) << 8; // = t >> 8
+                                                     // 1 == x1 * ROOT_OF_UNITY^(t_ << 16)
+        assert!(t_ < 0x10000);
+        let alpha = x2 * tab.g1[t_ & 0xFF] * tab.g2[t_ >> 8];
+
+        // i = 3
+        t_ += (tab.inv[alpha.hash()] as usize) << 16; // = t
+                                                      // 1 == x2 * ROOT_OF_UNITY^(t_ << 8)
+        assert!(t_ < 0x1000000);
+        let alpha = x3 * tab.g0[t_ & 0xFF] * tab.g1[(t_ >> 8) & 0xFF] * tab.g2[t_ >> 16];
+
+        t_ += (tab.inv[alpha.hash()] as usize) << 24; // = t << 1
+                                                      // 1 == x3 * ROOT_OF_UNITY^t_
+        t_ = (t_ + 1) >> 1;
+        assert!(t_ <= 0x80000000);
+        let res = *uv
+            * tab.g0[t_ & 0xFF]
+            * tab.g1[(t_ >> 8) & 0xFF]
+            * tab.g2[(t_ >> 16) & 0xFF]
+            * tab.g3[t_ >> 24];
+
+        let sqdiv = res.square() * div;
+        let is_square = (sqdiv - num).ct_is_zero();
+        let is_nonsquare = (sqdiv - Self::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)
+    }
+
+    /// Returns a perfect hash of this element for use with inv.
+    fn hash(&self) -> usize {
+        ((self.get_lower_32() ^ Self::HASH_XOR) as usize) % Self::HASH_MOD
+    }
+
     /// Exponentiates `self` by `by`, where `by` is a little-endian order
     /// integer exponent.
     fn pow(&self, by: &[u64; 4]) -> Self {
@@ -77,6 +208,10 @@ pub trait FieldExt:
     /// canonically.
     fn get_lower_128(&self) -> u128;
 
+    /// Gets the lower 32 bits of this field element when expressed
+    /// canonically.
+    fn get_lower_32(&self) -> u32;
+
     /// Performs a batch inversion using Montgomery's trick, returns the product
     /// of every inverse. Zero inputs are ignored.
     fn batch_invert(v: &mut [Self]) -> Self {
@@ -103,6 +238,53 @@ pub trait FieldExt:
     }
 }
 
+/// Tables used for square root computation.
+#[derive(Debug)]
+pub struct SqrtTables<F: FieldExt> {
+    inv: Vec<u8>,
+    g0: [F; 256],
+    g1: [F; 256],
+    g2: [F; 256],
+    g3: [F; 129],
+}
+
+impl<F: FieldExt> SqrtTables<F> {
+    /// Build tables given parameters for the perfect hash.
+    pub fn init() -> Self {
+        let gtab: Vec<Vec<F>> = (0..4)
+            .scan(F::ROOT_OF_UNITY, |gi, _| {
+                // gi == ROOT_OF_UNITY^(256^i)
+                let gtab_i: Vec<F> = (0..256)
+                    .scan(F::one(), |acc, _| {
+                        let res = *acc;
+                        *acc *= *gi;
+                        Some(res)
+                    })
+                    .collect();
+                *gi = gtab_i[255] * *gi;
+                Some(gtab_i)
+            })
+            .collect();
+
+        // Now invert gtab[3].
+        let mut inv: Vec<u8> = vec![1; F::HASH_MOD];
+        for j in 0..256 {
+            let hash = gtab[3][j].hash();
+            // 1 is the last value to be assigned, so this ensures there are no collisions.
+            assert!(inv[hash] == 1);
+            inv[hash] = ((256 - j) & 0xFF) as u8;
+        }
+
+        SqrtTables::<F> {
+            inv,
+            g0: gtab[0][..].try_into().unwrap(),
+            g1: gtab[1][..].try_into().unwrap(),
+            g2: gtab[2][..].try_into().unwrap(),
+            g3: gtab[3][0..129].try_into().unwrap(),
+        }
+    }
+}
+
 /// Compute a + b + carry, returning the result and the new carry over.
 #[inline(always)]
 pub(crate) const fn adc(a: u64, b: u64, carry: u64) -> (u64, u64) {

src/pasta/fields/fp.rs

@@ -2,10 +2,11 @@ use bitvec::{array::BitArray, order::Lsb0};
 use core::convert::TryInto;
 use core::fmt;
 use core::ops::{Add, Mul, Neg, Sub};
+use lazy_static::lazy_static;
 use rand::RngCore;
 use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
 
-use crate::arithmetic::{adc, mac, sbb, FieldExt, Group};
+use crate::arithmetic::{adc, mac, sbb, FieldExt, Group, SqrtTables};
 
 /// This represents an element of $\mathbb{F}_p$ where
 ///
@@ -643,6 +644,12 @@ impl FieldExt for Fp {
         0xb4ed8e647196dad1,
         0x2cd5282c53116b5c,
     ]);
+    const T_MINUS1_OVER2: [u64; 4] = [
+        0x04a67c8dcc969876,
+        0x0000000011234c7e,
+        0x0000000000000000,
+        0x20000000,
+    ];
     const DELTA: Self = DELTA;
     const TWO_INV: Self = Fp::from_raw([
         0xcc96987680000001,
@@ -664,6 +671,16 @@ impl FieldExt for Fp {
         0x12ccca834acdba71,
     ]);
 
+    const HASH_XOR: u32 = 0x11BE;
+    const HASH_MOD: usize = 1098;
+
+    fn get_tables() -> &'static SqrtTables<Self> {
+        lazy_static! {
+            static ref FP_TABLES: SqrtTables<Fp> = SqrtTables::init();
+        }
+        &FP_TABLES
+    }
+
     fn ct_is_zero(&self) -> Choice {
         self.ct_eq(&Self::zero())
     }
@@ -740,6 +757,13 @@ impl FieldExt for Fp {
 
         u128::from(tmp.0[0]) | (u128::from(tmp.0[1]) << 64)
     }
+
+    fn get_lower_32(&self) -> u32 {
+        // TODO: don't reduce, just hash the Montgomery form. (Requires rebuilding perfect hash table.)
+        let tmp = Fp::montgomery_reduce(self.0[0], self.0[1], self.0[2], self.0[3], 0, 0, 0, 0);
+
+        tmp.0[0] as u32
+    }
 }
 
 #[cfg(test)]
@@ -778,6 +802,39 @@ fn test_sqrt() {
     assert!(v == Fp::TWO_INV || (-v) == Fp::TWO_INV);
 }
 
+#[test]
+fn test_sqrt_ratio() {
+    // (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 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);
+
+    // (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();
+    let (is_square, v) = Fp::sqrt_ratio(&num, &div);
+    assert!(!bool::from(is_square));
+    assert!(v == expected || (-v) == expected);
+
+    // (true, 0), if num is zero
+    let num = Fp::zero();
+    let expected = Fp::zero();
+    let (is_square, v) = Fp::sqrt_ratio(&num, &div);
+    assert!(bool::from(is_square));
+    assert!(v == expected);
+
+    // (false, 0), if num is nonzero and div is zero
+    let num = (Fp::TWO_INV).square();
+    let div = Fp::zero();
+    let expected = Fp::zero();
+    let (is_square, v) = Fp::sqrt_ratio(&num, &div);
+    assert!(!bool::from(is_square));
+    assert!(v == expected);
+}
+
 #[test]
 fn test_zeta() {
     assert_eq!(

src/pasta/fields/fq.rs

@@ -2,10 +2,11 @@ use bitvec::{array::BitArray, order::Lsb0};
 use core::convert::TryInto;
 use core::fmt;
 use core::ops::{Add, Mul, Neg, Sub};
+use lazy_static::lazy_static;
 use rand::RngCore;
 use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
 
-use crate::arithmetic::{adc, mac, sbb, FieldExt, Group};
+use crate::arithmetic::{adc, mac, sbb, FieldExt, Group, SqrtTables};
 
 /// This represents an element of $\mathbb{F}_q$ where
 ///
@@ -643,6 +644,12 @@ impl FieldExt for Fq {
         0xf4c8f353124086c1,
         0x2235e1a7415bf936,
     ]);
+    const T_MINUS1_OVER2: [u64; 4] = [
+        0x04ca546ec6237590,
+        0x0000000011234c7e,
+        0x0000000000000000,
+        0x20000000,
+    ];
     const DELTA: Self = DELTA;
     const TWO_INV: Self = Fq::from_raw([
         0xc623759080000001,
@@ -664,6 +671,16 @@ impl FieldExt for Fq {
         0x06819a58283e528e,
     ]);
 
+    const HASH_XOR: u32 = 0x116A9E;
+    const HASH_MOD: usize = 1206;
+
+    fn get_tables() -> &'static SqrtTables<Self> {
+        lazy_static! {
+            static ref FQ_TABLES: SqrtTables<Fq> = SqrtTables::init();
+        }
+        &FQ_TABLES
+    }
+
     fn ct_is_zero(&self) -> Choice {
         self.ct_eq(&Self::zero())
     }
@@ -740,6 +757,13 @@ impl FieldExt for Fq {
 
         u128::from(tmp.0[0]) | (u128::from(tmp.0[1]) << 64)
     }
+
+    fn get_lower_32(&self) -> u32 {
+        // TODO: don't reduce, just hash the Montgomery form. (Requires rebuilding perfect hash table.)
+        let tmp = Fq::montgomery_reduce(self.0[0], self.0[1], self.0[2], self.0[3], 0, 0, 0, 0);
+
+        tmp.0[0] as u32
+    }
 }
 
 #[cfg(test)]
@@ -778,6 +802,39 @@ fn test_sqrt() {
     assert!(v == Fq::TWO_INV || (-v) == Fq::TWO_INV);
 }
 
+#[test]
+fn test_sqrt_ratio() {
+    // (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 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);
+
+    // (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();
+    let (is_square, v) = Fq::sqrt_ratio(&num, &div);
+    assert!(!bool::from(is_square));
+    assert!(v == expected || (-v) == expected);
+
+    // (true, 0), if num is zero
+    let num = Fq::zero();
+    let expected = Fq::zero();
+    let (is_square, v) = Fq::sqrt_ratio(&num, &div);
+    assert!(bool::from(is_square));
+    assert!(v == expected);
+
+    // (false, 0), if num is nonzero and div is zero
+    let num = (Fq::TWO_INV).square();
+    let div = Fq::zero();
+    let expected = Fq::zero();
+    let (is_square, v) = Fq::sqrt_ratio(&num, &div);
+    assert!(!bool::from(is_square));
+    assert!(v == expected);
+}
+
 #[test]
 fn test_zeta() {
     assert_eq!(