arcium-primitives 0.4.2

use core::iter::{Product, Sum};
use std::{
    cmp::Ordering,
    fmt::Debug,
    hash::{Hash, Hasher},
    marker::PhantomData,
    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    str,
};

use derive_more::{AsMut, AsRef};
use ff::Field;
use hybrid_array::{Array, ArraySize, AssocArraySize};
use itertools::{izip, Itertools};
use itybity::ToBits;
use num_traits::{One, Zero};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};
use typenum::{Prod, Unsigned, U64, U8};
use wincode::{SchemaRead, SchemaWrite};

use crate::{
    algebra::{
        field::{binary::Gf2, FieldExtension},
        ops::{AccReduce, DefaultDotProduct, DotProduct, IntoWide, MulAccReduce, ReduceWide},
        uniform_bytes::FromUniformBytes,
    },
    random::{CryptoRngCore, Random},
    types::{HeapArray, Positive},
    utils::IntoExactSizeIterator,
};

#[serde_as]
#[derive(Clone, Copy, Debug, Eq, AsRef, AsMut, Serialize, Deserialize, SchemaWrite, SchemaRead)]
#[repr(transparent)]
pub struct Gf2Ext<P: Gf2ExtParams, const LIMBS: usize> {
    #[serde_as(as = "[_; LIMBS]")]
    pub(crate) data: [u64; LIMBS],
    _id: PhantomData<P>,
}

