p3_field/extension/

packed_binomial_extension.rs

use alloc::vec::Vec;
use core::array;
use core::iter::{Product, Sum};
use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};

use p3_util::{flatten_to_base, reconstitute_from_base};
use rand::distr::{Distribution, StandardUniform};

use super::{BinomialExtensionField, PackedExtField, binomial_mul, vector_add, vector_sub};
use crate::extension::{Binomial, BinomiallyExtendable, binomial_square};
use crate::{
    Algebra, BasedVectorSpace, Dup, Field, PackedField, PackedFieldExtension, PackedValue, Powers,
    PrimeCharacteristicRing, field_to_array,
};

/// Packed binomial extension field `F[X] / (X^D - W)`, one element per SIMD lane.
///
/// Type alias for the unified [`PackedExtField`] with `Shape = Binomial<F>`.
pub type PackedBinomialExtensionField<F, PF, const D: usize> =
    PackedExtField<F, PF, D, Binomial<F>>;

impl<F: Field, PF: PackedField<Scalar = F>, const D: usize> Default
    for PackedBinomialExtensionField<F, PF, D>
{
    #[inline]
    fn default() -> Self {
        Self::new(array::from_fn(|_| PF::ZERO))
    }
}

impl<F: Field, PF: PackedField<Scalar = F>, const D: usize> From<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
{
    #[inline]
    fn from(x: BinomialExtensionField<F, D>) -> Self {
        Self::new(x.value.map(Into::into))
    }
}

impl<F: Field, PF: PackedField<Scalar = F>, const D: usize> From<PF>
    for PackedBinomialExtensionField<F, PF, D>
{
    #[inline]
    fn from(x: PF) -> Self {
        Self::new(field_to_array(x))
    }
}

impl<F: Field, PF: PackedField<Scalar = F>, const D: usize>
    Distribution<PackedBinomialExtensionField<F, PF, D>> for StandardUniform
where
    Self: Distribution<PF>,
{
    #[inline]
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> PackedBinomialExtensionField<F, PF, D> {
        PackedBinomialExtensionField::new(array::from_fn(|_| self.sample(rng)))
    }
}

impl<F: BinomiallyExtendable<D>, PF: PackedField<Scalar = F>, const D: usize>
    Algebra<BinomialExtensionField<F, D>> for PackedBinomialExtensionField<F, PF, D>
{
}

impl<F: BinomiallyExtendable<D>, PF: PackedField<Scalar = F>, const D: usize> Algebra<PF>
    for PackedBinomialExtensionField<F, PF, D>
{
    #[inline]
    fn mixed_dot_product<const N: usize>(a: &[Self; N], f: &[PF; N]) -> Self
    where
        PF: Dup,
    {
        // Output container; each coordinate is filled independently below.
        let mut result = Self::default();

        // One base-field dot product per output coordinate.
        for k in 0..D {
            // Strided gather of the k-th coordinate from each extension input:
            //
            //     coord_k = [ a_0[k], a_1[k], ..., a_{N-1}[k] ]
            let coord_k: [PF; N] = core::array::from_fn(|i| a[i].value[k]);

            // Base-level dot product.
            //
            // - For Monty-31 packings this is the delayed-reduction primitive;
            // - For other packings it falls back to the eager default and the override is a no-op gain.
            result.value[k] = PF::dot_product::<N>(&coord_k, f);
        }

        result
    }
}

