krypteia-quantica 0.1.0

//! Compressed polynomial storage — 3 bytes per coefficient (24 bits).
//!
//! Since `q = 8_380_417 < 2^23`, fully reduced polynomial coefficients
//! fit in 23 bits. We store them in **3 bytes** (24 bits) per coefficient
//! instead of the standard 4-byte `i32`, saving 25% of RAM on each
//! polynomial vector.
//!
//! This module is compiled only when the `compressed-poly` Cargo feature
//! is enabled. The compressed format is used for intermediate polynomial
//! vectors in the ML-DSA Sign rejection loop (`w`, `cs2`, `w_minus_cs2`,
//! `r0`, `ct0`, etc.) that don't need full i32 precision between
//! operations.
//!
//! ## Layout
//!
//! A single compressed polynomial is `[u8; 768]` (256 coefficients × 3 bytes).
//! Coefficients must be in `[0, q-1]` (canonical form) before packing.

use super::ntt::mod_q;
use super::params::{N, Q};
use alloc::vec::Vec;

/// Bytes per compressed polynomial: 256 coefficients × 3 bytes = 768.
pub const COMPRESSED_POLY_BYTES: usize = N * 3;

/// Pack a polynomial into 3-byte-per-coefficient compressed form.
///
/// The input coefficients are reduced to `[0, q-1]` before packing.
#[inline]
pub fn poly_pack_24(out: &mut [u8], poly: &[i32; N]) {
    debug_assert!(out.len() >= COMPRESSED_POLY_BYTES);
    for i in 0..N {
        let c = mod_q(poly[i]) as u32;
        out[i * 3] = c as u8;
        out[i * 3 + 1] = (c >> 8) as u8;
        out[i * 3 + 2] = (c >> 16) as u8;
    }
}

/// Unpack a compressed polynomial back to `[i32; N]`.
///
/// Output coefficients are in `[0, q-1]`.
#[inline]
pub fn poly_unpack_24(poly: &mut [i32; N], data: &[u8]) {
    debug_assert!(data.len() >= COMPRESSED_POLY_BYTES);
    for i in 0..N {
        let c = data[i * 3] as u32 | ((data[i * 3 + 1] as u32) << 8) | ((data[i * 3 + 2] as u32) << 16);
        poly[i] = c as i32;
    }
}

/// A vector of K compressed polynomials (each 768 bytes).
///
/// Provides pack/unpack of individual polynomials and a
/// compressed-domain subtraction (`sub_unpack`) that reads two
/// compressed polys, subtracts, and writes the result to an `[i32; N]`.
pub struct CompressedVecK {
    /// Raw byte storage: `k` polynomials × 768 bytes each.
    data: Vec<u8>,
    /// Number of polynomials actually used.
    len: usize,
}

impl CompressedVecK {
    /// Allocate a zero-initialized compressed vector for `len` polynomials.
    pub fn new(len: usize) -> Self {
        Self {
            data: vec![0u8; len * COMPRESSED_POLY_BYTES],
            len,
        }
    }

    /// Pack a full polynomial into slot `idx`.
    pub fn pack(&mut self, idx: usize, poly: &[i32; N]) {
        let off = idx * COMPRESSED_POLY_BYTES;
        poly_pack_24(&mut self.data[off..off + COMPRESSED_POLY_BYTES], poly);
    }

    /// Unpack slot `idx` into a full polynomial.
    pub fn unpack(&self, idx: usize, poly: &mut [i32; N]) {
        let off = idx * COMPRESSED_POLY_BYTES;
        poly_unpack_24(poly, &self.data[off..off + COMPRESSED_POLY_BYTES]);
    }

    /// Read the raw bytes of slot `idx` (for hashing / encoding).
    pub fn slot(&self, idx: usize) -> &[u8] {
        let off = idx * COMPRESSED_POLY_BYTES;
        &self.data[off..off + COMPRESSED_POLY_BYTES]
    }

