//! Macros for generating wrappers for `fiat-crypto` synthesized field implementations, currently
//! specialized to support `word-by-word-montgomery` implementations.

/// Add `fiat-crypto` synthesized arithmetic impls to the given field element.
///
/// This macro wraps a field implementation generated by the `word-by-word-montgomery` backend
/// (other field implementations are not supported) and uses it in conjunction with this crate's
/// `MontyField` type as the internal representation.
#[macro_export]
macro_rules! fiat_monty_field_arithmetic {
    (
        name: $fe:tt,
        params: $params:ty,
        uint: $uint:ty,
        non_mont: $non_mont_type:expr,
        mont: $mont_type:expr,
        from_mont: $from_mont:ident,
        to_mont: $to_mont:ident,
        add: $add:ident,
        sub: $sub:ident,
        mul: $mul:ident,
        neg: $neg:ident,
        square: $square:ident,
        divstep_precomp: $divstep_precomp:ident,
        divstep: $divstep:ident,
        msat: $msat:ident,
        selectnz: $selectznz:ident
    ) => {
        impl $fe {
            /// Decode [`
            #[doc = stringify!($fe)]
            /// `] from [`
            #[doc = stringify!($uint)]
            /// `] converting it into Montgomery form.
            ///
            /// Does *not* perform a check that the field element does not overflow the order.
            ///
            /// Used incorrectly this can lead to invalid results!
            #[inline]
            pub(crate) const fn from_uint_unchecked(w: $uint) -> Self {
                let mut out = $mont_type([0; <$uint>::LIMBS]);
                $to_mont(&mut out, &$non_mont_type(w.to_words()));
                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        out.0,
                    ),
                )
            }

            /// Translate [`
            #[doc = stringify!($fe)]
            /// `] out of the Montgomery domain, returning a [`
            #[doc = stringify!($uint)]
            /// `] in canonical form.
            #[inline]
            pub const fn to_canonical(self) -> $uint {
                let mut out = $non_mont_type([0; <$uint>::LIMBS]);
                $from_mont(&mut out, &$mont_type(self.0.to_montgomery_words()));
                <$uint>::from_words(out.0)
            }

            /// Add elements.
            #[inline]
            pub const fn add(&self, rhs: &Self) -> Self {
                let mut out = $mont_type([0; <$uint>::LIMBS]);
                $add(
                    &mut out,
                    &$mont_type(self.0.to_montgomery_words()),
                    &$mont_type(rhs.0.to_montgomery_words()),
                );
                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        out.0,
                    ),
                )
            }

            /// Double element (add it to itself).
            #[inline]
            #[must_use]
            pub const fn double(&self) -> Self {
                self.add(self)
            }

            /// Subtract elements.
            #[inline]
            pub const fn sub(&self, rhs: &Self) -> Self {
                let mut out = $mont_type([0; <$uint>::LIMBS]);
                $sub(
                    &mut out,
                    &$mont_type(self.0.to_montgomery_words()),
                    &$mont_type(rhs.0.to_montgomery_words()),
                );
                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        out.0,
                    ),
                )
            }

            /// Multiply elements.
            #[inline]
            pub const fn multiply(&self, rhs: &Self) -> Self {
                let mut out = $mont_type([0; <$uint>::LIMBS]);
                $mul(
                    &mut out,
                    &$mont_type(self.0.to_montgomery_words()),
                    &$mont_type(rhs.0.to_montgomery_words()),
                );
                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        out.0,
                    ),
                )
            }

            /// Negate element.
            #[inline]
            pub const fn neg(&self) -> Self {
                let mut out = $mont_type([0; <$uint>::LIMBS]);
                $neg(&mut out, &$mont_type(self.0.to_montgomery_words()));
                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        out.0,
                    ),
                )
            }

            /// Compute modular square.
            #[inline]
            #[must_use]
            pub const fn square(&self) -> Self {
                let mut out = $mont_type([0; <$uint>::LIMBS]);
                $square(&mut out, &$mont_type(self.0.to_montgomery_words()));
                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        out.0,
                    ),
                )
            }

            /// Compute
            #[doc = stringify!($fe)]
            /// inversion: `1 / self`.
            #[inline]
            pub fn invert(&self) -> $crate::subtle::CtOption<Self> {
                $crate::subtle::CtOption::new(self.invert_unchecked(), !self.is_zero())
            }

            /// Compute field inversion as a `const fn`. Panics if `self` is zero.
            ///
            /// This is mainly intended for inverting constants at compile time.
            pub const fn const_invert(&self) -> Self {
                assert!(
                    !self.0.as_montgomery().cmp_vartime(&<$uint>::ZERO).is_eq(),
                    "input to invert should be non-zero"
                );
                self.invert_unchecked()
            }

            /// Returns the multiplicative inverse of self.
            ///
            /// Does not check that self is non-zero.
            const fn invert_unchecked(&self) -> Self {
                let words = $crate::fiat_bernstein_yang_invert!(
                    a: &$mont_type(self.0.to_montgomery_words()),
                    one: &$mont_type(Self::ONE.0.to_montgomery_words()),
                    d: <$fe as $crate::ff::PrimeField>::NUM_BITS as usize,
                    nlimbs: <$uint>::LIMBS,
                    word: $crate::bigint::Word,
                    non_mont: $non_mont_type,
                    mont: $mont_type,
                    from_mont: $from_mont,
                    mul: $mul,
                    neg: $neg,
                    divstep_precomp: $divstep_precomp,
                    divstep: $divstep,
                    msat: $msat,
                    selectnz: $selectznz
                );

                Self(
                    $crate::MontyFieldElement::<$params, { <$uint>::LIMBS }>::from_montgomery_words(
                        words,
                    ),
                )
            }
        }
    };
}

