quantica/ml_kem/

encode.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Cédric Mesnil <cslashm@pm.me>

// Index-based loops in this module deliberately mirror the reference
// algorithm (parallel-array access `a[i]`/`b[i]`, or a position-indexed
// packing/reduction state) and keep the constant-time data layout explicit;
// the `iter().enumerate()` rewrite would obscure it. Documented, single-lint
// allow per CLAUDE.md §6 (scoped to this module, not a blanket suppression).
#![allow(clippy::needless_range_loop)]
// Complete FIPS-algorithm / countermeasure family: with the module now
// crate-internal (API re-tiering, v0.2), not every entry point has an
// in-crate caller in every feature combination. All of them light up under
// --all-features and are exercised by the KAT suites; keeping the full,
// audited family beats pruning it per-combo. Documented allow (CLAUDE.md §6).
#![allow(dead_code)]

//! # Stability (hazmat tier)
//!
//! This module is part of the `hazmat` expert building-block layer: it is
//! public only with the `hazmat` Cargo feature, for KAT/test-vector tooling
//! and research. Read each function's invariants before use. **No stability,
//! misuse-resistance, or side-channel promise** applies to this tier — its
//! API may change in any release. The typed facade is the guaranteed API.

/// Encoding, decoding, compression, and decompression algorithms.
///
/// Implements FIPS 203 Section 4.2.1, Algorithms 3-6:
///
/// - [`byte_encode`] / [`byte_decode`] -- bit-pack/unpack 256 integers into/from bytes
/// - [`compress`] / [`decompress`] -- lossy compression between Z_q and Z_{2^d}
/// - [`compress_poly`] / [`decompress_poly`] -- vectorized compression over full polynomials
use super::params::{N, Q};

/// Bit-pack 256 integers into `32*d` bytes (Algorithm 5: ByteEncode_d).
///
/// Each of the 256 coefficients in `f` is encoded using `d` bits.
/// For `d < 12`, coefficients are treated modulo `2^d`.
/// For `d = 12`, coefficients are treated modulo q = 3329.
///
/// # Arguments
///
/// * `d` - The bit-width per coefficient (1..=12).
/// * `f` - Array of 256 unsigned 16-bit coefficients.
/// * `out` - Output buffer of exactly `32 * d` bytes.
///
/// # Panics
///
/// Debug-asserts that `out.len() == 32 * d`.
pub fn byte_encode(d: usize, f: &[u16; N], out: &mut [u8]) {
    debug_assert_eq!(out.len(), 32 * d);

    if d < 12 {
        // Bit-packing: each coefficient is d bits
        let mut bit_idx = 0usize;
        for i in 0..N {
            let mut val = f[i] as u32;
            for _ in 0..d {
                let byte_pos = bit_idx >> 3;
                let bit_pos = bit_idx & 7;
                if bit_pos == 0 && byte_pos < out.len() {
                    out[byte_pos] = 0;
                }
                if byte_pos < out.len() {
                    out[byte_pos] |= ((val & 1) as u8) << bit_pos;
                }
                val >>= 1;
                bit_idx += 1;
            }
        }
    } else {
        // d = 12: same logic but coefficients are mod q
        let mut bit_idx = 0usize;
        for b in out.iter_mut() {
            *b = 0;
        }
        for i in 0..N {
            let mut val = f[i] as u32;
            for _ in 0..12 {
                let byte_pos = bit_idx >> 3;
                let bit_pos = bit_idx & 7;
                if byte_pos < out.len() {
                    out[byte_pos] |= ((val & 1) as u8) << bit_pos;
                }
                val >>= 1;
                bit_idx += 1;
            }
        }
    }
}