impl<F, PF, const D: usize> PrimeCharacteristicRing for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type PrimeSubfield = PF::PrimeSubfield;

    const ZERO: Self = Self::new([PF::ZERO; D]);

    const ONE: Self = Self::new(field_to_array(PF::ONE));

    const TWO: Self = Self::new(field_to_array(PF::TWO));

    const NEG_ONE: Self = Self::new(field_to_array(PF::NEG_ONE));

    #[inline]
    fn from_prime_subfield(val: Self::PrimeSubfield) -> Self {
        PF::from_prime_subfield(val).into()
    }

    #[inline]
    fn from_bool(b: bool) -> Self {
        PF::from_bool(b).into()
    }

    #[inline]
    fn halve(&self) -> Self {
        Self::new(self.value.map(|x| x.halve()))
    }

    #[inline(always)]
    fn square(&self) -> Self {
        let mut res = Self::default();
        let w = F::W;
        binomial_square(&self.value, &mut res.value, w);
        res
    }

    #[inline]
    fn mul_2exp_u64(&self, exp: u64) -> Self {
        Self::new(self.value.map(|x| x.mul_2exp_u64(exp)))
    }

    #[inline]
    fn div_2exp_u64(&self, exp: u64) -> Self {
        Self::new(self.value.map(|x| x.div_2exp_u64(exp)))
    }

    #[inline]
    fn zero_vec(len: usize) -> Vec<Self> {
        // SAFETY: this is a repr(transparent) wrapper around an array.
        unsafe { reconstitute_from_base(PF::zero_vec(len * D)) }
    }
}

impl<F, PF, const D: usize> BasedVectorSpace<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    const DIMENSION: usize = D;

    #[inline]
    fn as_basis_coefficients_slice(&self) -> &[PF] {
        &self.value
    }

    #[inline]
    fn from_basis_coefficients_fn<Fn: FnMut(usize) -> PF>(f: Fn) -> Self {
        Self::new(array::from_fn(f))
    }

    #[inline]
    fn from_basis_coefficients_iter<I: ExactSizeIterator<Item = PF>>(mut iter: I) -> Option<Self> {
        (iter.len() == D).then(|| Self::new(array::from_fn(|_| iter.next().unwrap()))) // The unwrap is safe as we just checked the length of iter.
    }

    #[inline]
    fn flatten_to_base(vec: Vec<Self>) -> Vec<PF> {
        unsafe {
            // Safety:
            // As `Self` is a `repr(transparent)`, it is stored identically in memory to `[PF; D]`
            flatten_to_base(vec)
        }
    }

    #[inline]
    fn reconstitute_from_base(vec: Vec<PF>) -> Vec<Self> {
        unsafe {
            // Safety:
            // As `Self` is a `repr(transparent)`, it is stored identically in memory to `[PF; D]`
            reconstitute_from_base(vec)
        }
    }
}

impl<F, const D: usize> PackedFieldExtension<F, BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, F::Packing, D>
where
    F: BinomiallyExtendable<D>,
{
    #[inline]
    fn from_ext_fn(f: impl Fn(usize) -> BinomialExtensionField<F, D>) -> Self {
        Self::new(F::Packing::pack_columns_fn(|lane| f(lane).value))
    }

    #[inline]
    fn add_assign_lane(&mut self, lane: usize, value: BinomialExtensionField<F, D>) {
        // Add each coefficient into its own base packing at the shared lane.
        for (coeff, v) in self.value.iter_mut().zip(value.value) {
            coeff.as_slice_mut()[lane] += v;
        }
    }

    #[inline]
    fn packed_ext_powers(base: BinomialExtensionField<F, D>) -> crate::Powers<Self> {
        let width = F::Packing::WIDTH;
        let powers = base.powers().collect_n(width + 1);
        // Transpose first WIDTH powers
        let current = Self::from_ext_slice(&powers[..width]);

        // Broadcast self^WIDTH
        let multiplier = powers[width].into();

        Powers {
            base: multiplier,
            current,
        }
    }
}

impl<F, PF, const D: usize> Neg for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn neg(self) -> Self {
        Self::new(self.value.map(PF::neg))
    }
}

impl<F, PF, const D: usize> Add for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn add(self, rhs: Self) -> Self {
        let value = vector_add(&self.value, &rhs.value);
        Self::new(value)
    }
}

impl<F, PF, const D: usize> Add<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn add(self, rhs: BinomialExtensionField<F, D>) -> Self {
        let value = vector_add(&self.value, &rhs.value);
        Self::new(value)
    }
}

impl<F, PF, const D: usize> Add<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn add(mut self, rhs: PF) -> Self {
        self.value[0] += rhs;
        self
    }
}

impl<F, PF, const D: usize> AddAssign for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn add_assign(&mut self, rhs: Self) {
        for i in 0..D {
            self.value[i] += rhs.value[i];
        }
    }
}

