primitives/algebra/field/mersenne/

m107.rs

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)
    }
}

// Fixed-exponent helpers used by the `sqrt`/`sqrt_ratio` closed forms (p ≡ 3 mod 4).
impl Mersenne107 {
    /// Returns `self^(2^n)`, i.e. `n` repeated squarings.
    #[inline]
    fn pow2_pow(self, n: u32) -> Self {
        (0..n).fold(self, |acc, _| acc.square())
    }

    /// Returns `self^(2^k - 1)` via an Itoh–Tsujii-style addition chain: ≈`k` squarings and
    /// `O(log k)` multiplications, versus the ~`2k` operations of plain square-and-multiply on the
    /// all-ones exponent.
    fn pow2_minus_1(self, k: u32) -> Self {
        debug_assert!(k >= 1);
        // Invariant: `acc == self^(2^cur - 1)`; grow `cur` from 1 to `k` following the bits of `k`.
        let mut acc = self;
        let mut cur = 1u32;
        for i in (0..(u32::BITS - 1 - k.leading_zeros())).rev() {
            // self^(2^{2·cur} - 1) = (self^(2^cur - 1))^(2^cur) · self^(2^cur - 1)
            acc = acc.pow2_pow(cur) * acc;
            cur *= 2;
            if (k >> i) & 1 == 1 {
                // self^(2^{cur+1} - 1) = (self^(2^cur - 1))^2 · self
                acc = acc.square() * self;
                cur += 1;
            }
        }
        debug_assert_eq!(cur, k);
        acc
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// 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) {
        // p = 2^107 - 1 ≡ 3 (mod 4), so sqrt(num/div) has the closed form
        //     y = num·div·(num·div³)^((p-3)/4)
        // computed with a single fixed exponentiation and no modular inversion (unlike the generic
        // Tonelli–Shanks `sqrt_ratio`). `y` is a true square root of num/div iff y²·div == num. The
        // exponent (p-3)/4 = 2^105 - 1.
        let uv = *num * div;
        let uv3 = div.square() * uv;
        let y = uv * uv3.pow2_minus_1(105);
        let is_qr = (y.square() * div).ct_eq(num);
        (is_qr, y)
    }

    fn sqrt(&self) -> CtOption<Self> {
        // p ≡ 3 (mod 4): sqrt(a) = a^((p+1)/4) = a^(2^105), i.e. 105 repeated squarings. The result
        // is a valid square root iff it squares back to `a` (rejects non-residues; zero maps to
        // zero).
        let root = self.pow2_pow(105);
        CtOption::new(root, root.square().ct_eq(self))
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// 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) -> Array<Self::Subfield, Self::Degree> {
        Array([*self])
    }

    fn from_subfield_elements(elems: Array<Self::Subfield, Self::Degree>) -> Self {
        elems[0]
    }

    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_sqrt_ratio() {
        fn test_internal(num: Mersenne107, div: Mersenne107) {
            let (is_qr, y) = Mersenne107::sqrt_ratio(&num, &div);

            // QR-ness must agree with the generic ff reference implementation.
            let num_ff: super::ff_impl::Mersenne107FF = (&num).into();
            let div_ff: super::ff_impl::Mersenne107FF = (&div).into();
            let (is_qr_ref, _) = super::ff_impl::Mersenne107FF::sqrt_ratio(&num_ff, &div_ff);
            assert_eq!(
                bool::from(is_qr),
                bool::from(is_qr_ref),
                "QR mismatch for num={num:?} div={div:?}"
            );

            // When num/div is a square, y must satisfy y² · div == num.
            if bool::from(is_qr) {
                assert_eq!(
                    y.square() * div,
                    num,
                    "wrong root for num={num:?} div={div:?}"
                );
            }
        }

        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let num = Mersenne107::random(&mut rng);
            let div = Mersenne107::random(&mut rng);
            if div == Mersenne107::ZERO {
                continue;
            }
            test_internal(num, div);
            // The daBits use case: inverse square root sqrt_ratio(1, y) == y^{-1/2}.
            test_internal(Mersenne107::ONE, div);
        }

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

    #[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!(-);
    }
}