arcium-primitives 0.8.1

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

use ff::Field;
use hybrid_array::Array;
use rand::RngCore;
use serde::{Deserialize, Serialize};
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
use typenum::{U1, U128, U16};

use crate::{
    algebra::{
        field::{
            binary::{
                gf2_ext::{Gf2LimbsWide, MulWide},
                Gf2_128,
            },
            FieldExtension,
        },
        ops::{AccReduce, DefaultDotProduct, IntoWide, MulAccReduce, ReduceWide},
        uniform_bytes::FromUniformBytes,
    },
    errors::PrimitiveError,
    random::{CryptoRngCore, Random},
    types::{HeapArray, Positive},
    utils::codec::InPlaceCodec,
};

/// GF(2^128) presented as a standalone field (`Subfield = Self`).
///
/// Unlike [`Gf2_128`], whose `FieldExtension` impl has `Subfield = Gf2` (so
/// `FieldShare<Gf2_128>` = `BitShare` authenticates a single bit), this wrapper is its own
/// subfield: `FieldShare<Gf2_128Field>` authenticates a full 128-bit secret.
///
/// Caveat:
/// - Order-dependent plaintext circuit ops (`Gt`/`Ge`, signed `BitExtract`, `EuclDiv`/`Mod`) are
///   well-defined on the byte representation but carry no algebraic meaning in char 2.
#[derive(
    Copy, Clone, Default, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize,
)]
#[serde(transparent)]
#[repr(transparent)]
pub struct Gf2_128Field(pub Gf2_128);

// SAFETY: `Gf2_128Field` is `#[repr(transparent)]` over `Gf2_128`; its encoding is exactly
// `Gf2_128`'s `InPlaceCodec` encoding.
unsafe impl InPlaceCodec for Gf2_128Field {
    const ENCODED_SIZE: usize = Gf2_128::ENCODED_SIZE;

    fn write_le_bytes(&self, out: &mut [MaybeUninit<u8>]) {
        self.0.write_le_bytes(out);
    }

    fn read_le_bytes(bytes: &[u8]) -> Result<Self, PrimitiveError> {
        Gf2_128::read_le_bytes(bytes).map(Gf2_128Field)
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Arithmetic ops
///////////////////////////////////////////////////////////////////////////////////////////////////

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

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

    #[inline]
    fn mul(mut self, rhs: &'a Gf2_128Field) -> Self::Output {
        self.mul_assign(rhs);
        self
    }
}

#[macros::op_variants(owned)]
impl<'a> AddAssign<&'a Gf2_128Field> for Gf2_128Field {
    #[inline]
    fn add_assign(&mut self, rhs: &'a Gf2_128Field) {
        self.0 += rhs.0;
    }
}

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

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

#[macros::op_variants(owned)]
impl<'a> SubAssign<&'a Gf2_128Field> for Gf2_128Field {
    #[inline]
    fn sub_assign(&mut self, rhs: &'a Gf2_128Field) {
        self.0 -= rhs.0;
    }
}

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

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

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

    #[inline]
    fn neg(self) -> Self::Output {
        self
    }
}

impl Sum for Gf2_128Field {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(<Self as Field>::ZERO, |acc, x| acc + x)
    }
}

impl<'a> Sum<&'a Self> for Gf2_128Field {
    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.fold(<Self as Field>::ZERO, |acc, x| acc + x)
    }
}

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

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

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

impl ConditionallySelectable for Gf2_128Field {
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self(Gf2_128::conditional_select(&a.0, &b.0, choice))
    }
}

