krypteia-quantica 0.1.0

//! Small-polynomial representation for ML-DSA secret vectors.
//!
//! The secret polynomials `s1`, `s2` (and the challenge `c`) have
//! coefficients bounded by `eta` (2 or 4), which fits in `i16`.
//! Instead of storing them as `[i32; 256]` (1 024 bytes), we store
//! them as `[i16; 256]` (512 bytes) — a 50% savings.
//!
//! For NTT-domain multiplication, we reuse the ML-KEM NTT which
//! operates mod `q' = 3329` with i16 arithmetic. This works because
//! the product `c · s` has coefficients bounded by `TAU × eta ≤ 240`,
//! well below `q'/2 = 1664`, so no modular wraparound occurs.
//!
//! ## Pipeline
//!
//! ```text
//! s ∈ [-eta, eta]^256  →  i16  →  small_ntt (mod 3329)  →  i16 NTT domain
//! c ∈ {-1,0,+1}^256   →  i16  →  small_ntt (mod 3329)  →  i16 NTT domain
//!                             basemul (mod 3329)  →  i16 NTT domain
//!                             small_invntt        →  i16 time domain
//!                             widen to i32        →  result in [-q'/2, q'/2]
//! ```
//!
//! Enabled by the `small-secret` Cargo feature.

use super::params::N;

/// A polynomial with i16 coefficients (512 bytes instead of 1 024).
///
/// Used for secret vectors `s1`, `s2` and the challenge `c` when the
/// `small-secret` feature is enabled.
pub struct SmallPoly {
    pub coeffs: [i16; N],
}

impl SmallPoly {
    /// All-zero polynomial.
    pub const fn zero() -> Self {
        Self { coeffs: [0i16; N] }
    }

    /// Convert from i32 polynomial (narrowing — caller must ensure
    /// coefficients fit in i16).
    pub fn from_i32(src: &[i32; N]) -> Self {
        let mut s = Self::zero();
        for i in 0..N {
            s.coeffs[i] = src[i] as i16;
        }
        s
    }
}

// =====================================================================
// NTT operations — delegate to the ML-KEM NTT (q=3329, i16)
// =====================================================================

/// Forward NTT in-place (mod 3329, BitRev_7 partial).
pub fn small_ntt(p: &mut SmallPoly) {
    crate::ml_kem::ntt::ntt(&mut p.coeffs);
}

/// Inverse NTT in-place (mod 3329).
pub fn small_invntt(p: &mut SmallPoly) {
    crate::ml_kem::ntt::ntt_inv(&mut p.coeffs);
}

/// Basemul two small NTT-domain polynomials (mod 3329, len-2 base case)
/// and store in `out`.
pub fn small_basemul(out: &mut SmallPoly, a: &SmallPoly, b: &SmallPoly) {
    crate::ml_kem::ntt::multiply_ntts(&a.coeffs, &b.coeffs, &mut out.coeffs);
}

/// Basemul + accumulate: `out += a * b` in NTT domain (mod 3329).
pub fn small_basemul_acc(out: &mut SmallPoly, a: &SmallPoly, b: &SmallPoly) {
    let mut tmp = [0i16; N];
    crate::ml_kem::ntt::multiply_ntts(&a.coeffs, &b.coeffs, &mut tmp);
    for i in 0..N {
        out.coeffs[i] = crate::ml_kem::ntt::barrett_reduce(out.coeffs[i].wrapping_add(tmp[i]));
    }
}

/// Multiply small NTT-domain `a` by small NTT-domain `b`, inverse NTT,
/// and widen to i32. Result coefficients are in `[-(q'-1)/2, (q'-1)/2]`.
pub fn small_basemul_invntt_widen(a: &SmallPoly, b: &SmallPoly) -> [i32; N] {
    let mut tmp = SmallPoly::zero();
    small_basemul(&mut tmp, a, b);
    small_invntt(&mut tmp);
    let mut out = [0i32; N];
    for i in 0..N {
        out[i] = tmp.coeffs[i] as i32;
    }
    out
}

#[cfg(test)]
mod tests {
    use super::super::ntt as dsa_ntt;
    use super::super::params::Q;
    use super::*;

    /// Verify that small NTT multiplication matches the full i32 NTT
    /// for the product of a challenge-like poly × a secret-like poly.
    #[test]
    fn small_mul_matches_full_ntt() {
        // Build a challenge: TAU=39 nonzero coefficients in {-1, +1}
        let mut c_i32 = [0i32; N];
        for i in (N - 39)..N {
            c_i32[i] = if i % 3 == 0 { -1 } else { 1 };
        }
        // Build a secret: coefficients in [-2, 2]
        let mut s_i32 = [0i32; N];
        for i in 0..N {
            s_i32[i] = ((i as i32 * 7 + 3) % 5) - 2; // values in [-2, 2]
        }

        // --- Reference: full i32 ML-DSA NTT path ---
        let mut c_ntt = c_i32;
        for c in c_ntt.iter_mut() {
            *c = dsa_ntt::mod_q(*c);
        }
        dsa_ntt::ntt(&mut c_ntt);
        let mut s_ntt = s_i32;
        for c in s_ntt.iter_mut() {
            *c = dsa_ntt::mod_q(*c);
        }
        dsa_ntt::ntt(&mut s_ntt);
        let mut prod_full = dsa_ntt::pointwise_mul(&c_ntt, &s_ntt);
        dsa_ntt::ntt_inv(&mut prod_full);
        for c in prod_full.iter_mut() {
            *c = dsa_ntt::mod_q(*c);
        }

        // --- Small i16 path ---
        let c_small = SmallPoly::from_i32(&c_i32);
        let s_small = SmallPoly::from_i32(&s_i32);
        let mut c_small_ntt = c_small;
        small_ntt(&mut c_small_ntt);
        let mut s_small_ntt = s_small;
        small_ntt(&mut s_small_ntt);
        let prod_small = small_basemul_invntt_widen(&c_small_ntt, &s_small_ntt);

        // Compare: normalize both to [0, q-1]
        for i in 0..N {
            let expected = dsa_ntt::mod_q(prod_full[i]);
            let got_raw = prod_small[i];
            // The small path gives result mod 3329, which for small
            // values (< 3329/2) is the true integer value. Normalize.
            let got = dsa_ntt::mod_q(got_raw);
            assert_eq!(
                got, expected,
                "mismatch at i={}: small={} (raw={}), full={}",
                i, got, got_raw, expected
            );
        }
    }
}