diff --git a/CHANGELOG.md b/CHANGELOG.md index b26567698..8087c0027 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ - #122 (ark-poly) Add infrastructure for benchmarking `FFT`s. - #125 (ark-poly) Add parallelization to applying coset shifts within `coset_fft`. - #126 (ark-ec) Use `ark_ff::batch_inversion` for point normalization +- #131 (ark-ff) Speedup `sqrt` on `PrimeField` when a square root exists. (And slows it down when doesn't exist) ### Bug fixes - #36 (ark-ec) In Short-Weierstrass curves, include an infinity bit in `ToConstraintField`. diff --git a/ff/src/fields/arithmetic.rs b/ff/src/fields/arithmetic.rs index 7b6208d96..0ef8e848c 100644 --- a/ff/src/fields/arithmetic.rs +++ b/ff/src/fields/arithmetic.rs @@ -229,53 +229,64 @@ macro_rules! sqrt_impl { // Actually this is just normal Tonelli-Shanks; since `P::Generator` // is a quadratic non-residue, `P::ROOT_OF_UNITY = P::GENERATOR ^ t` // is also a quadratic non-residue (since `t` is odd). - match $self.legendre() { - Zero => Some(*$self), - QuadraticNonResidue => None, - QuadraticResidue => { - let mut z = $Self::qnr_to_t(); - let mut w = $self.pow($P::T_MINUS_ONE_DIV_TWO); - let mut x = w * $self; - let mut b = x * &w; - - let mut v = $P::TWO_ADICITY as usize; - // t = self^t - #[cfg(debug_assertions)] - { - let mut check = b; - for _ in 0..(v - 1) { - check.square_in_place(); - } - if !check.is_one() { - panic!("Input is not a square root, but it passed the QR test") - } - } - - while !b.is_one() { - let mut k = 0usize; + if $self.is_zero() { + return Some($Self::zero()); + } + // Try computing the square root (x at the end of the algorithm) + // Check at the end of the algorithm if x was a square root + // Begin Tonelli-Shanks + let mut z = $Self::qnr_to_t(); + let mut w = $self.pow($P::T_MINUS_ONE_DIV_TWO); + let mut x = w * $self; + let mut b = x * &w; + + let mut v = $P::TWO_ADICITY as usize; + // t = self^t + #[cfg(debug_assertions)] + { + let mut check = b; + for _ in 0..(v - 1) { + check.square_in_place(); + } + if !check.is_one() { + panic!("Input is not a square root, but it passed the QR test") + } + } - let mut b2k = b; - while !b2k.is_one() { - // invariant: b2k = b^(2^k) after entering this loop - b2k.square_in_place(); - k += 1; - } + while !b.is_one() { + let mut k = 0usize; - let j = v - k - 1; - w = z; - for _ in 0..j { - w.square_in_place(); - } + let mut b2k = b; + while !b2k.is_one() { + // invariant: b2k = b^(2^k) after entering this loop + b2k.square_in_place(); + k += 1; + } - z = w.square(); - b *= &z; - x *= &w; - v = k; - } + let j = v - k - 1; + w = z; + for _ in 0..j { + w.square_in_place(); + } - Some(x) + z = w.square(); + b *= &z; + x *= &w; + v = k; + } + // Is x the square root? If so, return it. + if (x * x == *$self) { + return Some(x); + } + // Consistency check that if no square root is found, + // it is because none exists. + #[cfg(debug_assertions)] + { + if ($self.legendre() != QuadraticNonResidue) { + panic!("Input has a square root per its legendre symbol, but it was not found") } } + None }}; }