arcium-primitives 0.6.0

use std::{
    iter::Sum,
    marker::PhantomData,
    ops::{Add, AddAssign, Mul},
};

use ff::Field;
use rayon::prelude::*;
use typenum::{PartialDiv, PowerOfTwo};

// This module defines a multivariate polynomial ring $$ F_q[X_1, ..., X_n] / (X_1^2 - 1, ...,
// X_n^2 - 1) $$ together with basic operations of this ring.
use crate::{
    algebra::field::{FieldExtension, SubfieldElement},
    random::{CryptoRngCore, Random},
    types::{heap_array::SubfieldElements, Positive},
};

/// Minimum number of ring elements one rayon task processes; below this the
/// element-wise loops run effectively sequentially.
const PAR_MIN_LEN: usize = 1 << 13;

/// The form of multivariate polynomial
#[derive(Default, Copy, Clone, Debug, PartialEq)]
pub struct Coefficient;
#[derive(Default, Copy, Clone, Debug, PartialEq)]
pub struct Evaluation;

/// An element `P(X_0, ..., X_{N-1})` of the multivariate polynomial ring with `N` variables.
///
/// The polynomial can be either in coefficient or in evaluation form, given by `Form`
/// parameter. The number of coefficients/evaluations is `M = 2 ^ N`.
///
/// In _coefficient_ representation `data` are the coefficients of the polynomial:
///     `\sum_{i=0..2^N-1} c_i * \Prod_{j=0..N-1} X_j ^ bin(i)_j`
/// `bin(i)` is the  binary-decomposition of `i` (LSB first).
/// For example for `N=2` the 2-variate polynomial is `c_0 + c_1 . X_0 + c_2 . X_1 + c_3 . X_0 .
/// X_1` and `data = [c_0, c_1, c_2, c_3]`.
///
/// In _evaluation_ representation `data` are the evaluations:
///     `e_i = P(bin_signed(i))`
/// `bin_signed(i)` is the signed binary-decomposition of `i` (LSB first).
/// For example for `N=2` the evaluations are:
///     - `e_0 = P(-1, -1)`
///     - `e_1 = P(-1,  1)`
///     - `e_2 = P( 1, -1)`
///     - `e_3 = P( 1,  1)`
/// and `data = [e_0, e_1, e_2, e_3]`
#[derive(Clone, Default, Debug, PartialEq)]
pub struct MultivariateRing<F: FieldExtension, M: Positive, Form> {
    data: SubfieldElements<F, M>,
    _form: PhantomData<Form>,
}

pub type MultivariateRingCoefForm<F, M> = MultivariateRing<F, M, Coefficient>;
pub type MultivariateRingEvalForm<F, M> = MultivariateRing<F, M, Evaluation>;

impl<F: FieldExtension, M: Positive + PowerOfTwo, Form> MultivariateRing<F, M, Form> {
    pub fn new(data: SubfieldElements<F, M>) -> Self {
        Self {
            data,
            _form: PhantomData,
        }
    }

    pub fn random(rng: impl CryptoRngCore) -> Self {
        Self::new(SubfieldElements::<F, M>::random(rng))
    }

    pub fn data(&self) -> &SubfieldElements<F, M> {
        &self.data
    }

    pub fn into_data(self) -> SubfieldElements<F, M> {
        self.data
    }
}

impl<F: FieldExtension, M: Positive + PowerOfTwo> MultivariateRing<F, M, Coefficient> {
    /// Transform polynomial from coefficient representation to evaluation representation using the
    /// Walsh-Hadamard transform.
    pub fn to_eval_repr(self) -> MultivariateRing<F, M, Evaluation> {
        let mut data: SubfieldElements<F, M> = self.data;
        walsh_hadamard_inplace(&mut data);
        MultivariateRing::<_, _, Evaluation>::new(data)
    }

    pub fn nb_coefs(&self) -> usize {
        self.data.len()
    }
}

