lib-q-ml-dsa 0.0.2

use crate::arithmetic::shift_left_then_reduce;
use crate::constants::BITS_IN_LOWER_PART_OF_T;
use crate::ntt::{
    invert_ntt_montgomery,
    ntt,
    ntt_multiply_montgomery,
    reduce,
};
use crate::polynomial::PolynomialRingElement;
use crate::simd::traits::Operations;

// Not inlining this makes key generation 3x slower for avx2. Only `inline` this
// function costs 30% performance too.
/// Compute InvertNTT(Â ◦ ŝ₁) + s₂
#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn compute_as1_plus_s2<SIMDUnit: Operations>(
    rows_in_a: usize,
    columns_in_a: usize,
    a_as_ntt: &mut [PolynomialRingElement<SIMDUnit>],
    s1_ntt: &[PolynomialRingElement<SIMDUnit>],
    s1_s2: &[PolynomialRingElement<SIMDUnit>],
    result: &mut [PolynomialRingElement<SIMDUnit>],
) {
    for i in 0..rows_in_a {
        for j in 0..columns_in_a {
            // Create a copy of the matrix element to avoid modifying the original matrix
            let mut product = a_as_ntt[i * columns_in_a + j];
            ntt_multiply_montgomery::<SIMDUnit>(&mut product, &s1_ntt[j]);
            PolynomialRingElement::add(&mut result[i], &product);
        }
    }

    for i in 0..result.len() {
        // We do a Barrett reduction here, since the absolute value of
        // `columns_in_a` additions might be as large as `columns_in_a
        // * FIELD_MODULUS`, and `invert_ntt_montgomery` expects
        // coefficients of size at most `FIELD_MODULUS`.
        reduce(&mut result[i]);
        invert_ntt_montgomery::<SIMDUnit>(&mut result[i]);
        PolynomialRingElement::add(&mut result[i], &s1_s2[columns_in_a + i]);
    }
}

/// Compute InvertNTT(Â ◦ ŷ)
#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn compute_matrix_x_mask<SIMDUnit: Operations>(
    rows_in_a: usize,
    columns_in_a: usize,
    matrix: &[PolynomialRingElement<SIMDUnit>],
    mask: &[PolynomialRingElement<SIMDUnit>],
    result: &mut [PolynomialRingElement<SIMDUnit>],
) {
    for i in 0..rows_in_a {
        for j in 0..columns_in_a {
            // We could make `matrix` mutable here and avoid copying.
            // But that would require sampling the matrix multiple times.
            // That's not worth it.
            let mut product = mask[j];
            ntt_multiply_montgomery(&mut product, &matrix[i * columns_in_a + j]);
            PolynomialRingElement::<SIMDUnit>::add(&mut result[i], &product);
        }
        // We do a Barrett reduction here, since the absolute value of
        // `columns_in_a` additions might be as large as `columns_in_a
        // * FIELD_MODULUS`, and `invert_ntt_montgomery` expects
        // coefficients of size at most `FIELD_MODULUS`.
        reduce(&mut result[i]);
        invert_ntt_montgomery(&mut result[i]);
    }
}

#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn vector_times_ring_element<SIMDUnit: Operations>(
    vector: &mut [PolynomialRingElement<SIMDUnit>],
    ring_element: &PolynomialRingElement<SIMDUnit>,
) {
    for element in vector {
        ntt_multiply_montgomery(element, ring_element);
        invert_ntt_montgomery(element);
    }
}

#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn add_vectors<SIMDUnit: Operations>(
    dimension: usize,
    lhs: &mut [PolynomialRingElement<SIMDUnit>],
    rhs: &[PolynomialRingElement<SIMDUnit>],
) {
    for i in 0..dimension {
        PolynomialRingElement::<SIMDUnit>::add(&mut lhs[i], &rhs[i]);
    }
}

#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn subtract_vectors<SIMDUnit: Operations>(
    dimension: usize,
    lhs: &mut [PolynomialRingElement<SIMDUnit>],
    rhs: &[PolynomialRingElement<SIMDUnit>],
) {
    for i in 0..dimension {
        PolynomialRingElement::<SIMDUnit>::subtract(&mut lhs[i], &rhs[i]);
    }
}

