//! 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(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(),

    )

}

Unexecuted instantiation: ff::helpers::sqrt_ratio_generic::<p256::arithmetic::field::FieldElement>

Unexecuted instantiation: ff::helpers::sqrt_ratio_generic::<p256::arithmetic::scalar::Scalar>