//! Field arithmetic modulo p = 2^{224}(2^{32} − 1) + 2^{192} + 2^{96} − 1

#![allow(clippy::assign_op_pattern, clippy::op_ref)]

use elliptic_curve::{bigint::cpubits, ops::BatchInvert};

cpubits! {
    32 => {
        #[path = "field/field32.rs"]
        mod field_impl;
    }
    64 => {
        #[path = "field/field64.rs"]
        mod field_impl;
    }
}

use core::ops::Mul;
use elliptic_curve::{
    bigint::U256,
    ff::PrimeField,
    subtle::{Choice, ConstantTimeEq, CtOption},
};

#[cfg(doc)]
use {
    core::ops::{Add, Neg, Sub},
    elliptic_curve::{
        ff::{self, Field},
        subtle::ConditionallySelectable,
    },
};

/// Constant representing the modulus: p = 2^{224}(2^{32} − 1) + 2^{192} + 2^{96} − 1
const MODULUS_HEX: &str = "ffffffff00000001000000000000000000000000ffffffffffffffffffffffff";

primefield::monty_field_params!(
    name: FieldParams,
    modulus: MODULUS_HEX,
    uint: U256,
    byte_order: primefield::ByteOrder::BigEndian,
    multiplicative_generator: 6,
    doc: "Montgomery parameters for the NIST P-256 field modulus: p = 2^{224}(2^{32} − 1) + 2^{192} + 2^{96} − 1."
);

primefield::monty_field_element!(
    name: FieldElement,
    params: FieldParams,
    uint: U256,
    doc: "Element in the finite field modulo p = 2^{224}(2^{32} − 1) + 2^{192} + 2^{96} − 1."
);

impl FieldElement {
    /// Decode [`FieldElement`] from [`U256`] converting it into Montgomery form.
    ///
    /// Does *not* perform a check that the field element does not overflow the modulus.
    ///
    /// Used incorrectly this can lead to invalid results!
    pub(crate) const fn from_uint_unchecked(w: U256) -> Self {
        Self::multiply(
            &Self::from_montgomery(w),
            &Self::from_montgomery(*FieldParams::PARAMS.r2()),
        )
    }

    /// Returns self + rhs mod p
    pub const fn add(&self, rhs: &Self) -> Self {
        Self::from_montgomery(field_impl::add(
            self.0.as_montgomery(),
            rhs.0.as_montgomery(),
        ))
    }

    /// Returns 2 * self.
    pub const fn double(&self) -> Self {
        self.add(self)
    }

    /// Returns self - rhs mod p
    pub const fn sub(&self, rhs: &Self) -> Self {
        Self::from_montgomery(field_impl::sub(
            self.0.as_montgomery(),
            rhs.0.as_montgomery(),
        ))
    }

    /// Negate element.
    pub const fn neg(&self) -> Self {
        Self::sub(&Self::ZERO, self)
    }

    /// Translate a field element out of the Montgomery domain.
    #[inline]
    pub(crate) const fn to_canonical(self) -> U256 {
        field_impl::to_canonical(self.0.as_montgomery())
    }

    /// Returns self * rhs mod p
    pub const fn multiply(&self, rhs: &Self) -> Self {
        let (lo, hi): (U256, U256) = self.0.as_montgomery().widening_mul(rhs.0.as_montgomery());
        Self::from_montgomery(field_impl::montgomery_reduce(&lo, &hi))
    }

    /// Returns self * self mod p
    pub const fn square(&self) -> Self {
        // Schoolbook multiplication.
        self.multiply(self)
    }

    /// Returns the multiplicative inverse of self, if self is non-zero.
    pub fn invert(&self) -> CtOption<Self> {
        self.0.invert().map(Self)
    }

    /// Returns the multiplicative inverse of self, if self is non-zero.
    pub fn invert_vartime(&self) -> CtOption<Self> {
        self.0.invert_vartime().map(Self)
    }

    /// Returns the square root of self mod p, or `None` if no square root exists.
    pub fn sqrt(&self) -> CtOption<Self> {
        // We need to find alpha such that alpha^2 = beta mod p. For secp256r1,
        // p ≡ 3 mod 4. By Euler's Criterion, beta^(p-1)/2 ≡ 1 mod p. So:
        //
        //     alpha^2 = beta beta^((p - 1) / 2) mod p ≡ beta^((p + 1) / 2) mod p
        //     alpha = ± beta^((p + 1) / 4) mod p
        //
        // Thus sqrt can be implemented with a single exponentiation.

        let t11 = self.mul(&self.square());
        let t1111 = t11.mul(&t11.sqn(2));
        let t11111111 = t1111.mul(t1111.sqn(4));
        let x16 = t11111111.sqn(8).mul(t11111111);
        let sqrt = x16
            .sqn(16)
            .mul(x16)
            .sqn(32)
            .mul(self)
            .sqn(96)
            .mul(self)
            .sqn(94);

        CtOption::new(
            sqrt,
            (&sqrt * &sqrt).ct_eq(self), // Only return Some if it's the square root.
        )
    }