/// Emit wrapper function for a `fiat-crypto` generated implementation of the Bernstein-Yang
/// (a.k.a. safegcd) modular inversion algorithm.
#[macro_export]
macro_rules! fiat_bernstein_yang_invert {
    (
        a: $a:expr,
        one: $one:expr,
        d: $d:expr,
        nlimbs: $nlimbs:expr,
        word: $word:ty,
        non_mont: $non_mont_type: expr,
        mont: $mont_type: expr,
        from_mont: $from_mont:ident,
        mul: $mul:ident,
        neg: $neg:ident,
        divstep_precomp: $divstep_precomp:ident,
        divstep: $divstep:ident,
        msat: $msat:ident,
        selectnz: $selectznz:ident
    ) => {{
        // See Bernstein-Yang 2019 p.366
        const ITERATIONS: usize = (49 * $d + 57) / 17;

        let mut a = $non_mont_type([0; $nlimbs]);
        $from_mont(&mut a, $a);
        let mut d = 1;
        let mut f = [0; $nlimbs + 1];
        $msat(&mut f);
        let mut g = [0; $nlimbs + 1];
        let mut v = [0; $nlimbs];
        let mut r = $one.0;
        let mut i = 0;
        let mut j = 0;

        while j < $nlimbs {
            g[j] = a.0[j];
            j += 1;
        }

        while i < ITERATIONS - ITERATIONS % 2 {
            let mut out1 = 0;
            let mut out2 = [0; $nlimbs + 1];
            let mut out3 = [0; $nlimbs + 1];
            let mut out4 = [0; $nlimbs];
            let mut out5 = [0; $nlimbs];

            $divstep(
                &mut out1, &mut out2, &mut out3, &mut out4, &mut out5, d, &f, &g, &v, &r,
            );
            $divstep(
                &mut d, &mut f, &mut g, &mut v, &mut r, out1, &out2, &out3, &out4, &out5,
            );
            i += 2;
        }

        if ITERATIONS % 2 != 0 {
            let mut out1 = 0;
            let mut out2 = [0; $nlimbs + 1];
            let mut out3 = [0; $nlimbs + 1];
            let mut out4 = [0; $nlimbs];
            let mut out5 = [0; $nlimbs];
            $divstep(
                &mut out1, &mut out2, &mut out3, &mut out4, &mut out5, d, &f, &g, &v, &r,
            );
            f = out2;
            v = out4;
        }

        let s = ((f[f.len() - 1] >> <$word>::BITS - 1) & 1) as u8;
        let mut neg_v = $mont_type([0; $nlimbs]);
        $neg(&mut neg_v, &$mont_type(v));

        let mut v2 = $mont_type([0; $nlimbs]);
        $selectznz(&mut v2.0, s, &v, &neg_v.0);

        let mut precomp = $mont_type([0; $nlimbs]);
        $divstep_precomp(&mut precomp.0);

        let mut out = $mont_type([0; $nlimbs]);
        $mul(&mut out, &v2, &precomp);
        out.0
    }};
}

/// Test that the wrapper generated for a `word-by-word-montgomery` fiat-crypto implementation
/// matches the internal representation used by `MontyField*`, e.g. do they have the same saturated
/// modulus representation and agree on the value of `R`.
///
/// If these tests do not pass, the `fiat-crypto` synthesized implementation is incompatible with
/// `MontyField*`.
///
/// See also: mit-plv/fiat-crypto#2166
#[macro_export]
macro_rules! test_fiat_monty_field_arithmetic {
    (
        name: $fe:tt,
        params: $params:ty,
        uint: $uint:ty,
        non_mont: $non_mont_type:expr,
        mont: $mont_type:expr,
        to_mont: $to_mont:ident,
        msat: $msat:ident
    ) => {
        use $crate::bigint::modular::ConstMontyParams as _;

        /// Compute `fiat-crypto`'s selection of `R` by converting `1` into Montgomery form using
        /// `$to_mont`, and then compare it to `ConstMontyParams::PARAMS.one()` as computed by
        /// `crypto_bigint::modular` to ensure they're equal.
        #[test]
        fn fiat_montgomery_r_constant() {
            let mut one = $non_mont_type([0; <$uint>::LIMBS]);
            one[0] = 1;

            let mut R = $mont_type([0; <$uint>::LIMBS]);
            $to_mont(&mut R, &one);

            assert_eq!(<$params>::PARAMS.one().as_words(), &R.0)
        }

        /// Compute `fiat-crypto`'s saturated modulus representation using `$msat`, and then compare
        /// it to `ConstMontyParams::PARAMS.modulus()` as computed by `crypto_bigint::modular`
        /// to ensure they're equal.
        #[test]
        fn fiat_msat_constant() {
            use $crate::bigint::Word;

            let mut out = [0 as Word; <$uint>::LIMBS + 1];
            $msat(&mut out);
            assert_eq!(out[<$uint>::LIMBS], 0);
            assert_eq!(
                <$params>::PARAMS.modulus().as_words(),
                &out[..<$uint>::LIMBS]
            );
        }
    };
}