impl<F: FieldExtension, M: Positive + PowerOfTwo> MultivariateRing<F, M, Evaluation> {
    /// Realize a `T`-sparse coefficient vector directly in evaluation form.
    ///
    /// The Walsh transform of `Σ_k c_k·X^{i_k}` is evaluated point-wise — the
    /// monomial `X^{i}` at the signed point `bin_signed(x)` is
    /// `(−1)^{popcount(i & !x)}` — costing `O(M·T)`.
    ///
    /// # Crossover with the dense path
    ///
    /// Densifying the coefficients and running the full Walsh–Hadamard transform
    /// ([`Self::from_coeffs`], i.e. `to_eval_repr` on the coefficient form) costs `O(M·log M) =
    /// O(M·N)`, where `N = log2(M)` is the number of variables. This sparse path is therefore
    /// only cheaper while `T ≲ N`; for larger `T` it is strictly worse (per element it does `T`
    /// work versus the dense path's `N`), so callers with `T > N` should densify and use the
    /// dense path instead.
    pub fn new_from_sparse_coef(
        coefs: impl IntoIterator<Item = (usize, SubfieldElement<F>)>,
    ) -> Self {
        let m = M::to_usize();
        let coefs: Vec<(usize, SubfieldElement<F>)> = coefs
            .into_iter()
            .inspect(|(i, _)| debug_assert!(*i < m))
            .map(|(i, c_i)| (i & (m - 1), c_i)) // Reduce `i` modulo `M`
            .collect();

        let mut data = SubfieldElements::<F, M>::default();
        data.par_iter_mut()
            .enumerate()
            .with_min_len(PAR_MIN_LEN)
            .for_each(|(x, e)| {
                let not_x = !x & (m - 1);
                for (i, c_i) in &coefs {
                    if (i & not_x).count_ones() & 1 == 1 {
                        *e -= *c_i;
                    } else {
                        *e += *c_i;
                    }
                }
            });
        Self::new(data)
    }

    pub fn from_coeffs(coefs: SubfieldElements<F, M>) -> Self {
        MultivariateRing::new(coefs).to_eval_repr()
    }

    /// Complete a transform whose `BlockSize`-aligned blocks are already in the Walsh domain: only
    /// the cross-block butterfly levels are applied (the transform factors per bit, so levels
    /// commute). `BlockSize == M` is a pure relabel; `BlockSize == 1` is the full transform.
    pub fn from_blockwise_walsh<BlockSize>(mut data: SubfieldElements<F, M>) -> Self
    where
        BlockSize: Positive + PowerOfTwo,
        M: PartialDiv<BlockSize>,
    {
        walsh_hadamard::walsh_hadamard_from_step(&mut data, BlockSize::USIZE);
        Self::new(data)
    }

    pub fn square(&self) -> Self {
        Self {
            data: self.data.iter().map(SubfieldElement::<F>::square).collect(),
            _form: PhantomData,
        }
    }

    pub fn nb_evals(&self) -> usize {
        self.data.len()
    }
}

#[macros::op_variants(owned, borrowed, flipped)]
impl<F: FieldExtension, M: Positive + PowerOfTwo> Mul<&MultivariateRing<F, M, Evaluation>>
    for MultivariateRing<F, M, Evaluation>
{
    type Output = MultivariateRing<F, M, Evaluation>;

    fn mul(mut self, rhs: &MultivariateRing<F, M, Evaluation>) -> Self::Output {
        self.data
            .par_iter_mut()
            .zip(rhs.data.par_iter())
            .with_min_len(PAR_MIN_LEN)
            .for_each(|(a, b)| *a *= b);
        self
    }
}

impl<F: FieldExtension, M: Positive> AddAssign for MultivariateRing<F, M, Evaluation> {
    fn add_assign(&mut self, rhs: Self) {
        self.data
            .par_iter_mut()
            .zip(rhs.data.par_iter())
            .with_min_len(PAR_MIN_LEN)
            .for_each(|(a, b)| *a += b);
    }
}

impl<F: FieldExtension, M: Positive + PowerOfTwo> Add for MultivariateRing<F, M, Evaluation> {
    type Output = Self;

    fn add(mut self, rhs: Self) -> Self::Output {
        self += rhs;
        self
    }
}

impl<F: FieldExtension, M: Positive + PowerOfTwo> Sum for MultivariateRing<F, M, Evaluation> {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::default(), |acc, x| acc + x)
    }
}