    /// Returns self^(2^n) mod p
    const fn sqn(&self, n: usize) -> Self {
        Self(self.0.sqn_vartime(n))
    }

    /// Construct a field element from a [`U256`] in Montgomery form.
    #[inline]
    pub(crate) const fn from_montgomery(uint: U256) -> Self {
        Self(primefield::MontyFieldElement::<
            FieldParams,
            { FieldParams::LIMBS },
        >::from_montgomery(uint))
    }
}

impl BatchInvert for FieldElement {}

#[cfg(test)]
mod tests {
    use super::{FieldElement, FieldParams, cpubits};
    use crate::{FieldBytes, U256, test_vectors::field::DBL_TEST_VECTORS};
    use elliptic_curve::{array::Array, bigint::modular::ConstMontyParams};

    cpubits! {
        64 => { use proptest::{num::u64::ANY, prelude::*}; }
    }

    primefield::test_primefield!(FieldElement, U256);

    /// Ensures the legacy `R2` constant is computed the same way as the `crypto-bigint`
    /// implementation.
    // TODO(tarcieri): since we know this works, it can probably be removed in future refactoring
    #[test]
    fn r2() {
        let expected_r2 =
            U256::from_be_hex("00000004fffffffdfffffffffffffffefffffffbffffffff0000000000000003");
        assert_eq!(FieldParams::PARAMS.r2(), &expected_r2);
    }

    #[test]
    fn from_bytes() {
        assert_eq!(
            FieldElement::from_bytes(&FieldBytes::default()).unwrap(),
            FieldElement::ZERO
        );
        assert_eq!(
            FieldElement::from_bytes(&Array([
                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 1
            ]))
            .unwrap(),
            FieldElement::ONE
        );
        assert!(bool::from(
            FieldElement::from_bytes(&Array([0xff; 32])).is_none()
        ));
    }

    #[test]
    fn to_bytes() {
        assert_eq!(FieldElement::ZERO.to_bytes(), FieldBytes::default());
        assert_eq!(
            FieldElement::ONE.to_bytes(),
            [
                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0, 1
            ]
        );
    }

    #[test]
    fn repeated_add() {
        let mut r = FieldElement::ONE;
        for item in DBL_TEST_VECTORS {
            assert_eq!(r.to_bytes().as_slice(), item);
            r = r + &r;
        }
    }

    #[test]
    fn repeated_double() {
        let mut r = FieldElement::ONE;
        for item in DBL_TEST_VECTORS {
            assert_eq!(r.to_bytes().as_slice(), item);
            r = r.double();
        }
    }

    #[test]
    fn repeated_mul() {
        let mut r = FieldElement::ONE;
        let two = r + &r;
        for item in DBL_TEST_VECTORS {
            assert_eq!(r.to_bytes().as_slice(), item);
            r = r * &two;
        }
    }

    #[test]
    fn negation() {
        let two = FieldElement::ONE.double();
        let neg_two = -two;
        assert_eq!(two + &neg_two, FieldElement::ZERO);
        assert_eq!(-neg_two, two);
    }

    #[test]
    fn pow_vartime() {
        let one = FieldElement::ONE;
        let two = one + &one;
        let four = two.square();
        assert_eq!(two.pow_vartime(&U256::from_u64(2)), four);
    }

    cpubits! {
        64 => {
            proptest! {
                /// This checks behaviour well within the field ranges, because it doesn't set the
                /// highest limb.
                #[test]
                fn add_then_sub(
                    a0 in ANY,
                    a1 in ANY,
                    a2 in ANY,
                    b0 in ANY,
                    b1 in ANY,
                    b2 in ANY,
                ) {
                    let a = FieldElement::from_montgomery(U256::from_words([a0, a1, a2, 0]));
                    let b = FieldElement::from_montgomery(U256::from_words([b0, b1, b2, 0]));
                    assert_eq!(a.add(&b).sub(&a), b);
                }
            }
        }
    }
}