    /// Number of polynomial slots.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Subtract: `out[i] = self[idx][i] - b[i]` for each coefficient.
    ///
    /// Reads from compressed storage, subtracts `b` (a full polynomial),
    /// and writes to `out` (a full polynomial). Avoids allocating a
    /// full decompressed copy of `self[idx]`.
    pub fn sub_into(&self, idx: usize, b: &[i32; N], out: &mut [i32; N]) {
        let off = idx * COMPRESSED_POLY_BYTES;
        let data = &self.data[off..];
        for i in 0..N {
            let c = data[i * 3] as u32 | ((data[i * 3 + 1] as u32) << 8) | ((data[i * 3 + 2] as u32) << 16);
            out[i] = c as i32 - b[i];
        }
    }

    /// Add: `out[i] = self[idx][i] + b[i]` for each coefficient.
    pub fn add_into(&self, idx: usize, b: &[i32; N], out: &mut [i32; N]) {
        let off = idx * COMPRESSED_POLY_BYTES;
        let data = &self.data[off..];
        for i in 0..N {
            let c = data[i * 3] as u32 | ((data[i * 3 + 1] as u32) << 8) | ((data[i * 3 + 2] as u32) << 16);
            out[i] = c as i32 + b[i];
        }
    }
}

// =====================================================================
// Compressed challenge polynomial (68 bytes)
// =====================================================================

/// Maximum TAU across all ML-DSA parameter sets (ML-DSA-87: TAU=60).
const MAX_TAU: usize = 60;

/// Size of a compressed challenge: TAU indices + 8-byte sign mask.
/// Fixed at 68 bytes to accommodate all parameter sets (MAX_TAU=60,
/// padded to 60 indices + 8 bytes signs).
pub const COMPRESSED_CHALLENGE_BYTES: usize = MAX_TAU + 8;

/// Compress a challenge polynomial `c` (output of SampleInBall) into
/// 68 bytes: the first `tau` bytes are the indices of the non-zero
/// coefficients, and bytes 60..68 are the 64-bit sign mask.
///
/// `c` must have exactly `tau` non-zero coefficients in {-1, +1}.
pub fn challenge_compress(out: &mut [u8; COMPRESSED_CHALLENGE_BYTES], c: &[i32; N], tau: usize) {
    for b in out.iter_mut() {
        *b = 0;
    }
    let mut signs: u64 = 0;
    let mut mask: u64 = 1;
    let mut pos = 0;
    for i in 0..N {
        if c[i] != 0 {
            out[pos] = i as u8;
            pos += 1;
            if c[i] == -1 {
                signs |= mask;
            }
            mask <<= 1;
        }
    }
    debug_assert_eq!(pos, tau);
    for i in 0..8 {
        out[MAX_TAU + i] = (signs >> (8 * i)) as u8;
    }
}

/// Decompress a challenge polynomial from its 68-byte compressed form.
pub fn challenge_decompress(c: &mut [i32; N], comp: &[u8; COMPRESSED_CHALLENGE_BYTES], tau: usize) {
    for coeff in c.iter_mut() {
        *coeff = 0;
    }
    let mut signs: u64 = 0;
    for i in 0..8 {
        signs |= (comp[MAX_TAU + i] as u64) << (8 * i);
    }
    for idx in 0..tau {
        let pos = comp[idx] as usize;
        if signs & 1 == 1 {
            c[pos] = -1;
        } else {
            c[pos] = 1;
        }
        signs >>= 1;
    }
}