impl<F, PF, const D: usize> AddAssign<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn add_assign(&mut self, rhs: BinomialExtensionField<F, D>) {
        for i in 0..D {
            self.value[i] += rhs.value[i];
        }
    }
}

impl<F, PF, const D: usize> AddAssign<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn add_assign(&mut self, rhs: PF) {
        self.value[0] += rhs;
    }
}

impl<F, PF, const D: usize> Sum for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.reduce(|acc, x| acc + x).unwrap_or(Self::ZERO)
    }
}

impl<F, PF, const D: usize> Sub for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn sub(self, rhs: Self) -> Self {
        let value = vector_sub(&self.value, &rhs.value);
        Self::new(value)
    }
}

impl<F, PF, const D: usize> Sub<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn sub(self, rhs: BinomialExtensionField<F, D>) -> Self {
        let value = vector_sub(&self.value, &rhs.value);
        Self::new(value)
    }
}

impl<F, PF, const D: usize> Sub<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn sub(self, rhs: PF) -> Self {
        let mut res = self.value;
        res[0] -= rhs;
        Self::new(res)
    }
}

impl<F, PF, const D: usize> SubAssign for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl<F, PF, const D: usize> SubAssign<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn sub_assign(&mut self, rhs: BinomialExtensionField<F, D>) {
        *self = *self - rhs;
    }
}

impl<F, PF, const D: usize> SubAssign<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn sub_assign(&mut self, rhs: PF) {
        *self = *self - rhs;
    }
}

impl<F, PF, const D: usize> Mul for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn mul(self, rhs: Self) -> Self {
        let a = self.value;
        let b = rhs.value;
        let mut res = Self::default();
        let w = F::W;

        binomial_mul::<F, PF, PF, D>(&a, &b, &mut res.value, w);

        res
    }
}

impl<F, PF, const D: usize> Mul<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn mul(self, rhs: BinomialExtensionField<F, D>) -> Self {
        let a = self.value;
        let b = rhs.value;
        let mut res = Self::default();
        let w = F::W;

        binomial_mul(&a, &b, &mut res.value, w);

        res
    }
}

impl<F, PF, const D: usize> Mul<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[inline]
    fn mul(self, rhs: PF) -> Self {
        Self::new(self.value.map(|x| x * rhs))
    }
}

impl<F, PF, const D: usize> Product for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.reduce(|acc, x| acc * x).unwrap_or(Self::ONE)
    }
}

impl<F, PF, const D: usize> MulAssign for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl<F, PF, const D: usize> MulAssign<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn mul_assign(&mut self, rhs: BinomialExtensionField<F, D>) {
        *self = *self * rhs;
    }
}

impl<F, PF, const D: usize> MulAssign<PF> for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn mul_assign(&mut self, rhs: PF) {
        *self = *self * rhs;
    }
}

impl<F, PF, const D: usize> Div<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    type Output = Self;

    #[allow(clippy::suspicious_arithmetic_impl)]
    #[inline]
    fn div(self, rhs: BinomialExtensionField<F, D>) -> Self {
        self * Self::from(rhs.inverse())
    }
}

impl<F, PF, const D: usize> DivAssign<BinomialExtensionField<F, D>>
    for PackedBinomialExtensionField<F, PF, D>
where
    F: BinomiallyExtendable<D>,
    PF: PackedField<Scalar = F>,
{
    #[inline]
    fn div_assign(&mut self, rhs: BinomialExtensionField<F, D>) {
        *self = *self / rhs;
    }
}

impl<F: BinomiallyExtendable<D>, const D: usize> Div
    for PackedBinomialExtensionField<F, F::Packing, D>
{
    type Output = Self;

    #[allow(clippy::suspicious_arithmetic_impl)]
    #[inline]
    fn div(self, rhs: Self) -> Self {
        self * crate::invert_packed_extension::<F, BinomialExtensionField<F, D>>(rhs)
    }
}

impl<F: BinomiallyExtendable<D>, const D: usize> DivAssign
    for PackedBinomialExtensionField<F, F::Packing, D>
{
    #[inline]
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs;
    }
}