/// After two masked `vector_times_ring_element` passes (`acc` and `rhs`), fold `rhs` into `acc` and
/// Barrett-reduce so per-coefficient values match the canonical representatives from a single pass.
/// Without this, `InvNTT(c·s_a)+InvNTT(c·s_b)` can differ from `InvNTT(c·s)` as raw `i32`s (same mod
/// `q`), breaking infinity-norm rejection and NIST KAT byte-for-byte signatures under `hardened`.
#[cfg(feature = "hardened")]
#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn merge_masked_ntt_products<SIMDUnit: Operations>(
    dimension: usize,
    acc: &mut [PolynomialRingElement<SIMDUnit>],
    rhs: &[PolynomialRingElement<SIMDUnit>],
) {
    for i in 0..dimension {
        PolynomialRingElement::<SIMDUnit>::add(&mut acc[i], &rhs[i]);
        reduce(&mut acc[i]);
    }
}

/// Compute InvertNTT(Â ◦ ẑ - ĉ ◦ NTT(t₁2ᵈ))
#[cfg_attr(tarpaulin, inline(never))]
#[cfg_attr(not(tarpaulin), inline(always))]
pub(crate) fn compute_w_approx<SIMDUnit: Operations>(
    rows_in_a: usize,
    columns_in_a: usize,
    matrix: &[PolynomialRingElement<SIMDUnit>],
    signer_response: &[PolynomialRingElement<SIMDUnit>],
    verifier_challenge_as_ntt: &PolynomialRingElement<SIMDUnit>,
    t1: &mut [PolynomialRingElement<SIMDUnit>],
) {
    for i in 0..rows_in_a {
        let mut inner_result = PolynomialRingElement::<SIMDUnit>::zero();
        for j in 0..columns_in_a {
            let mut product = matrix[i * columns_in_a + j];
            ntt_multiply_montgomery(&mut product, &signer_response[j]);
            PolynomialRingElement::<SIMDUnit>::add(&mut inner_result, &product);
        }

        shift_left_then_reduce::<SIMDUnit, { BITS_IN_LOWER_PART_OF_T as i32 }>(&mut t1[i]);
        ntt(&mut t1[i]);
        ntt_multiply_montgomery(&mut t1[i], verifier_challenge_as_ntt);
        PolynomialRingElement::<SIMDUnit>::subtract(&mut inner_result, &t1[i]);
        t1[i] = inner_result;
        // We do a Barrett reduction here, since the absolute value of
        // `columns_in_a` additions might be as large as `columns_in_a
        // * FIELD_MODULUS`, and `invert_ntt_montgomery` expects
        // coefficients of size at most `FIELD_MODULUS`.
        reduce(&mut t1[i]);
        invert_ntt_montgomery(&mut t1[i]);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::polynomial::PolynomialRingElement;
    use crate::simd::portable::PortableSIMDUnit;

    type P = PolynomialRingElement<PortableSIMDUnit>;

    #[test]
    fn compute_as1_plus_s2_2x2_zeros() {
        let mut a = [P::zero(); 4];
        let s1_ntt = [P::zero(); 2];
        let s1_s2 = [P::zero(); 4];
        let mut result = [P::zero(); 2];
        compute_as1_plus_s2(2, 2, &mut a, &s1_ntt, &s1_s2, &mut result);
    }

    #[test]
    fn matrix_x_mask_and_vector_ops_smoke() {
        let matrix = [P::zero(); 4];
        let mask = [P::zero(); 2];
        let mut out = [P::zero(); 2];
        compute_matrix_x_mask(2, 2, &matrix, &mask, &mut out);

        let mut v = [P::zero(); 2];
        let r = P::zero();
        vector_times_ring_element(&mut v, &r);

        let mut lhs = [P::zero(); 2];
        let rhs = [P::zero(); 2];
        add_vectors(2, &mut lhs, &rhs);
        subtract_vectors(2, &mut lhs, &rhs);
    }

    #[test]
    fn compute_w_approx_2x2_zeros() {
        let matrix = [P::zero(); 4];
        let signer = [P::zero(); 2];
        let c = P::zero();
        let mut t1 = [P::zero(); 2];
        compute_w_approx(2, 2, &matrix, &signer, &c, &mut t1);
    }
}