arcium-primitives 0.4.2

use std::{
    iter::{Product, Sum},
    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
};

use crypto_bigint::rand_core::RngCore;
use ff::{Field, PrimeField};
use hybrid_array::Array;
use rand::Rng;
use serde::{Deserialize, Serialize};
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
use typenum::{U1, U14, U16};
use wincode::{SchemaRead, SchemaWrite};

use crate::{
    algebra::{
        field::FieldExtension,
        ops::{AccReduce, ReduceWide},
        uniform_bytes::FromUniformBytes,
    },
    random::{CryptoRngCore, Random},
    types::{HeapArray, Positive},
};

mod ff_impl {
    use ff::PrimeField;
    use serde::{Deserialize, Serialize};

    #[derive(PrimeField, Serialize, Deserialize)]
    #[PrimeFieldModulus = "162259276829213363391578010288127"]
    #[PrimeFieldGenerator = "3"]
    #[PrimeFieldReprEndianness = "little"]
    pub struct Mersenne107FF([u64; 2]);
}

#[derive(
    Copy, Clone, Default, Debug, SchemaRead, SchemaWrite, PartialEq, Eq, Hash, Ord, PartialOrd,
)]
#[repr(C)]
pub struct Mersenne107(pub(super) u128);

impl Serialize for Mersenne107 {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        self.as_le_array().serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for Mersenne107 {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let arr = <[u8; 14]>::deserialize(deserializer)?;
        Self::from_canonical_bytes(&arr).ok_or_else(|| {
            serde::de::Error::custom("Invalid Mersenne107 canonical byte representation")
        })
    }
}

impl Mersenne107 {
    pub const NUM_BITS: usize = 107;
    pub const MODULUS: u128 = (1u128 << Self::NUM_BITS) - 1;
    pub const MAX: u128 = Self::MODULUS - 1;

    fn as_le_array(&self) -> [u8; 14] {
        let mut arr = [0u8; 14];
        arr[..14].copy_from_slice(&self.0.to_le_bytes()[..14]);
        arr
    }

