arcium-primitives 0.4.2

use std::{
    fmt::{Display, Formatter, Result as FmtResult},
    iter::{Product, Sum},
    ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign},
};

use derive_more::derive::{AsMut, AsRef};
use ff::Field;
use hybrid_array::Array;
use num_traits::{One, Zero};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
use wincode::{SchemaRead, SchemaWrite};

use crate::{
    algebra::{
        field::{
            binary::Gf2_128,
            mersenne::Mersenne107,
            ByteSize,
            FieldElement,
            FieldExtension,
            PrimeFieldExtension,
        },
        ops::{AccReduce, DefaultDotProduct, IntoWide, MulAccReduce, ReduceWide},
        uniform_bytes::FromUniformBytes,
    },
    errors::PrimitiveError,
    random::{CryptoRngCore, Random},
    sharing::unauthenticated::AdditiveShares,
    types::HeapArray,
};

pub type Bit = SubfieldElement<Gf2_128>;
pub type Bits<N> = HeapArray<Bit, N>;

pub type Mersenne107Element = SubfieldElement<Mersenne107>;
pub type Mersenne107Elements<N> = HeapArray<Mersenne107Element, N>;

/// A subfield element wrapper.
#[derive(
    Copy, Clone, Debug, PartialOrd, PartialEq, Eq, Hash, AsRef, AsMut, SchemaRead, SchemaWrite,
)]
#[repr(transparent)]
pub struct SubfieldElement<F: FieldExtension>(pub F::Subfield);

impl<F: FieldExtension> SubfieldElement<F> {
    /// Construct a subfield element from an inner subfield element
    #[inline]
    pub fn new(inner: F::Subfield) -> Self {
        SubfieldElement(inner)
    }

    /// Get the inner value of the subfield element
    #[inline]
    pub fn inner(&self) -> F::Subfield {
        self.0
    }

    /// Compute the exponentiation of the given subfield element
    #[inline]
    pub fn pow<S: AsRef<[u64]>>(&self, exp: S) -> Self {
        SubfieldElement::new(self.0.pow(exp))
    }

    /// Construct a subfield element from the given bytes
    #[inline]
    pub fn from_be_bytes(bytes: &[u8]) -> Result<SubfieldElement<F>, PrimitiveError> {
        let mut bytes = bytes.to_vec();
        bytes.reverse();
        Ok(SubfieldElement(
            F::Subfield::from_le_bytes(&bytes).ok_or_else(|| {
                PrimitiveError::DeserializationFailed(
                    "Invalid subfield element encoding".to_string(),
                )
            })?,
        ))
    }

    pub fn from_le_bytes(bytes: &[u8]) -> Result<SubfieldElement<F>, PrimitiveError> {
        Ok(SubfieldElement(
            F::Subfield::from_le_bytes(bytes).ok_or_else(|| {
                PrimitiveError::DeserializationFailed(
                    "Invalid subfield element encoding".to_string(),
                )
            })?,
        ))
    }

    /// Convert the subfield element to little-endian bytes
    #[inline]
    pub fn to_le_bytes(&self) -> Array<u8, ByteSize<F::Subfield>> {
        self.0.to_le_bytes()
    }

    /// Convert the subfield element to big-endian bytes
    #[inline]
    pub fn to_be_bytes(&self) -> Array<u8, ByteSize<F::Subfield>> {
        let mut rev = self.0.to_le_bytes();
        rev.as_mut().reverse();
        rev
    }

    pub fn to_biguint(&self) -> num_bigint::BigUint {
        num_bigint::BigUint::from_bytes_le(self.to_le_bytes().as_ref())
    }

    pub fn to_bigint(&self) -> num_bigint::BigInt {
        self.to_biguint().into()
    }
}

impl<F: FieldExtension> Random for SubfieldElement<F> {
    fn random(mut rng: impl CryptoRngCore) -> Self {
        SubfieldElement(Random::random(&mut rng))
    }
}

impl<F: FieldExtension> Default for SubfieldElement<F> {
    fn default() -> Self {
        Self::zero()
    }
}

impl<F: FieldExtension> Serialize for SubfieldElement<F> {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let bytes = self
            .to_le_bytes()
            .into_iter()
            .collect::<Array<u8, ByteSize<F::Subfield>>>();
        serde_bytes::serialize(AsRef::<[u8]>::as_ref(&bytes), serializer)
    }
}

impl<'de, F: FieldExtension> Deserialize<'de> for SubfieldElement<F> {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let bytes: &[u8] = serde_bytes::deserialize(deserializer)?;
        let field_elem = SubfieldElement::from_le_bytes(bytes).map_err(|err| {
            serde::de::Error::custom(format!("Failed to deserialize field element: {err:?}"))
        })?;
        Ok(field_elem)
    }
}

impl<F: FieldExtension> Display for SubfieldElement<F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        write!(f, "{self:?}")
    }
}

