primitives/algebra/field/mersenne/

m107.rs

use std::{
    iter::{Product, Sum},
    mem::MaybeUninit,
    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 crate::{
    algebra::{
        field::{
            exponentiation::{pow2_minus_1, pow2_pow},
            FieldExtension,
        },
        ops::{AccReduce, ReduceWide},
        uniform_bytes::FromUniformBytes,
    },
    errors::PrimitiveError,
    random::{CryptoRngCore, Random},
    types::{HeapArray, Positive},
    utils::codec::InPlaceCodec,
};

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, 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) {
        // 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 * pow2_minus_1(uv3, 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 = pow2_pow(*self, 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 {}

// SAFETY: `write_le_bytes`/`read_le_bytes` use a 14-byte canonical little-endian encoding
// (architecture-independent); `write_le_bytes` initializes all 14 bytes and `read_le_bytes`
// range-validates via `from_canonical_bytes`, so the round-trip is unbiased.
unsafe impl InPlaceCodec for Mersenne107 {
    const ENCODED_SIZE: usize = 14;

    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
        let bytes = self.as_le_array();
        // SAFETY: `out.len() == ENCODED_SIZE == bytes.len()`; `copy_nonoverlapping` initializes
        // every byte of `out` without reading it.
        unsafe {
            std::ptr::copy_nonoverlapping(
                bytes.as_ptr(),
                out.as_mut_ptr().cast::<u8>(),
                bytes.len(),
            );
        }
    }

    // Per-element canonical-range check. Batching the check across a `HeapArray` (single pass or
    // branchless accumulator) was measured slower: the branch is near-always true for real field
    // elements and thus near-perfectly predicted, while a batch pass only adds memory traffic.
    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
        let arr = bytes
            .try_into()
            .map_err(|_| PrimitiveError::InvalidSize(Self::ENCODED_SIZE, bytes.len()))?;
        Self::from_canonical_bytes(arr).ok_or_else(|| {
            PrimitiveError::DeserializationFailed("non-canonical Mersenne107 encoding".into())
        })
    }

    // Vectorized write: 8 elements still take `8 * 14 = 112` bytes (same total as the per-element
    // path, so `PACK_BYTES`/`read_pack` fall back to the defaults), but use wide overlapping
    // `u128` stores instead of 8 strided 14-byte copies.
    const PACK: usize = 8;

    fn write_pack(items: &[Self], out: &mut [MaybeUninit<u8>]) {
        // `items.len() == 8`, `out.len() == 112`. Emit each element's full 16 little-endian bytes;
        // the top 2 bytes are 0 for canonical values (< 2^107) and get overwritten by the next
        // element, so the overlap is harmless. The last element writes only 14 bytes to stay in
        // bounds. This reproduces the per-element 14-byte layout exactly.
        let dst = out.as_mut_ptr().cast::<u8>();
        for (i, item) in items.iter().enumerate().take(Self::PACK - 1) {
            let le = item.0.to_le_bytes();
            // SAFETY: `i * 14 + 16 <= 6 * 14 + 16 = 100 <= 112`; `le` has 16 bytes.
            unsafe { std::ptr::copy_nonoverlapping(le.as_ptr(), dst.add(i * 14), 16) };
        }
        let last = items[Self::PACK - 1].0.to_le_bytes();
        // SAFETY: `(PACK - 1) * 14 + 14 == 112 == out.len()`; `last` has >= 14 bytes.
        unsafe { std::ptr::copy_nonoverlapping(last.as_ptr(), dst.add((Self::PACK - 1) * 14), 14) };
    }
}

impl FromUniformBytes for Mersenne107 {
    type UniformBytes = U16;
    fn from_uniform_bytes(bytes: &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!(-);
    }
}