    fn from_canonical_bytes(arr: &[u8; 14]) -> Option<Self> {
        let mut tmp = [0u8; 16];
        tmp[..14].copy_from_slice(arr);
        let val = u128::from_le_bytes(tmp);
        (val < Self::MODULUS).then_some(Self(val))
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Arithmetic ops - multiplication
///////////////////////////////////////////////////////////////////////////////////////////////////

#[macros::op_variants(owned)]
impl<'a> MulAssign<&'a Mersenne107> for Mersenne107 {
    #[inline]
    fn mul_assign(&mut self, rhs: &'a Mersenne107) {
        self.0 = super::m107_ops::mul(self.0, rhs.0);
    }
}

#[macros::op_variants(owned)]
impl<'a> Mul<&'a Mersenne107> for Mersenne107 {
    type Output = Self;
    #[inline]
    fn mul(self, rhs: &'a Mersenne107) -> Self::Output {
        let mut res = self;
        res.mul_assign(rhs);
        res
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Arithmetic ops - addition
///////////////////////////////////////////////////////////////////////////////////////////////////

#[macros::op_variants(owned)]
impl<'a> AddAssign<&'a Mersenne107> for Mersenne107 {
    #[inline]
    fn add_assign(&mut self, rhs: &'a Mersenne107) {
        self.0 += rhs.0;
        super::m107_ops::reduce_mod_1bit_inplace(&mut self.0);
    }
}
#[macros::op_variants(owned)]
impl<'a> Add<&'a Mersenne107> for Mersenne107 {
    type Output = Self;

    #[inline]
    fn add(self, rhs: &'a Mersenne107) -> Self::Output {
        let mut res = self;
        res.add_assign(rhs);
        res
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Arithmetic ops - substraction
///////////////////////////////////////////////////////////////////////////////////////////////////

#[macros::op_variants(owned)]
impl<'a> SubAssign<&'a Mersenne107> for Mersenne107 {
    #[inline]
    fn sub_assign(&mut self, rhs: &'a Mersenne107) {
        self.0 += Self::MODULUS - rhs.0;
        super::m107_ops::reduce_mod_1bit_inplace(&mut self.0);
    }
}

#[macros::op_variants(owned)]
impl<'a> Sub<&'a Mersenne107> for Mersenne107 {
    type Output = Self;

    #[inline]
    fn sub(mut self, rhs: &'a Mersenne107) -> Self::Output {
        self.sub_assign(rhs);
        self
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Arithmetic ops - negation
///////////////////////////////////////////////////////////////////////////////////////////////////

#[macros::op_variants(borrowed)]
impl Neg for Mersenne107 {
    type Output = Mersenne107;

    fn neg(self) -> Self::Output {
        Self(super::m107_ops::reduce_mod_1bit(Self::MODULUS - self.0))
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Constant time
///////////////////////////////////////////////////////////////////////////////////////////////////

impl ConditionallySelectable for Mersenne107 {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self(u128::conditional_select(&a.0, &b.0, choice))
    }
}

impl ConstantTimeEq for Mersenne107 {
    fn ct_eq(&self, other: &Self) -> Choice {
        self.0.ct_eq(&other.0)
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Iterator operations
///////////////////////////////////////////////////////////////////////////////////////////////////

impl Sum for Mersenne107 {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        let t = iter.fold(<Self as AccReduce>::zero_wide(), |mut acc, x| {
            Self::acc(&mut acc, &x);
            acc
        });
        Self::reduce_mod_order(t)
    }
}

impl<'a> Sum<&'a Self> for Mersenne107 {
    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        let t = iter.fold(<Self as AccReduce>::zero_wide(), |mut acc, x| {
            Self::acc(&mut acc, x);
            acc
        });
        Self::reduce_mod_order(t)
    }
}

impl<'a> Product<&'a Self> for Mersenne107 {
    fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.fold(Self::ONE, |acc, x| acc * x)
    }
}

impl Product for Mersenne107 {
    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::ONE, |acc, x| acc * x)
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Field trait
///////////////////////////////////////////////////////////////////////////////////////////////////

impl Field for Mersenne107 {
    const ZERO: Self = Mersenne107(0);
    const ONE: Self = Mersenne107(1);

    fn random(mut rng: impl RngCore) -> Self {
        let tmp = rng.gen::<u128>();
        Self(super::m107_ops::reduce_mod(tmp)) // the probability is skewed because M107 modulus
                                               // does not
                                               // divide 2^128 exactly
    }

    fn square(&self) -> Self {
        *self * self // TODO: optimize ?
    }

    fn double(&self) -> Self {
        Self(super::m107_ops::reduce_mod_1bit(self.0 << 1))
    }

    fn invert(&self) -> CtOption<Self> {
        // Fallback to ff implementation
        // TODO: see if we can optimize this without ff
        let val: ff_impl::Mersenne107FF = self.into();
        let inv = val.invert();
        inv.map(|v| v.into())
    }

    fn sqrt_ratio(num: &Self, div: &Self) -> (Choice, Self) {
        let num = num.into();
        let div = div.into();
        let (choice, val) = ff_impl::Mersenne107FF::sqrt_ratio(&num, &div);
        (choice, val.into())
    }

    fn sqrt(&self) -> CtOption<Self> {
        // Fallback to ff implementation
        // TODO: see if we can optimize this without ff
        let val: ff_impl::Mersenne107FF = self.into();
        let inv = val.sqrt();
        inv.map(|v| v.into())
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Field extension trait
///////////////////////////////////////////////////////////////////////////////////////////////////

impl FieldExtension for Mersenne107 {
    type Subfield = Self;
    type Degree = U1;
    type FieldBitSize = typenum::U<{ Mersenne107::NUM_BITS }>;
    type FieldBytesSize = U14;

    fn to_subfield_elements(&self) -> impl ExactSizeIterator<Item = Self::Subfield> {
        std::iter::once(*self)
    }

    fn from_subfield_elements(elems: &[Self::Subfield]) -> Option<Self> {
        if elems.len() == 1 {
            elems.first().copied()
        } else {
            None
        }
    }

    fn to_le_bytes(&self) -> Array<u8, Self::FieldBytesSize> {
        self.as_le_array().into()
    }

    fn from_le_bytes(bytes: &[u8]) -> Option<Self> {
        if bytes.len() == 14 {
            let arr: &[u8; 14] = bytes.try_into().expect("This should never fail");
            Self::from_canonical_bytes(arr)
        } else {
            None
        }
    }

    fn mul_by_subfield(&self, other: &Self::Subfield) -> Self {
        *self * other
    }

    fn generator() -> Self {
        Self(3u128) // 3^(2^107-2) == 1 mod 2^107-1
    }
}

impl Random for Mersenne107 {
    fn random(mut rng: impl CryptoRngCore) -> Self {
        let tmp = rng.gen::<u128>();
        Self(super::m107_ops::reduce_mod(tmp))
    }

    fn random_array<M: Positive>(mut rng: impl CryptoRngCore) -> HeapArray<Self, M> {
        let mut buf = HeapArray::<Self, M>::default().into_box_bytes();
        rng.fill_bytes(&mut buf);
        let mut tmp = HeapArray::from_box_bytes(buf);
        tmp.iter_mut()
            .for_each(|v: &mut Self| super::m107_ops::reduce_mod_inplace(&mut v.0));
        tmp
    }
}

unsafe impl bytemuck::Zeroable for Mersenne107 {}
unsafe impl bytemuck::Pod for Mersenne107 {}

impl FromUniformBytes for Mersenne107 {
    type UniformBytes = U16;
    fn from_uniform_bytes(bytes: &hybrid_array::Array<u8, Self::UniformBytes>) -> Self {
        let mut val = u128::from_le_bytes(bytes.0);
        super::m107_ops::reduce_mod_inplace(&mut val);
        Self(val)
    }
}

impl From<u64> for Mersenne107 {
    fn from(val: u64) -> Self {
        Self(val as u128)
    }
}

impl From<u128> for Mersenne107 {
    fn from(val: u128) -> Self {
        Self(super::m107_ops::reduce_mod(val))
    }
}

impl From<ff_impl::Mersenne107FF> for Mersenne107 {
    fn from(val: ff_impl::Mersenne107FF) -> Self {
        Self::from_le_bytes(&val.to_repr().as_ref()[..14]).unwrap()
    }
}

impl<'a> From<&'a Mersenne107> for ff_impl::Mersenne107FF {
    fn from(val: &'a Mersenne107) -> Self {
        Self::from_repr(ff_impl::Mersenne107FFRepr(val.0.to_le_bytes())).unwrap()
    }
}

#[cfg(test)]
mod test {
    use ff::Field;
    use num_bigint::BigInt;
    use typenum::Unsigned;

    use crate::{
        algebra::field::{
            mersenne::{m107::Mersenne107, test::bigint_to_m107},
            FieldExtension,
        },
        random::test_rng,
    };

    type M = typenum::U1000;

    #[test]
    fn test_neg() {
        fn test_internal(a: Mersenne107) {
            let exp = bigint_to_m107(-BigInt::from(a.0));
            let act = -a;
            assert_eq!(exp, act, "a = {a:?}");
        }

        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let a = Mersenne107::random(&mut rng);
            test_internal(a);
        }

        // Corner cases
        test_internal(Mersenne107::ZERO);
        test_internal(Mersenne107::ONE);
        test_internal(Mersenne107(Mersenne107::MAX));
    }

    #[test]
    fn test_invert() {
        fn test_internal(a: Mersenne107) {
            let a_inv = a.invert().unwrap();
            let act = a * a_inv;
            assert_eq!(Mersenne107::ONE, act, "a = {a:?}");
        }

        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let a = Mersenne107::random(&mut rng);
            if a == Mersenne107::ZERO {
                continue;
            }
            test_internal(a);
        }

        // Corner cases
        test_internal(Mersenne107::ONE);
        test_internal(Mersenne107(Mersenne107::MAX));
    }

    #[test]
    fn test_sqrt() {
        fn test_internal(a: Mersenne107) {
            let a_sqrt = a.sqrt();
            if a_sqrt.into_option().is_none() {
                return;
            }

            let a_sqrt = a_sqrt.unwrap();
            let act = a_sqrt * a_sqrt;
            assert_eq!(a, act, "a = {a:?}");
        }

        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let a = Mersenne107::random(&mut rng);
            test_internal(a);
        }

        // Corner cases
        test_internal(Mersenne107::ZERO);
        test_internal(Mersenne107::ONE);
        test_internal(Mersenne107(Mersenne107::MAX));
    }

    #[test]
    fn test_canonical_bytes_decoding() {
        let value = Mersenne107::from(123456789u64);
        let bytes = value.to_le_bytes();
        assert_eq!(Mersenne107::from_le_bytes(&bytes), Some(value));

        // Encoding of the modulus p = 2^107 - 1 is non-canonical and must be rejected.
        let modulus_bytes = [
            0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x07,
        ];
        assert_eq!(Mersenne107::from_le_bytes(&modulus_bytes), None);
    }

    macro_rules! test_op {
        ($op:tt) => {
                fn test_internal(a: Mersenne107, b: Mersenne107) {
                    let exp = bigint_to_m107(BigInt::from(a.0) $op BigInt::from(b.0));
                    let act = a $op b;
                    assert_eq!(exp, act, "a = {a:?}, b = {b:?}");
                }

                let mut rng = test_rng();
                for _ in 0..M::to_usize() {
                    let a = Mersenne107::random(&mut rng);
                    let b = Mersenne107::random(&mut rng);
                    test_internal(a, b);
                }

                // Corner cases
                test_internal(Mersenne107::ZERO, Mersenne107::ZERO);
                test_internal(Mersenne107::ZERO, Mersenne107::ONE);
                test_internal(Mersenne107::ONE, Mersenne107::ZERO);
                test_internal(Mersenne107::ONE, Mersenne107::ONE);
                test_internal(Mersenne107::ZERO, Mersenne107(Mersenne107::MAX));
                test_internal(Mersenne107::ONE, Mersenne107(Mersenne107::MAX));
                test_internal(Mersenne107(Mersenne107::MAX), Mersenne107::ZERO);
                test_internal(Mersenne107(Mersenne107::MAX), Mersenne107::ONE);
                test_internal(Mersenne107(Mersenne107::MAX), Mersenne107(Mersenne107::MAX));
            }
    }

    #[test]
    fn test_mul() {
        test_op!(*);
    }
    #[test]
    fn test_add() {
        test_op!(+);
    }
    #[test]
    fn test_sub() {
        test_op!(-);
    }
}