/// Schoolbook multiply: `out += c_compressed * b`, where `c` is in
/// compressed form (68 bytes) and `b` is a full polynomial.
///
/// This replaces `NTT(c) + pointwise_mul + iNTT` with a direct
/// time-domain multiplication, operating at O(N × TAU) instead of
/// O(N log N). For TAU ≤ 60 and N = 256, this is comparable in
/// cost to a single NTT.
///
/// `out` is **accumulated** (not zeroed) — call with a zeroed
/// polynomial if a fresh product is needed.
pub fn schoolbook_mul_add(out: &mut [i32; N], c_comp: &[u8; COMPRESSED_CHALLENGE_BYTES], b: &[i32; N], tau: usize) {
    let mut signs: u64 = 0;
    for i in 0..8 {
        signs |= (c_comp[MAX_TAU + i] as u64) << (8 * i);
    }

    for idx in 0..tau {
        let ci = c_comp[idx] as usize; // position of the non-zero coeff in c
        if signs & 1 == 0 {
            // c[ci] = +1 → out = out + X^ci * b  (mod X^N + 1)
            for j in 0..N {
                if ci + j < N {
                    out[ci + j] += b[j];
                } else {
                    out[ci + j - N] -= b[j];
                }
            }
        } else {
            // c[ci] = -1 → out = out - X^ci * b  (mod X^N + 1)
            for j in 0..N {
                if ci + j < N {
                    out[ci + j] -= b[j];
                } else {
                    out[ci + j - N] += b[j];
                }
            }
        }
        signs >>= 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pack_unpack_roundtrip() {
        let mut poly = [0i32; N];
        for i in 0..N {
            poly[i] = ((i as i32 * 32771 + 13) % Q) as i32;
        }
        let mut buf = [0u8; COMPRESSED_POLY_BYTES];
        poly_pack_24(&mut buf, &poly);
        let mut recovered = [0i32; N];
        poly_unpack_24(&mut recovered, &buf);
        // mod_q in pack normalizes; compare normalized
        for i in 0..N {
            assert_eq!(recovered[i], mod_q(poly[i]), "mismatch at i={}", i);
        }
    }

    #[test]
    fn challenge_compress_decompress_roundtrip() {
        // Build a challenge-like polynomial: TAU=39 non-zero coeffs in {-1,+1}
        let tau = 39;
        let mut c = [0i32; N];
        // Place TAU non-zero coefficients
        for i in (N - tau)..N {
            c[i] = if i % 3 == 0 { -1 } else { 1 };
        }
        let mut comp = [0u8; COMPRESSED_CHALLENGE_BYTES];
        challenge_compress(&mut comp, &c, tau);
        let mut recovered = [0i32; N];
        challenge_decompress(&mut recovered, &comp, tau);
        assert_eq!(c, recovered);
    }

    #[test]
    fn schoolbook_mul_matches_ntt_pointwise() {
        use super::super::ntt;
        use super::super::sample;
        use super::super::sha3;

        // Generate a challenge polynomial via SampleInBall
        let c_tilde = [0x42u8; 32];
        let tau = 39; // ML-DSA-44
        let c = sample::sample_in_ball::<super::super::params::MlDsa44>(&c_tilde);

        // Generate a random-ish polynomial b
        let mut b = [0i32; N];
        for i in 0..N {
            b[i] = ((i as i32 * 32771 + 17) % Q) as i32;
        }

        // Reference: NTT path
        let mut c_ntt = c;
        for coeff in c_ntt.iter_mut() {
            *coeff = mod_q(*coeff);
        }
        ntt::ntt(&mut c_ntt);
        let mut b_ntt = b;
        ntt::ntt(&mut b_ntt);
        let mut prod_ntt = ntt::pointwise_mul(&c_ntt, &b_ntt);
        ntt::ntt_inv(&mut prod_ntt);
        // Normalize to [0, q-1]
        for coeff in prod_ntt.iter_mut() {
            *coeff = mod_q(*coeff);
        }

        // Schoolbook path
        let mut comp = [0u8; COMPRESSED_CHALLENGE_BYTES];
        challenge_compress(&mut comp, &c, tau);
        let mut prod_school = [0i32; N];
        schoolbook_mul_add(&mut prod_school, &comp, &b, tau);
        // Normalize to [0, q-1]
        for coeff in prod_school.iter_mut() {
            *coeff = mod_q(*coeff);
        }

        assert_eq!(prod_ntt, prod_school, "schoolbook mul must match NTT pointwise mul");
    }

    #[test]
    fn compressed_vec_sub() {
        let mut v = CompressedVecK::new(2);
        let mut a = [0i32; N];
        let mut b = [0i32; N];
        for i in 0..N {
            a[i] = ((i as i32 * 999 + 7) % Q) as i32;
            b[i] = ((i as i32 * 333 + 3) % Q) as i32;
        }
        v.pack(0, &a);
        let mut out = [0i32; N];
        v.sub_into(0, &b, &mut out);
        for i in 0..N {
            assert_eq!(out[i], mod_q(a[i]) - b[i], "sub mismatch at i={}", i);
        }
    }
}