impl<P: Gf2ExtParams, const LIMBS: usize> AsRef<[u8]> for Gf2Ext<P, LIMBS> {
    fn as_ref(&self) -> &[u8] {
        // SAFETY: This is safe because:
        // 1. We're only reading the bytes
        // 2. The memory layout of [u64; LIMBS] is well-defined
        // 3. The slice length is exactly LIMBS * 8 bytes
        unsafe {
            std::slice::from_raw_parts(
                self.data.as_ptr() as *const u8,
                LIMBS * std::mem::size_of::<u64>(),
            )
        }
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> AsMut<[u8]> for Gf2Ext<P, LIMBS> {
    fn as_mut(&mut self) -> &mut [u8] {
        // SAFETY: This is safe because:
        // 1. We have exclusive access to the bytes
        // 2. The memory layout of [u64; LIMBS] is well-defined
        // 3. The slice length is exactly LIMBS * 8 bytes
        unsafe {
            std::slice::from_raw_parts_mut(
                self.data.as_mut_ptr() as *mut u8,
                LIMBS * std::mem::size_of::<u64>(),
            )
        }
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> PartialEq<Self> for Gf2Ext<P, LIMBS> {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> PartialOrd<Self> for Gf2Ext<P, LIMBS> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Ord for Gf2Ext<P, LIMBS> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.data.cmp(&other.data)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Hash for Gf2Ext<P, LIMBS> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        let bytes: &[u8] = self.as_ref();
        bytes.iter().for_each(|x| {
            x.hash(state);
        });
    }
}

// Datatype which can store a dot product result without modulus reduction
#[derive(Clone, Copy)]
pub struct Gf2LimbsWide<const LIMBS: usize> {
    pub(crate) low: [u64; LIMBS],
    pub(crate) high: [u64; LIMBS],
}

pub trait Gf2ExtParams: Copy + Debug + Eq + PartialEq + Sync + Send + Unpin + 'static {
    /// The extension degree, a multiple of 64
    type Degree: ArraySize + Positive;

    /// Number of bytes needed for the extension, a multiple of 8
    type Bytes: ArraySize + Positive;
    //  Modulus polynomial non-zero coefficients (x^0 is not included) positions
    const POLY_MOD_ONES: &[usize];
}

impl<P: Gf2ExtParams, const LIMBS: usize> Gf2Ext<P, LIMBS> {
    pub const fn new(data: [u64; LIMBS]) -> Self {
        Self {
            data,
            _id: PhantomData,
        }
    }

    const ZERO: Self = Self::new([0u64; LIMBS]);
    const ONE: Self = Self::new({
        let mut tmp = [0u64; LIMBS];
        tmp[0] = 1;
        tmp
    });
}

impl<P: Gf2ExtParams, const LIMBS: usize> Gf2Ext<P, LIMBS> {
    pub fn as_mut_ne_bytes_slice(&mut self) -> &mut [u8] {
        bytemuck::bytes_of_mut(&mut self.data)
    }

    pub fn as_ne_bytes_slice(&self) -> &[u8] {
        bytemuck::bytes_of(&self.data)
    }

    pub fn from_limbs(val: [u64; LIMBS]) -> Self {
        Self::new(val)
    }

    pub fn from_u64(val: u64) -> Self {
        let mut tmp = [0u64; LIMBS];
        tmp[0] = val;
        Self::from_limbs(tmp)
    }

    pub fn from_u128(val: u128) -> Self {
        let mut tmp = [0u64; LIMBS];
        tmp[0] = val as u64;
        tmp[1] = (val >> 64) as u64;
        Self::from_limbs(tmp)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Default for Gf2Ext<P, LIMBS> {
    fn default() -> Self {
        Self::ZERO
    }
}

impl<const LIMBS: usize> Default for Gf2LimbsWide<LIMBS> {
    fn default() -> Self {
        Self {
            low: [0u64; LIMBS],
            high: [0u64; LIMBS],
        }
    }
}

// === Conversion traits

impl<P: Gf2ExtParams, const LIMBS: usize> From<bool> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn from(value: bool) -> Self {
        Self::from_u64(value as u64)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> From<u8> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn from(value: u8) -> Self {
        Self::from_u64(value as u64)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> From<u16> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn from(value: u16) -> Self {
        Self::from_u64(value as u64)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> From<u32> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn from(value: u32) -> Self {
        Self::from_u64(value as u64)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> From<u64> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn from(value: u64) -> Self {
        Self::from_u64(value)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> From<u128> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn from(value: u128) -> Self {
        Self::from_u128(value)
    }
}

// === Implementation of Field trait methods

impl<P: Gf2ExtParams, const LIMBS: usize> Field for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
{
    const ZERO: Self = Self::ZERO;
    const ONE: Self = Self::ONE;

    fn random(mut rng: impl RngCore) -> Self {
        let mut tmp = Self::default();
        rng.fill_bytes(tmp.as_mut_ne_bytes_slice());
        tmp
    }

    fn square(&self) -> Self {
        self * self
    }

    fn double(&self) -> Self {
        self + self
    }

    fn invert(&self) -> CtOption<Self> {
        unimplemented!()
    }

    fn sqrt_ratio(_num: &Self, _div: &Self) -> (Choice, Self) {
        unimplemented!()
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> ConditionallySelectable for Gf2Ext<P, LIMBS> {
    #[inline]
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self::from_limbs(std::array::from_fn(|k| {
            u64::conditional_select(&a.data[k], &b.data[k], choice)
        }))
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> ConstantTimeEq for Gf2Ext<P, LIMBS> {
    fn ct_eq(&self, other: &Self) -> Choice {
        izip!(self.data, other.data).fold(1u8.into(), |r, ab| r & ab.0.ct_eq(&ab.1))
    }
}

// === Implementation of arithmetic operators === //

// Negation

#[macros::op_variants(borrowed)]
impl<P: Gf2ExtParams, const LIMBS: usize> Neg for Gf2Ext<P, LIMBS> {
    type Output = Gf2Ext<P, LIMBS>;

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

// Addition

#[macros::op_variants(owned, borrowed, flipped_commutative)]
impl<P: Gf2ExtParams, const LIMBS: usize> Add<&Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS> {
    type Output = Gf2Ext<P, LIMBS>;

    #[inline]
    #[allow(clippy::op_ref)]
    fn add(mut self, rhs: &Self::Output) -> Self::Output {
        self += rhs;
        self
    }
}

#[macros::op_variants(owned)]
impl<P: Gf2ExtParams, const LIMBS: usize> AddAssign<&Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS> {
    #[allow(clippy::suspicious_op_assign_impl)]
    fn add_assign(&mut self, rhs: &Self) {
        izip!(&mut self.data, rhs.data).for_each(|(a, b)| *a ^= b);
    }
}

// Subtraction

#[macros::op_variants(owned, borrowed, flipped)]
impl<P: Gf2ExtParams, const LIMBS: usize> Sub<&Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS> {
    type Output = Gf2Ext<P, LIMBS>;

    #[inline]
    #[allow(clippy::suspicious_arithmetic_impl)]
    fn sub(self, rhs: &Self::Output) -> Self::Output {
        self + rhs
    }
}

#[macros::op_variants(owned)]
impl<P: Gf2ExtParams, const LIMBS: usize> SubAssign<&Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS> {
    #[inline]
    #[allow(clippy::suspicious_op_assign_impl)]
    fn sub_assign(&mut self, rhs: &Self) {
        *self += rhs;
    }
}

// Multiplication

#[macros::op_variants(owned, borrowed)]
impl<P: Gf2ExtParams, const LIMBS: usize> Mul<&Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
{
    type Output = Gf2Ext<P, LIMBS>;

    #[inline]
    fn mul(mut self, rhs: &Self::Output) -> Self::Output {
        self *= rhs;
        self
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Mul<Gf2Ext<P, LIMBS>> for &Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
{
    type Output = Gf2Ext<P, LIMBS>;

    #[inline]
    fn mul(self, rhs: Self::Output) -> Self::Output {
        rhs * self
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> MulAssign<&Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulAssign,
{
    #[inline]
    fn mul_assign(&mut self, rhs: &Gf2Ext<P, LIMBS>) {
        *self *= *rhs;
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> MulAssign<Gf2> for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulAssign,
{
    #[inline]
    fn mul_assign(&mut self, rhs: Gf2) {
        self.data
            .iter_mut()
            .for_each(|limb: &mut u64| *limb *= rhs.0 as u64)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> MulAssign<&Gf2> for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulAssign,
{
    #[inline]
    fn mul_assign(&mut self, rhs: &Gf2) {
        self.data
            .iter_mut()
            .for_each(|limb: &mut u64| *limb *= rhs.0 as u64)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> MulAssign for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
    Gf2Ext<P, LIMBS>: IntoWide<Gf2LimbsWide<LIMBS>>,
{
    fn mul_assign(&mut self, rhs: Self) {
        let res_wide = self.mul_wide(rhs);
        *self = Self::reduce_mod_order(res_wide)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Sum for Gf2Ext<P, LIMBS> {
    #[inline]
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Gf2Ext::ZERO, |a, b| a + b)
    }
}

impl<'a, P: Gf2ExtParams, const LIMBS: usize> Sum<&'a Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.fold(Gf2Ext::ZERO, |a, b| a + b)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Product for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
{
    #[inline]
    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Gf2Ext::ONE, |a, b| a * b)
    }
}

impl<'a, P: Gf2ExtParams, const LIMBS: usize> Product<&'a Gf2Ext<P, LIMBS>> for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
{
    #[inline]
    fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
        iter.fold(Gf2Ext::ONE, |a, b| a * b)
    }
}

// === Implementation of FieldExtension trait methods
impl<P: Gf2ExtParams, const LIMBS: usize> FieldExtension for Gf2Ext<P, LIMBS>
where
    Gf2Ext<P, LIMBS>: MulWide<Output = Gf2LimbsWide<LIMBS>>,
    [u8; LIMBS]: AssocArraySize,
    <[u8; LIMBS] as AssocArraySize>::Size: ArraySize + Positive + Mul<U8> + Mul<U64>,
    Prod<<[u8; LIMBS] as AssocArraySize>::Size, U8>: ArraySize + Positive,
    Prod<<[u8; LIMBS] as AssocArraySize>::Size, U64>: ArraySize + Positive,
{
    type Subfield = Gf2;

    type Degree = P::Degree;
    type FieldBitSize = Prod<<[u8; LIMBS] as AssocArraySize>::Size, U64>;
    type FieldBytesSize = Prod<<[u8; LIMBS] as AssocArraySize>::Size, U8>;

    fn to_subfield_elements(&self) -> impl ExactSizeIterator<Item = Self::Subfield> {
        self.data.iter_lsb0().take(P::Degree::USIZE).map(Gf2::from)
    }

    fn from_subfield_elements(elems: &[Self::Subfield]) -> Option<Self> {
        let len = elems.len();
        if len == 0 || len > P::Degree::to_usize() {
            return None;
        }
        let mut data = [0u64; LIMBS];
        for (i, elem) in elems.iter().enumerate() {
            let limb_idx = i / 64;
            let bit_idx = i % 64;
            data[limb_idx] |= (bool::from(*elem) as u64) << bit_idx;
        }
        Some(Self::new(data))
    }

    fn to_le_bytes(&self) -> Array<u8, Self::FieldBytesSize> {
        Array::from_iter(
            (0..LIMBS)
                .cartesian_product((0..64).step_by(8))
                .map(|(limb, shift)| ((self.data[limb] >> shift) & 0xFF) as u8),
        )
    }

    fn from_le_bytes(bytes: &[u8]) -> Option<Self> {
        if bytes.len() != Self::FieldBytesSize::USIZE {
            return None;
        }

        let mut it = bytes.chunks_exact(8).take(LIMBS);
        Some(Self::new(std::array::from_fn(|_| {
            u64::from_le_bytes(it.next().unwrap().try_into().unwrap())
        })))
    }

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

    fn generator() -> Self {
        Self::from_u64(2u64) // encoding of variable X
    }

    // Fast linear-orthomorphism for GF2 extensions of even degrees from "Minimizing the Two-Round
    // Even-Mansour Cipher" by Chen et al.
    fn linear_orthomorphism(&self) -> Self {
        let mut res = Self::default();
        for i in 0..LIMBS / 2 {
            res.data[i] = self.data[i] ^ self.data[LIMBS - i - 1];
            res.data[LIMBS - i - 1] = self.data[i];
        }
        if LIMBS % 2 == 1 {
            let k = LIMBS / 2;
            let (tl, th) = (self.data[k] as u32, (self.data[k] >> 32) as u32);

            let (rl, rh) = (tl ^ th, tl);
            res.data[k] = rl as u64 + ((rh as u64) << 32);
        }

        res
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Random for Gf2Ext<P, LIMBS> {
    fn random(mut rng: impl CryptoRngCore) -> Self {
        let mut tmp = Self::default();
        rng.fill_bytes(tmp.as_mut_ne_bytes_slice());
        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);
        HeapArray::from_box_bytes(buf)
    }
}

unsafe impl<P: Gf2ExtParams, const LIMBS: usize> bytemuck::Zeroable for Gf2Ext<P, LIMBS> {
    fn zeroed() -> Self {
        Self::ZERO
    }
}
unsafe impl<P: Gf2ExtParams, const LIMBS: usize> bytemuck::Pod for Gf2Ext<P, LIMBS> {}

impl<P: Gf2ExtParams, const LIMBS: usize> Mul<Gf2> for Gf2Ext<P, LIMBS> {
    type Output = Self;

    #[allow(clippy::suspicious_arithmetic_impl)]
    fn mul(mut self, rhs: Gf2) -> Self::Output {
        // Transform 0/1 value into a bit-mask 00..0/11..1
        let m = (-(rhs.0 as i64)) as u64;
        self.data.iter_mut().for_each(|v: &mut u64| *v &= m);
        self
    }
}

impl<'a, P: Gf2ExtParams, const LIMBS: usize> Mul<&'a Gf2> for Gf2Ext<P, LIMBS> {
    type Output = Self;

    #[inline]
    fn mul(self, rhs: &'a Gf2) -> Self::Output {
        self * *rhs
    }
}

// Lazy multiplication and modulus reduction

/// Add gf2 polynomials encoded in Gf2LimbsWide<LIMBS>
impl<const LIMBS: usize> AddAssign for Gf2LimbsWide<LIMBS> {
    fn add_assign(&mut self, rhs: Self) {
        izip!(&mut self.low, rhs.low).for_each(|(a, b)| *a ^= b);
        izip!(&mut self.high, rhs.high).for_each(|(a, b)| *a ^= b);
    }
}

// Wide type for Gf2Ext x Gf2Ext multiplication
impl<P: Gf2ExtParams, const LIMBS: usize> IntoWide<Gf2LimbsWide<LIMBS>> for Gf2Ext<P, LIMBS> {
    #[inline]
    fn to_wide(&self) -> Gf2LimbsWide<LIMBS> {
        Gf2LimbsWide {
            low: self.data,
            high: [0u64; LIMBS],
        }
    }

    #[inline]
    fn zero_wide() -> Gf2LimbsWide<LIMBS> {
        Default::default()
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> ReduceWide<Gf2LimbsWide<LIMBS>> for Gf2Ext<P, LIMBS> {
    /// Reduce a GF2 polynomial of 2*LIMBS down to LIMBS.
    /// * `poly` - polynomial to reduce
    fn reduce_mod_order(a: Gf2LimbsWide<LIMBS>) -> Self {
        let Gf2LimbsWide { mut low, mut high } = a;

        let ones = P::POLY_MOD_ONES;

        // Use a macro (instead of a function) to avoid rust borrowing rules in case `$inp` aliases
        // with `$out_high`
        macro_rules! reduce_1step {
            ($out_low:expr, $out_high:expr, $inp:expr) => {
                for k in ones {
                    $out_high ^= $inp >> (64 - k);
                }
                $out_low ^= $inp;
                for k in ones {
                    $out_low ^= $inp << k;
                }
            };
        }

        reduce_1step!(low[LIMBS - 1], high[0], high[LIMBS - 1]);
        for i in (0..LIMBS - 1).rev() {
            reduce_1step!(low[i], low[i + 1], high[i]);
        }

        Gf2Ext {
            data: low,
            _id: PhantomData,
        }
    }
}

// Dot product: Gf2Ext x Gf2Ext
impl<P: Gf2ExtParams, const LIMBS: usize> MulAccReduce for Gf2Ext<P, LIMBS>
where
    Self: MulWide<Output = Gf2LimbsWide<LIMBS>>,
{
    type WideType = Gf2LimbsWide<LIMBS>;

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

impl<P: Gf2ExtParams, const LIMBS: usize> DefaultDotProduct for Gf2Ext<P, LIMBS> where
    Self: MulAccReduce
{
}

// Dot product: Gf2Ext x &Gf2Ext
impl<'a, P: Gf2ExtParams, const LIMBS: usize> DotProduct<Self, &'a Self> for Gf2Ext<P, LIMBS>
where
    Self: DefaultDotProduct,
{
    #[inline]
    fn dot<I1, I2>(a: I1, b: I2) -> Self
    where
        I1: IntoExactSizeIterator<Item = Self>,
        I2: IntoExactSizeIterator<Item = &'a Self>,
    {
        Self::dot(a, b.into_iter().copied())
    }
}

// Dot product: &Gf2Ext x &Gf2Ext
impl<'a, 'b, P: Gf2ExtParams, const LIMBS: usize> DotProduct<&'a Self, &'b Self>
    for Gf2Ext<P, LIMBS>
where
    Self: DefaultDotProduct,
{
    #[inline]
    fn dot<I1, I2>(a: I1, b: I2) -> Self
    where
        I1: IntoExactSizeIterator<Item = &'a Self>,
        I2: IntoExactSizeIterator<Item = &'b Self>,
    {
        Self::dot(a.into_iter().copied(), b)
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> IntoWide for Gf2Ext<P, LIMBS> {
    #[inline]
    fn to_wide(&self) -> Self {
        *self
    }

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

impl<P: Gf2ExtParams, const LIMBS: usize> ReduceWide for Gf2Ext<P, LIMBS> {
    #[inline]
    fn reduce_mod_order(a: Self) -> Self {
        a
    }
}

// Dot product : Gf2Ext, Gf2
impl<P: Gf2ExtParams, const LIMBS: usize> MulAccReduce<Self, Gf2> for Gf2Ext<P, LIMBS> {
    type WideType = Self;

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

impl<P: Gf2ExtParams, const LIMBS: usize> DefaultDotProduct<Self, Gf2> for Gf2Ext<P, LIMBS> {}

// Dot product : &Gf2Ext, Gf2
impl<'a, P: Gf2ExtParams, const LIMBS: usize> MulAccReduce<&'a Self, Gf2> for Gf2Ext<P, LIMBS> {
    type WideType = Self;

    #[inline]
    fn mul_acc(acc: &mut Self, a: &'a Self, b: Gf2) {
        Self::mul_acc(acc, *a, b);
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> DefaultDotProduct<&Self, Gf2> for Gf2Ext<P, LIMBS> {}

// Dot product : &Gf2Ext, &Gf2
impl<'a, 'b, P: Gf2ExtParams, const LIMBS: usize> MulAccReduce<&'a Self, &'b Gf2>
    for Gf2Ext<P, LIMBS>
{
    type WideType = Self;

    #[inline]
    fn mul_acc(acc: &mut Self, a: &'a Self, b: &'b Gf2) {
        Self::mul_acc(acc, a, *b);
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> DefaultDotProduct<&Self, &Gf2> for Gf2Ext<P, LIMBS> {}

impl<P: Gf2ExtParams, const LIMBS: usize> AccReduce for Gf2Ext<P, LIMBS> {
    type WideType = Self;

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

impl<P: Gf2ExtParams, const LIMBS: usize> AccReduce<&Self> for Gf2Ext<P, LIMBS> {
    type WideType = Self;

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

pub trait MulWide<Rhs = Self> {
    type Output;
    fn mul_wide(self, rhs: Rhs) -> Self::Output;
}

impl<P: Gf2ExtParams, const LIMBS: usize> FromUniformBytes for Gf2Ext<P, LIMBS>
where
    [u8; LIMBS]: AssocArraySize,
    <[u8; LIMBS] as AssocArraySize>::Size: ArraySize + Positive + Mul<U8> + Mul<U64>,
    Prod<<[u8; LIMBS] as AssocArraySize>::Size, U8>: ArraySize + Positive,
{
    type UniformBytes = Prod<<[u8; LIMBS] as AssocArraySize>::Size, U8>;

    fn from_uniform_bytes(a: &Array<u8, Self::UniformBytes>) -> Self {
        let mut it = a.chunks_exact(8).take(LIMBS);
        Self::new(std::array::from_fn(|_| {
            u64::from_le_bytes(it.next().unwrap().try_into().unwrap())
        }))
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> Zero for Gf2Ext<P, LIMBS> {
    fn zero() -> Self {
        Self::ZERO
    }

    fn is_zero(&self) -> bool {
        self.ct_eq(&Self::ZERO).into()
    }
}

impl<P: Gf2ExtParams, const LIMBS: usize> One for Gf2Ext<P, LIMBS>
where
    Self: Mul<Output = Self>,
{
    fn one() -> Self {
        Self::ONE
    }

    fn is_one(&self) -> bool {
        self.ct_eq(&Self::ONE).into()
    }
}

// statically dispatch size-dependent MulWide implementation
macro_rules! impl_wide_mul {
    ($params:ty, 1) => {
        impl $crate::algebra::field::binary::gf2_ext::MulWide for Gf2Ext<$params, 1> {
            type Output = Gf2LimbsWide<1>;
            fn mul_wide(self, rhs: Self) -> Self::Output {
                let (low, high) =
                    $crate::algebra::ops::clmul::carry_less_mul_1limb(self.data, rhs.data);
                Gf2LimbsWide { low, high }
            }
        }
    };
    ($params:ty, 2) => {
        impl $crate::algebra::field::binary::gf2_ext::MulWide for Gf2Ext<$params, 2> {
            type Output = Gf2LimbsWide<2>;
            fn mul_wide(self, rhs: Self) -> Self::Output {
                let (low, high) =
                    $crate::algebra::ops::clmul::carry_less_mul_2limbs(self.data, rhs.data);
                Gf2LimbsWide { low, high }
            }
        }
    };
    ($params:ty, $limbs:literal) => {
        impl $crate::algebra::field::binary::gf2_ext::MulWide for Gf2Ext<$params, $limbs> {
            type Output = Gf2LimbsWide<$limbs>;
            fn mul_wide(self, rhs: Self) -> Self::Output {
                let (low, high) = $crate::algebra::ops::clmul::carry_less_mul(self.data, rhs.data);
                Gf2LimbsWide { low, high }
            }
        }
    };
}
pub(crate) use impl_wide_mul;

macro_rules! validate_modulus_poly_ones {
    ([$($val:literal),+ $(,)?]) => ($(
        static_assertions::const_assert!(($val < 64) & ($val > 0));
    )*);
}
pub(crate) use validate_modulus_poly_ones;

macro_rules! define_gf2_extension {
    ($ext_name:ident, $ext_params:ident, $limbs:tt, $modulus_poly_ones:tt) => {
        $crate::algebra::field::binary::gf2_ext::validate_modulus_poly_ones!($modulus_poly_ones);

        #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
        pub struct $ext_params;

        pub type $ext_name = $crate::algebra::field::binary::gf2_ext::Gf2Ext<$ext_params, $limbs>;

        // Note: do not move to generic implementation
        impl $ext_name {
            /// Casts the field element to a core byte array
            pub fn into_ne_bytes_array(self) -> [u8; { $limbs * 8 }] {
                bytemuck::cast(self.data)
            }

            pub fn from_ne_bytes_array(arr: [u8; { $limbs * 8 }]) -> Self {
                Self::new(bytemuck::cast(arr))
            }
        }

        $crate::algebra::field::binary::gf2_ext::impl_wide_mul!($ext_params, $limbs);

        impl $crate::algebra::field::binary::gf2_ext::Gf2ExtParams for $ext_params {
            type Degree = typenum::U<{ $limbs * 64 }>;
            type Bytes = typenum::U<{ $limbs * 8 }>;
            const POLY_MOD_ONES: &[usize] = &{ $modulus_poly_ones };
        }
    };
}

pub(crate) use define_gf2_extension;

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

    const N_TESTS: usize = 1;

    // sage instruction used to find minimal weight modulus polynomials:
    //  GF(2^128, modulus="minimal_weight", name="x").polynomial()
    define_gf2_extension!(Gf2_128, Gf2_128Params, 2, [1, 2, 7]);
    define_gf2_extension!(Gf2_192, Gf2_192Params, 3, [1, 2, 7]);
    define_gf2_extension!(Gf2_256, Gf2_256Params, 4, [2, 5, 10]);
    define_gf2_extension!(Gf2_64, Gf2_64Params, 1, [1, 3, 4]);

    #[test]
    fn test_gf2_ext_operations() {
        fn gf2_operations_test_case<F: FieldExtension>() {
            let mut rng = crate::random::test_rng();
            for _ in 0..N_TESTS {
                let a: F = Random::random(&mut rng);
                // a + a = 0
                assert_eq!(a + a, F::ZERO);

                // a * 234 = 234 * a
                assert_eq!(a * F::from(234u64), F::from(234u64) * a);

                // a * (a + 234) = a * a + 234 * a
                assert_eq!(a * (a + F::from(234u64)), a * a + F::from(234u64) * a);

                // a exp (2^degree) = a
                let mut cumul = a;
                for _ in 0..(F::Degree::to_usize()) {
                    cumul *= cumul;
                }
                assert_eq!(a, cumul);
            }
        }

        gf2_operations_test_case::<Gf2_128>();
        gf2_operations_test_case::<Gf2_256>();
        gf2_operations_test_case::<Gf2_192>();
        gf2_operations_test_case::<Gf2_64>();
    }

    #[test]
    fn test_gf2_128_prod() {
        macro_rules! gf2_128_prod_test_case {
            ($aval:expr, $bval:expr, $prod_red:expr) => {{
                let a: [u64; 2] = $aval;
                let b: [u64; 2] = $bval;
                let prod_red: [u64; 2] = $prod_red;

                {
                    let ae = Gf2_128::from_limbs(a);
                    let be = Gf2_128::from_limbs(b);
                    let prod_red_comp = ae * be;
                    assert_eq!(prod_red_comp, Gf2_128::from_limbs(prod_red));
                }
            }};
        }

        gf2_128_prod_test_case!(
            [0x9f418f3bffd84bba, 0x4a7c605645afdfb1],
            [0x80b7bd91cddc5be5, 0x3a97291035e41e1f],
            [0x46ca0b600a32c5f7, 0x823a605e0452082a]
        );

        gf2_128_prod_test_case!(
            [0x74ef862bc1b6d333, 0x3a88103b80d97b73],
            [0x753f4846eb020b5a, 0x8f108359ea25fa8f],
            [0x6947ab52b94f0ef9, 0xb2ec1b5a4553aa6d]
        );

        gf2_128_prod_test_case!(
            [0x6447b3dcaed62649, 0x6e4af40b2ee1b4c1],
            [0xbd7a4e12fdb29840, 0x8950f56742015f25],
            [0x38ae5eb860021fe9, 0x6f18457f05ac2506]
        );
    }
}