// --------------
// | Arithmetic |
// --------------

// === Addition === //

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<F: FieldExtension> Add<&SubfieldElement<F>> for SubfieldElement<F> {
    type Output = SubfieldElement<F>;

    #[inline]
    fn add(self, rhs: &SubfieldElement<F>) -> Self::Output {
        SubfieldElement(self.0 + rhs.0)
    }
}

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

// === Subtraction === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<F: FieldExtension> Sub<&SubfieldElement<F>> for SubfieldElement<F> {
    type Output = SubfieldElement<F>;

    #[inline]
    fn sub(self, rhs: &SubfieldElement<F>) -> Self::Output {
        SubfieldElement(self.0 - rhs.0)
    }
}

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

// === Multiplication === //

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<F: FieldExtension> Mul<&SubfieldElement<F>> for SubfieldElement<F> {
    type Output = SubfieldElement<F>;

    #[inline]
    fn mul(self, rhs: &SubfieldElement<F>) -> Self::Output {
        SubfieldElement(self.0 * rhs.0)
    }
}

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<F: FieldExtension> Mul<&FieldElement<F>> for SubfieldElement<F> {
    type Output = FieldElement<F>;

    #[inline]
    fn mul(self, rhs: &FieldElement<F>) -> Self::Output {
        FieldElement(rhs.0 * self.0)
    }
}

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

// === Negation === //

#[macros::op_variants(borrowed)]
impl<F: FieldExtension> Neg for SubfieldElement<F> {
    type Output = SubfieldElement<F>;

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

// === Division === //

#[macros::op_variants(owned, borrowed, flipped)]
impl<F: FieldExtension> Div<&SubfieldElement<F>> for SubfieldElement<F> {
    type Output = CtOption<SubfieldElement<F>>;

    #[inline]
    fn div(self, rhs: &SubfieldElement<F>) -> Self::Output {
        rhs.0.invert().map(|inv| SubfieldElement(self.0 * inv))
    }
}

// === Equality === //

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

impl<F: FieldExtension> ConditionallySelectable for SubfieldElement<F> {
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        let selected = F::Subfield::conditional_select(&a.0, &b.0, choice);
        SubfieldElement(selected)
    }
}

// === Other === //

impl<F: FieldExtension> AdditiveShares for SubfieldElement<F> {}

// ---------------
// | Conversions |
// ---------------

impl<F: FieldExtension> From<bool> for SubfieldElement<F> {
    #[inline]
    fn from(value: bool) -> Self {
        SubfieldElement(F::Subfield::from(value as u64))
    }
}

impl<F: FieldExtension> From<u8> for SubfieldElement<F> {
    #[inline]
    fn from(value: u8) -> Self {
        SubfieldElement(F::Subfield::from(value as u64))
    }
}

impl<F: FieldExtension> From<u16> for SubfieldElement<F> {
    #[inline]
    fn from(value: u16) -> Self {
        SubfieldElement(F::Subfield::from(value as u64))
    }
}

impl<F: FieldExtension> From<u32> for SubfieldElement<F> {
    #[inline]
    fn from(value: u32) -> Self {
        SubfieldElement(F::Subfield::from(value as u64))
    }
}

impl<F: FieldExtension> From<u64> for SubfieldElement<F> {
    #[inline]
    fn from(value: u64) -> Self {
        SubfieldElement(F::Subfield::from(value))
    }
}

impl<F: FieldExtension> From<u128> for SubfieldElement<F> {
    #[inline]
    fn from(value: u128) -> Self {
        SubfieldElement(F::Subfield::from(value))
    }
}

impl<F: PrimeFieldExtension> From<FieldElement<F>> for SubfieldElement<F> {
    #[inline]
    fn from(value: FieldElement<F>) -> Self {
        SubfieldElement(value.0)
    }
}

// -------------------
// | Iterator Traits |
// -------------------

impl<F: FieldExtension> Sum for SubfieldElement<F> {
    #[inline]
    fn sum<I: Iterator<Item = SubfieldElement<F>>>(iter: I) -> Self {
        let tmp = iter.fold(<F::Subfield as AccReduce>::zero_wide(), |mut acc, x| {
            F::Subfield::acc(&mut acc, x.0);
            acc
        });
        SubfieldElement(F::Subfield::reduce_mod_order(tmp))
    }
}

impl<'a, F: FieldExtension> Sum<&'a SubfieldElement<F>> for SubfieldElement<F> {
    #[inline]
    fn sum<I: Iterator<Item = &'a SubfieldElement<F>>>(iter: I) -> Self {
        let tmp = iter.fold(<F::Subfield as AccReduce>::zero_wide(), |mut acc, x| {
            F::Subfield::acc(&mut acc, x.0);
            acc
        });
        SubfieldElement(F::Subfield::reduce_mod_order(tmp))
    }
}