impl ConstantTimeEq for Gf2_128Field {
    #[inline]
    fn ct_eq(&self, other: &Self) -> Choice {
        self.0.ct_eq(&other.0)
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Field traits
///////////////////////////////////////////////////////////////////////////////////////////////////

impl Field for Gf2_128Field {
    const ZERO: Self = Self(<Gf2_128 as Field>::ZERO);
    const ONE: Self = Self(<Gf2_128 as Field>::ONE);

    fn random(rng: impl RngCore) -> Self {
        Self(<Gf2_128 as Field>::random(rng))
    }

    fn square(&self) -> Self {
        Self(self.0.square())
    }

    fn double(&self) -> Self {
        Self(self.0.double())
    }

    fn invert(&self) -> CtOption<Self> {
        self.0.invert().map(Self)
    }

    fn sqrt_ratio(num: &Self, div: &Self) -> (Choice, Self) {
        let (is_valid, root) = Gf2_128::sqrt_ratio(&num.0, &div.0);
        (is_valid, Self(root))
    }

    fn sqrt(&self) -> CtOption<Self> {
        self.0.sqrt().map(Self)
    }
}

impl FieldExtension for Gf2_128Field {
    type Subfield = Self;
    type Degree = U1;
    type FieldBitSize = U128;
    type FieldBytesSize = U16;

    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.0.to_le_bytes()
    }

    fn from_le_bytes(bytes: &[u8]) -> Option<Self> {
        Gf2_128::from_le_bytes(bytes).map(Self)
    }

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

    fn generator() -> Self {
        Self(Gf2_128::generator())
    }

    // Must delegate to the char-2 override: the trait default `*self * 2` is the zero map in a
    // binary field.
    fn linear_orthomorphism(&self) -> Self {
        Self(self.0.linear_orthomorphism())
    }
}

impl Random for Gf2_128Field {
    fn random(rng: impl CryptoRngCore) -> Self {
        Self(Random::random(rng))
    }

    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);
        HeapArray::from_box_bytes(buf)
    }
}

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

impl FromUniformBytes for Gf2_128Field {
    type UniformBytes = U16;

    fn from_uniform_bytes(bytes: &Array<u8, Self::UniformBytes>) -> Self {
        Self(Gf2_128::from_uniform_bytes(bytes))
    }
}

impl From<u64> for Gf2_128Field {
    fn from(val: u64) -> Self {
        Self(Gf2_128::from(val))
    }
}

impl From<u128> for Gf2_128Field {
    fn from(val: u128) -> Self {
        Self(Gf2_128::from(val))
    }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Wide ops
///////////////////////////////////////////////////////////////////////////////////////////////////

impl IntoWide<Gf2LimbsWide<2>> for Gf2_128Field {
    #[inline]
    fn to_wide(&self) -> Gf2LimbsWide<2> {
        <Gf2_128 as IntoWide<Gf2LimbsWide<2>>>::to_wide(&self.0)
    }

    #[inline]
    fn zero_wide() -> Gf2LimbsWide<2> {
        <Gf2_128 as IntoWide<Gf2LimbsWide<2>>>::zero_wide()
    }
}

impl ReduceWide<Gf2LimbsWide<2>> for Gf2_128Field {
    #[inline]
    fn reduce_mod_order(a: Gf2LimbsWide<2>) -> Self {
        Self(Gf2_128::reduce_mod_order(a))
    }
}

impl IntoWide for Gf2_128Field {
    #[inline]
    fn to_wide(&self) -> Self {
        *self
    }

    #[inline]
    fn zero_wide() -> Self {
        <Self as Field>::ZERO
    }
}

impl ReduceWide for Gf2_128Field {
    #[inline]
    fn reduce_mod_order(a: Self) -> Self {
        a
    }
}

impl MulAccReduce for Gf2_128Field {
    type WideType = Gf2LimbsWide<2>;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: Self, b: Self) {
        *acc += a.0.mul_wide(b.0);
    }
}

impl MulAccReduce<Self, &Self> for Gf2_128Field {
    type WideType = Gf2LimbsWide<2>;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: Self, b: &Self) {
        <Self as MulAccReduce>::mul_acc(acc, a, *b);
    }
}

impl MulAccReduce<&Self, Self> for Gf2_128Field {
    type WideType = Gf2LimbsWide<2>;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: &Self, b: Self) {
        <Self as MulAccReduce>::mul_acc(acc, *a, b);
    }
}

impl MulAccReduce<&Self, &Self> for Gf2_128Field {
    type WideType = Gf2LimbsWide<2>;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: &Self, b: &Self) {
        <Self as MulAccReduce>::mul_acc(acc, *a, *b);
    }
}

impl AccReduce for Gf2_128Field {
    type WideType = Self;

    #[inline]
    fn acc(acc: &mut Self, a: Self) {
        *acc += a;
    }
}