/// In-place Walsh-Hadamard transform: the coefficient → evaluation map of
/// [`MultivariateRing`]. Linear, an involution up to scaling by `M`.
pub fn walsh_hadamard_inplace<F: FieldExtension, M: Positive + PowerOfTwo>(
    data: &mut SubfieldElements<F, M>,
) {
    walsh_hadamard::walsh_hadamard_from_step(data, 1);
}

mod walsh_hadamard {
    use rayon::prelude::*;
    use typenum::PowerOfTwo;

    use super::PAR_MIN_LEN;
    use crate::{
        algebra::field::FieldExtension,
        types::{heap_array::SubfieldElements, Positive},
    };

    /// In-place Walsh-Hadamard butterfly levels with step ≥ `first_step`.
    ///
    /// The transform factors into one butterfly level per index bit and the
    /// levels commute, so starting at `first_step = B` completes a transform
    /// whose `B`-aligned blocks are already transformed; `first_step = 1` is
    /// the full transform. `first_step` must be a power of two.
    pub(super) fn walsh_hadamard_from_step<F: FieldExtension, M: Positive + PowerOfTwo>(
        data: &mut SubfieldElements<F, M>,
        first_step: usize,
    ) {
        let m = M::to_usize();
        debug_assert!(first_step.is_power_of_two() && first_step <= m);
        let mut step = first_step;
        while step < m {
            let step2 = step << 1;
            if step < PAR_MIN_LEN {
                // Many small blocks: parallelize across blocks.
                data.par_chunks_mut(step2)
                    .with_min_len((PAR_MIN_LEN / step2).max(1))
                    .for_each(|chunk| {
                        let (lo, hi) = chunk.split_at_mut(step);
                        for (a, b) in lo.iter_mut().zip(hi) {
                            let t = *a;
                            *a = t - *b;
                            *b = t + *b;
                        }
                    });
            } else {
                // Few large blocks: parallelize the butterflies inside each.
                for chunk in data.chunks_mut(step2) {
                    let (lo, hi) = chunk.split_at_mut(step);
                    lo.par_iter_mut()
                        .zip(hi.par_iter_mut())
                        .with_min_len(PAR_MIN_LEN)
                        .for_each(|(a, b)| {
                            let t = *a;
                            *a = t - *b;
                            *b = t + *b;
                        });
                }
            }
            step = step2;
        }
    }
}

#[cfg(test)]
mod tests {
    use itertools::izip;
    use itybity::IntoBits;
    use typenum::{PartialDiv, PowerOfTwo, Shleft, Unsigned};

    use super::{walsh_hadamard_inplace, Coefficient, Evaluation, MultivariateRing};
    use crate::{
        algebra::{
            elliptic_curve::{Curve25519Ristretto as C, ScalarField},
            field::{mersenne::Mersenne107, FieldExtension, SubfieldElement},
        },
        izip_eq,
        random::{test_rng, Random},
        types::{heap_array::SubfieldElements, Positive},
    };

    type Fq = ScalarField<C>;

    impl<F: FieldExtension, M: Positive + PowerOfTwo> MultivariateRing<F, M, Coefficient> {
        fn eval(&self, x_vals: Vec<SubfieldElement<F>>) -> SubfieldElement<F> {
            use num_traits::identities::One;
            debug_assert_eq!(1usize << x_vals.len(), self.data.len());
            self.data
                .iter()
                .enumerate()
                .map(|(i, c_i)| {
                    let monomial_i = izip!(x_vals.iter(), i.into_iter_lsb0())
                        .map(|(x_i, b_i)| {
                            if b_i {
                                *x_i
                            } else {
                                SubfieldElement::<F>::one()
                            }
                        })
                        .product::<SubfieldElement<F>>();
                    *c_i * monomial_i
                })
                .sum()
        }

        /// Multiply 2 polynomials in coefficient representation.
        ///
        /// __Remark:__ Use only in tests as it has O(M^2) complexity
        fn mul_slow(&self, other: &Self) -> Self {
            let mut out = Self::default();

            let n = self.nb_coefs();
            for i in 0..n {
                for j in 0..n {
                    let k = i ^ j;
                    out.data[k] += self.data[i] * other.data[j];
                }
            }

            out
        }
    }