impl<F: FieldExtension> Product for SubfieldElement<F> {
    #[inline]
    fn product<I: Iterator<Item = SubfieldElement<F>>>(iter: I) -> Self {
        iter.fold(SubfieldElement::one(), |acc, x| acc * x)
    }
}

impl<'a, F: FieldExtension> Product<&'a SubfieldElement<F>> for SubfieldElement<F> {
    #[inline]
    fn product<I: Iterator<Item = &'a SubfieldElement<F>>>(iter: I) -> Self {
        iter.fold(SubfieldElement::one(), |acc, x| acc * x)
    }
}

impl<F: FieldExtension> FromUniformBytes for SubfieldElement<F> {
    type UniformBytes = <F::Subfield as FromUniformBytes>::UniformBytes;

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

// Dot product: Subfield<F> x Subfield<F>
impl<F: FieldExtension> IntoWide<<F::Subfield as MulAccReduce>::WideType> for SubfieldElement<F> {
    #[inline]
    fn to_wide(&self) -> <F::Subfield as MulAccReduce>::WideType {
        <F::Subfield as MulAccReduce>::to_wide(&self.0)
    }

    #[inline]
    fn zero_wide() -> <F::Subfield as MulAccReduce>::WideType {
        <F::Subfield as MulAccReduce>::zero_wide()
    }
}

impl<F: FieldExtension> ReduceWide<<F::Subfield as MulAccReduce>::WideType> for SubfieldElement<F> {
    #[inline]
    fn reduce_mod_order(a: <F::Subfield as MulAccReduce>::WideType) -> Self {
        Self(F::Subfield::reduce_mod_order(a))
    }
}

impl<F: FieldExtension> MulAccReduce for SubfieldElement<F> {
    type WideType = <F::Subfield as MulAccReduce>::WideType;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: Self, b: Self) {
        F::Subfield::mul_acc(acc, a.0, b.0);
    }
}

impl<F: FieldExtension> DefaultDotProduct for SubfieldElement<F> {}

// Dot product: &Subfield<F> x Subfield<F>
impl<'a, F: FieldExtension> MulAccReduce<&'a Self, Self> for SubfieldElement<F> {
    type WideType = <F::Subfield as MulAccReduce>::WideType;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: &'a Self, b: Self) {
        F::Subfield::mul_acc(acc, a.0, b.0);
    }
}

impl<F: FieldExtension> DefaultDotProduct<&Self, Self> for SubfieldElement<F> {}

// Dot product: Subfield<F> x &Subfield<F>
impl<'a, F: FieldExtension> MulAccReduce<Self, &'a Self> for SubfieldElement<F> {
    type WideType = <F::Subfield as MulAccReduce>::WideType;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: Self, b: &'a Self) {
        F::Subfield::mul_acc(acc, a.0, b.0);
    }
}

impl<F: FieldExtension> DefaultDotProduct<Self, &Self> for SubfieldElement<F> {}

// Dot product: &Subfield<F> x &Subfield<F>
impl<'a, 'b, F: FieldExtension> MulAccReduce<&'a Self, &'b Self> for SubfieldElement<F> {
    type WideType = <F::Subfield as MulAccReduce>::WideType;

    #[inline]
    fn mul_acc(acc: &mut Self::WideType, a: &'a Self, b: &'b Self) {
        F::Subfield::mul_acc(acc, a.0, b.0);
    }
}

impl<F: FieldExtension> DefaultDotProduct<&Self, &Self> for SubfieldElement<F> {}

// ----------------
// | Zero and One |
// ----------------

impl<F: FieldExtension> Zero for SubfieldElement<F> {
    /// The subfield's additive identity
    fn zero() -> Self {
        SubfieldElement(F::Subfield::ZERO)
    }

    fn is_zero(&self) -> bool {
        self.0.is_zero().into()
    }
}

impl<F: FieldExtension> One for SubfieldElement<F> {
    /// The subfield's multiplicative identity
    fn one() -> Self {
        SubfieldElement(F::Subfield::ONE)
    }
}

// ---------------
// | Field trait |
// ---------------

impl<F: FieldExtension> Field for SubfieldElement<F> {
    const ZERO: Self = SubfieldElement(F::Subfield::ZERO);
    const ONE: Self = SubfieldElement(F::Subfield::ONE);

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

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

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

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

    fn sqrt_ratio(num: &Self, div: &Self) -> (Choice, Self) {
        let (choice, sqrt) = F::Subfield::sqrt_ratio(&num.0, &div.0);
        (choice, SubfieldElement(sqrt))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_subfield_mersenne107_wincode() {
        let elem = SubfieldElement::<Mersenne107>::new(Mersenne107::from(42u64));
        let bytes = wincode::serialize(&elem).unwrap();

        let decoded: SubfieldElement<Mersenne107> = wincode::deserialize(&bytes).unwrap();

        assert_eq!(elem, decoded);
    }
}