impl AccReduce<&Self> for Gf2_128Field {
    type WideType = Self;

    #[inline]
    fn acc(acc: &mut Self, a: &Self) {
        *acc += a;
    }
}

impl DefaultDotProduct for Gf2_128Field {}
impl DefaultDotProduct<Self, &Self> for Gf2_128Field {}
impl DefaultDotProduct<&Self, Self> for Gf2_128Field {}
impl DefaultDotProduct<&Self, &Self> for Gf2_128Field {}

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

    use super::Gf2_128Field;
    use crate::{
        algebra::{field::FieldExtension, ops::DotProduct},
        random::{test_rng, Random},
        utils::bincode_io,
    };

    type M = typenum::U100;

    #[test]
    fn test_field_axioms() {
        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let a: Gf2_128Field = Random::random(&mut rng);
            let b: Gf2_128Field = Random::random(&mut rng);

            // Char 2: a + a = 0, -a = a
            assert_eq!(a + a, Gf2_128Field::ZERO);
            assert_eq!(-a, a);
            assert_eq!(a - b, a + b);

            // Distributivity
            assert_eq!(a * (a + b), a * a + a * b);

            // Frobenius: a^(2^128) = a
            let mut cumul = a;
            for _ in 0..128 {
                cumul *= cumul;
            }
            assert_eq!(a, cumul);
        }
    }

    #[test]
    fn test_invert_and_sqrt() {
        let mut rng = test_rng();

        assert!(bool::from(Gf2_128Field::ZERO.invert().is_none()));

        for _ in 0..M::to_usize() {
            let a: Gf2_128Field = Random::random(&mut rng);
            if bool::from(a.is_zero()) {
                continue;
            }
            assert_eq!(a * a.invert().unwrap(), Gf2_128Field::ONE);

            let root = a.sqrt().unwrap();
            assert_eq!(root * root, a);

            let (is_valid, ratio_root) = Gf2_128Field::sqrt_ratio_ext(&Gf2_128Field::ONE, &a);
            assert!(bool::from(is_valid));
            assert_eq!(ratio_root * ratio_root * a, Gf2_128Field::ONE);
        }
    }

    #[test]
    fn test_linear_orthomorphism() {
        // σ and σ'(x) = σ(x) + x must both be permutations; in particular neither may be the zero
        // map (the trait default `x * 2` is zero in char 2).
        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let a: Gf2_128Field = Random::random(&mut rng);
            if bool::from(a.is_zero()) {
                continue;
            }
            let sigma = a.linear_orthomorphism();
            assert_ne!(sigma, Gf2_128Field::ZERO);
            assert_ne!(sigma + a, Gf2_128Field::ZERO);
        }
    }

    #[test]
    fn test_byte_roundtrips() {
        let mut rng = test_rng();
        for _ in 0..M::to_usize() {
            let a: Gf2_128Field = Random::random(&mut rng);

            let bytes = a.to_le_bytes();
            assert_eq!(Gf2_128Field::from_le_bytes(&bytes), Some(a));

            let serialized = bincode_io::serialize(&a).unwrap();
            assert_eq!(
                bincode_io::deserialize::<Gf2_128Field>(&serialized).unwrap(),
                a
            );
        }
    }

    #[test]
    fn test_dot_product() {
        let mut rng = test_rng();
        let a = Gf2_128Field::random_array::<typenum::U13>(&mut rng);
        let b = Gf2_128Field::random_array::<typenum::U13>(&mut rng);

        let expected = a
            .iter()
            .zip(b.iter())
            .fold(Gf2_128Field::ZERO, |acc, (x, y)| acc + *x * y);
        let actual = Gf2_128Field::dot(a, b);
        assert_eq!(expected, actual);

        assert_eq!(
            Gf2_128Field::ZERO,
            Gf2_128Field::dot(
                std::iter::empty::<Gf2_128Field>(),
                std::iter::empty::<Gf2_128Field>()
            )
        );
    }

    #[test]
    fn test_constant_time_eq() {
        let a = Gf2_128Field::from(42u64);
        assert!(bool::from(a.ct_eq(&Gf2_128Field::from(42u64))));
        assert!(!bool::from(a.ct_eq(&Gf2_128Field::from(43u64))));
    }
}