    fn signed_binary_decomposition<F: FieldExtension>(
        x: usize,
        n: usize,
    ) -> Vec<SubfieldElement<F>> {
        use num_traits::identities::One;
        x.into_iter_lsb0()
            .take(n)
            .map(|b_i| {
                if b_i {
                    SubfieldElement::<F>::one()
                } else {
                    -SubfieldElement::<F>::one()
                }
            })
            .collect()
    }

    fn test_multivariate_ring_walsh_hadamard_impl<F: FieldExtension, M: Positive + PowerOfTwo>(
        walsh_hadamard: fn(&mut SubfieldElements<F, M>),
    ) {
        let rng = test_rng();
        let poly_coef = MultivariateRing::<F, M, Coefficient>::random(rng);
        let mut data = poly_coef.data.clone();
        walsh_hadamard(&mut data);
        let poly_eval = MultivariateRing::<_, _, Evaluation>::new(data);

        // Check each evaluation values by manually evaluating the input polynomial at each
        // point
        let log2m = M::to_usize().ilog2() as usize;
        for i in 0..poly_eval.nb_evals() {
            // Signed binary-decomposition of `i`
            let x_vals: Vec<_> = signed_binary_decomposition::<F>(i, log2m);
            let e_i = poly_coef.eval(x_vals);
            assert_eq!(e_i, poly_eval.data[i])
        }
    }

    #[test]
    fn test_multivariate_ring_walsh_hadamard() {
        type M = Shleft<typenum::U1, typenum::U8>;

        test_multivariate_ring_walsh_hadamard_impl::<Mersenne107, M>(walsh_hadamard_inplace);
        test_multivariate_ring_walsh_hadamard_impl::<Fq, M>(walsh_hadamard_inplace);
    }

    fn test_multivariate_ring_new_from_sparse_coef_impl<
        F: FieldExtension,
        M: Positive + PowerOfTwo,
        T: Positive,
    >() {
        use num_traits::identities::One;

        let mut rng = test_rng();
        let coefs = SubfieldElements::<F, T>::random(&mut rng);
        let coefs_sparse: Vec<_> = izip!(
            (0..T::to_usize()).map(|_| usize::random(&mut rng) % M::to_usize()),
            coefs
        )
        .collect();

        let poly_eval_exp =
            MultivariateRing::<_, M, Evaluation>::new_from_sparse_coef(coefs_sparse.clone());

        let log2m = M::to_usize().ilog2() as usize;
        let coef_sparse_bin: Vec<(Vec<_>, SubfieldElement<F>)> = coefs_sparse
            .into_iter()
            .map(|(i, c_i)| (i.into_iter_lsb0().take(log2m).collect(), c_i))
            .collect();

        for i in 0..poly_eval_exp.nb_evals() {
            let x_vals: Vec<_> = signed_binary_decomposition::<F>(i, log2m);
            let e_i = coef_sparse_bin
                .iter()
                .map(|(i_bin, c_i)| {
                    izip_eq!(&x_vals, i_bin)
                        .map(|(x_j, i_bin_j)| {
                            if *i_bin_j {
                                *x_j
                            } else {
                                SubfieldElement::<F>::one()
                            }
                        })
                        .product::<SubfieldElement<F>>()
                        * *c_i
                })
                .sum::<SubfieldElement<F>>();

            assert_eq!(e_i, poly_eval_exp.data[i]);
        }
    }

    #[test]
    fn test_multivariate_ring_new_from_sparse_coef() {
        type M = Shleft<typenum::U1, typenum::U8>;
        test_multivariate_ring_new_from_sparse_coef_impl::<Mersenne107, M, typenum::U1>();
        test_multivariate_ring_new_from_sparse_coef_impl::<Mersenne107, M, typenum::U13>();
        test_multivariate_ring_new_from_sparse_coef_impl::<Fq, M, typenum::U1>();
        test_multivariate_ring_new_from_sparse_coef_impl::<Fq, M, typenum::U13>();
    }