/// Unpack `32*d` bytes into 256 integers (Algorithm 6: ByteDecode_d).
///
/// The inverse of [`byte_encode`]. Each coefficient is extracted from `d`
/// consecutive bits and reduced modulo `2^d` (for `d < 12`) or modulo
/// q = 3329 (for `d = 12`).
///
/// # Arguments
///
/// * `d` - The bit-width per coefficient (1..=12).
/// * `input` - Input buffer of exactly `32 * d` bytes.
/// * `f` - Output array of 256 unsigned 16-bit coefficients.
///
/// # Panics
///
/// Debug-asserts that `input.len() == 32 * d`.
pub fn byte_decode(d: usize, input: &[u8], f: &mut [u16; N]) {
    debug_assert_eq!(input.len(), 32 * d);

    let m = if d < 12 { 1u32 << d } else { Q as u32 };

    let mut bit_idx = 0usize;
    for i in 0..N {
        let mut val = 0u32;
        for j in 0..d {
            let byte_pos = bit_idx >> 3;
            let bit_pos = bit_idx & 7;
            val |= (((input[byte_pos] >> bit_pos) & 1) as u32) << j;
            bit_idx += 1;
        }
        f[i] = (val % m) as u16;
    }
}

/// Lossy compression from Z_q to Z_{2^d} (FIPS 203 eq. 4.7).
///
/// Computes `x -> round(2^d / q * x) mod 2^d` using integer-only
/// arithmetic: `floor((2^d * x + q/2) / q) mod 2^d`.
///
/// # Arguments
///
/// * `d` - Target bit-width (1..=12).
/// * `x` - A coefficient in `[0, q-1]`.
///
/// # Returns
///
/// The compressed value in `[0, 2^d - 1]`.
#[inline(always)]
pub fn compress(d: u32, x: u16) -> u16 {
    // ⌈(2^d / q) · x⌋ = ⌊(2^d · x + q/2) / q⌋ mod 2^d
    let shifted = ((x as u64) << d) + (Q as u64 / 2);
    ((shifted / Q as u64) & ((1u64 << d) - 1)) as u16
}

/// Decompression from Z_{2^d} back to Z_q (FIPS 203 eq. 4.8).
///
/// Computes `y -> round(q / 2^d * y)` using integer-only arithmetic:
/// `floor((q * y + 2^{d-1}) / 2^d)`.
///
/// This is the approximate inverse of [`compress`]. The round-trip
/// introduces a bounded quantization error.
///
/// # Arguments
///
/// * `d` - Source bit-width (1..=12).
/// * `y` - A compressed value in `[0, 2^d - 1]`.
///
/// # Returns
///
/// The decompressed value in `[0, q-1]`.
#[inline(always)]
pub fn decompress(d: u32, y: u16) -> u16 {
    // ⌈(q / 2^d) · y⌋ = ⌊(q · y + 2^(d-1)) / 2^d⌋
    let val = (Q as u32 * y as u32 + (1u32 << (d - 1))) >> d;
    val as u16
}

/// Compress all 256 coefficients of a polynomial from Z_q to Z_{2^d}.
///
/// Applies [`compress`] element-wise.
///
/// # Arguments
///
/// * `d` - Target bit-width.
/// * `f` - Input polynomial with coefficients in `[0, q-1]`.
/// * `out` - Output polynomial with coefficients in `[0, 2^d - 1]`.
pub fn compress_poly(d: u32, f: &[u16; N], out: &mut [u16; N]) {
    for i in 0..N {
        out[i] = compress(d, f[i]);
    }
}

/// Decompress all 256 coefficients of a polynomial from Z_{2^d} to Z_q.
///
/// Applies [`decompress`] element-wise.
///
/// # Arguments
///
/// * `d` - Source bit-width.
/// * `f` - Input polynomial with coefficients in `[0, 2^d - 1]`.
/// * `out` - Output polynomial with coefficients in `[0, q-1]`.
pub fn decompress_poly(d: u32, f: &[u16; N], out: &mut [u16; N]) {
    for i in 0..N {
        out[i] = decompress(d, f[i]);
    }
}