arcium-primitives 0.4.6

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

use ff::Field;
use typenum::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},
    izip_eq,
    random::{CryptoRngCore, Random},
    types::{heap_array::SubfieldElements, Positive},
};

/// 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.clone();
        walsh_hadamard::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> {
    pub fn new_from_sparse_coef(
        coefs: impl IntoIterator<Item = (usize, SubfieldElement<F>)>,
    ) -> Self {
        let mut data = SubfieldElements::<F, M>::default();
        for (mut i, c_i) in coefs {
            debug_assert!(i < M::to_usize());
            i &= M::to_usize() - 1; // Reduce `i` modulo `M`
            data[i] += c_i;
        }
        MultivariateRing::<_, _, Coefficient>::new(data).to_eval_repr()
    }

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

    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(self, rhs: &MultivariateRing<F, M, Evaluation>) -> Self::Output {
        Self::new(izip_eq!(self.data, &rhs.data).map(|(a, b)| a * b).collect())
    }
}

impl<F: FieldExtension, M: Positive> AddAssign for MultivariateRing<F, M, Evaluation> {
    fn add_assign(&mut self, rhs: Self) {
        izip_eq!(&mut self.data, &rhs.data).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)
    }
}

mod walsh_hadamard {
    use typenum::PowerOfTwo;

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

    // In-place Walsh-Hadamard transformation
    pub(super) fn walsh_hadamard_inplace<F: FieldExtension, M: Positive + PowerOfTwo>(
        data: &mut SubfieldElements<F, M>,
    ) {
        let m = M::to_usize();
        let mut step = 1;
        while step < m {
            let mut i = 0;
            let step2 = step << 1;
            while i < m {
                for j in i..i + step {
                    let t = data[j];
                    data[j] = t - data[j + step];
                    data[j + step] = t + data[j + step];
                }
                i += step2;
            }
            step = step2;
        }
    }
}

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

    use super::{
        walsh_hadamard::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>();
    }

    #[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
        );
    }
}