    fn test_multivariate_ring_mul_impl<F: FieldExtension, M: Positive + PowerOfTwo>() {
        let mut rng = test_rng();
        let a = MultivariateRing::<F, M, Coefficient>::random(&mut rng);
        let b = MultivariateRing::<F, M, Coefficient>::random(&mut rng);

        // Transform to evaluation representation
        let a_eval = a.clone().to_eval_repr();
        let b_eval = b.clone().to_eval_repr();

        // Multiplication in the evaluation domain
        let c_eval_exp = a_eval * b_eval;

        // Schoolbook multiplication in the coefficient representation
        let c = a.mul_slow(&b);
        let c_eval_act = c.to_eval_repr();

        assert_eq!(c_eval_act, c_eval_exp);
    }

    #[test]
    fn test_multivariate_ring_mul() {
        type M = Shleft<typenum::U1, typenum::U8>;

        test_multivariate_ring_mul_impl::<Mersenne107, M>();
        test_multivariate_ring_mul_impl::<Fq, M>();
    }

    /// Transforming `BlockSize`-aligned blocks first and completing with
    /// `from_blockwise_walsh` reproduces the full transform, for every block
    /// size (1 = nothing pre-transformed, M = pure relabel).
    #[test]
    fn test_multivariate_ring_from_blockwise_walsh() {
        type F = Mersenne107;
        type M = Shleft<typenum::U1, typenum::U8>;

        // Pre-transform `BlockSize`-aligned blocks of `coeffs`, complete with
        // `from_blockwise_walsh::<BlockSize>`, and check it matches the full transform.
        fn check<BlockSize>(
            coeffs: &SubfieldElements<F, M>,
            expected: &MultivariateRing<F, M, Evaluation>,
        ) where
            BlockSize: Positive + PowerOfTwo,
            M: PartialDiv<BlockSize>,
        {
            let block_size = BlockSize::USIZE;
            let mut data = coeffs.clone();
            for block in data.chunks_mut(block_size) {
                // Blockwise transform via the same butterflies, on a copy.
                let mut step = 1;
                while step < block_size {
                    let step2 = step << 1;
                    for pair in block.chunks_mut(step2) {
                        let (lo, hi) = pair.split_at_mut(step);
                        for (a, b) in lo.iter_mut().zip(hi) {
                            let t = *a;
                            *a = t - *b;
                            *b = t + *b;
                        }
                    }
                    step = step2;
                }
            }
            let act = MultivariateRing::<F, M, Evaluation>::from_blockwise_walsh::<BlockSize>(data);
            assert_eq!(&act, expected, "BlockSize={block_size}");
        }

        let mut rng = test_rng();
        let coeffs = SubfieldElements::<F, M>::random(&mut rng);
        let expected = MultivariateRing::<F, M, Evaluation>::from_coeffs(coeffs.clone());

        // M = 2^8: every power-of-two block size from 1 to M.
        check::<typenum::U1>(&coeffs, &expected);
        check::<typenum::U2>(&coeffs, &expected);
        check::<typenum::U4>(&coeffs, &expected);
        check::<typenum::U8>(&coeffs, &expected);
        check::<typenum::U16>(&coeffs, &expected);
        check::<typenum::U32>(&coeffs, &expected);
        check::<typenum::U64>(&coeffs, &expected);
        check::<typenum::U128>(&coeffs, &expected);
        check::<typenum::U256>(&coeffs, &expected);
    }

    #[test]
    #[ignore]
    fn test_multivariate_ring_to_eval_repr_large() {
        type F = Mersenne107;
        type N = typenum::U25;
        type M = Shleft<typenum::U1, N>;
        const NB_ITER: usize = 1;

        let data_orig = SubfieldElements::<F, M>::random(test_rng());

        let mut data = data_orig.clone();
        let start = std::time::Instant::now();
        for _ in 0..NB_ITER {
            walsh_hadamard_inplace(&mut data);
            std::hint::black_box(&mut data);
        }
        let duration = start.elapsed();
        println!(
            "Walsh-Hadamard transformation of size 2^{} execution time: {:?} sec",
            N::to_usize(),
            duration.as_secs_f32() / NB_ITER as f32
        );
    }
}