krypteia-quantica 0.1.0

//! Sampling algorithms for ML-DSA (FIPS 204, Algorithms 29-34).
//!
//! Provides functions to sample polynomials and polynomial vectors from
//! various distributions using SHAKE-based extendable output functions.
//! All sampling routines use rejection sampling to ensure uniform output.
//! All returned data lives on the stack (no heap allocations).

use super::params::{MAX_K, MAX_L, N, Params, Q};
use super::sha3;

/// CoeffFromThreeBytes (Algorithm 14): rejection-sample a coefficient from 3 bytes.
/// Returns Some(coeff) if valid, None otherwise.
#[inline]
fn coeff_from_three_bytes(b0: u8, b1: u8, b2: u8) -> Option<i32> {
    let b2_masked = b2 & 0x7F; // top bit is discarded
    let z = (b0 as i32) | ((b1 as i32) << 8) | ((b2_masked as i32) << 16);
    if z < Q { Some(z) } else { None }
}

/// CoeffFromHalfByte (Algorithm 15): map a half-byte to a coefficient in [-eta, eta].
/// Returns Some(coeff) if valid, None otherwise.
#[inline]
fn coeff_from_half_byte(b: u8, eta: usize) -> Option<i32> {
    let b = b as i32;
    if eta == 2 {
        if b < 15 { Some(2 - (b % 5)) } else { None }
    } else {
        // eta == 4
        if b < 9 { Some(4 - b) } else { None }
    }
}

/// Sample a sparse challenge polynomial with exactly tau non-zero entries.
///
/// Implements Algorithm 29 of FIPS 204 (SampleInBall). The output polynomial
/// `c` has exactly `P::TAU` coefficients equal to +/-1 (the rest are 0).
/// Signs are determined by squeezing 8 bytes of sign bits from SHAKE256,
/// and positions are chosen via rejection sampling to ensure uniformity.
///
/// - `c_tilde`: commitment hash seed (lambda/4 bytes).
///
/// Returns a polynomial with coefficients in {-1, 0, 1}.
pub fn sample_in_ball<P: Params>(c_tilde: &[u8]) -> [i32; N] {
    let tau = P::TAU;
    let mut c = [0i32; N];

    let mut state = sha3::shake256();
    state.absorb(c_tilde);

    // Get 8 bytes for sign bits
    let mut sign_bytes = [0u8; 8];
    state.squeeze(&mut sign_bytes);
    let signs = u64::from_le_bytes(sign_bytes);

    let mut sign_idx = 0usize;
    for i in (N - tau)..N {
        // Sample j uniformly from [0, i]
        let mut j;
        loop {
            let mut buf = [0u8; 1];
            state.squeeze(&mut buf);
            j = buf[0] as usize;
            if j <= i {
                break;
            }
        }
        c[i] = c[j];
        let sign_bit = (signs >> sign_idx) & 1;
        c[j] = if sign_bit == 1 { -1 } else { 1 };
        sign_idx += 1;
    }

    c
}

/// Generate an NTT-domain polynomial via rejection sampling.
///
/// Implements Algorithm 30 of FIPS 204 (RejNTTPoly). Samples coefficients
/// uniformly from [0, q) by reading 3 bytes at a time from a SHAKE128
/// stream seeded with `rho || j1 || j2`. Candidates >= q are rejected.
///
/// - `rho`: 32-byte public seed.
/// - `j1`: column index of the matrix entry (s index).
/// - `j2`: row index of the matrix entry (r index).
///
/// Returns a polynomial in NTT domain with coefficients in [0, q-1].
pub fn rej_ntt_poly(rho: &[u8; 32], j1: u8, j2: u8) -> [i32; N] {
    let mut a = [0i32; N];
    let mut state = sha3::shake128();
    state.absorb(rho);
    state.absorb(&[j1, j2]);

    let mut idx = 0;
    while idx < N {
        let mut buf = [0u8; 3];
        state.squeeze(&mut buf);
        if let Some(coeff) = coeff_from_three_bytes(buf[0], buf[1], buf[2]) {
            a[idx] = coeff;
            idx += 1;
        }
    }
    a
}

/// Generate a polynomial with small coefficients via rejection sampling.
///
/// Implements Algorithm 31 of FIPS 204 (RejBoundedPoly). Samples
/// coefficients in [-eta, eta] by reading half-bytes from a SHAKE256
/// stream seeded with `rho_prime || nonce`. Invalid half-byte values are
/// rejected.
///
/// - `rho_prime`: 64-byte secret seed.
/// - `nonce`: 16-bit counter distinguishing different polynomials.
/// - `eta`: coefficient bound (2 or 4 depending on parameter set).
///
/// Returns a polynomial with coefficients in [-eta, eta].
pub fn rej_bounded_poly(rho_prime: &[u8; 64], nonce: u16, eta: usize) -> [i32; N] {
    let mut a = [0i32; N];
    let mut state = sha3::shake256();
    state.absorb(rho_prime);
    state.absorb(&nonce.to_le_bytes());

    let mut idx = 0;
    while idx < N {
        let mut buf = [0u8; 1];
        state.squeeze(&mut buf);
        let z0 = coeff_from_half_byte(buf[0] & 0x0F, eta);
        let z1 = coeff_from_half_byte(buf[0] >> 4, eta);
        if let Some(c) = z0 {
            if idx < N {
                a[idx] = c;
                idx += 1;
            }
        }
        if let Some(c) = z1 {
            if idx < N {
                a[idx] = c;
                idx += 1;
            }
        }
    }
    a
}

/// Expand the public matrix A in NTT domain from a seed.
///
/// Implements Algorithm 32 of FIPS 204 (ExpandA). Generates a k-by-l matrix
/// of NTT-domain polynomials by calling [`rej_ntt_poly`] for each entry
/// with indices `(s, r)`.
///
/// - `rho`: 32-byte public seed (part of the public key).
///
/// Returns the matrix A-hat as a `[[[i32; N]; MAX_L]; MAX_K]` with valid
/// entries in rows 0..k and columns 0..l.
pub fn expand_a<P: Params>(rho: &[u8; 32]) -> [[[i32; N]; MAX_L]; MAX_K] {
    let k = P::K;
    let l = P::L;
    let mut a_hat = [[[0i32; N]; MAX_L]; MAX_K];
    for r in 0..k {
        for s in 0..l {
            a_hat[r][s] = rej_ntt_poly(rho, s as u8, r as u8);
        }
    }
    a_hat
}

/// Expand the secret vectors s1 and s2 from a seed.
///
/// Implements Algorithm 33 of FIPS 204 (ExpandS). Generates the secret
/// vectors by calling [`rej_bounded_poly`] with incrementing nonces:
/// nonces 0..l for s1, and l..(l+k) for s2.
///
/// - `rho_prime`: 64-byte secret seed derived during key generation.
///
/// Returns `(s1, s2)` where s1 has MAX_L polynomials (valid 0..l) and s2
/// has MAX_K polynomials (valid 0..k), each with coefficients in [-eta, eta].
pub fn expand_s<P: Params>(rho_prime: &[u8; 64]) -> ([[i32; N]; MAX_L], [[i32; N]; MAX_K]) {
    let k = P::K;
    let l = P::L;
    let eta = P::ETA;
    let mut s1 = [[0i32; N]; MAX_L];
    for s in 0..l {
        s1[s] = rej_bounded_poly(rho_prime, s as u16, eta);
    }
    let mut s2 = [[0i32; N]; MAX_K];
    for s in 0..k {
        s2[s] = rej_bounded_poly(rho_prime, (l + s) as u16, eta);
    }
    (s1, s2)
}

/// Expand the masking vector y from a seed and counter.
///
/// Implements Algorithm 34 of FIPS 204 (ExpandMask). Generates l polynomials
/// with coefficients in [-(gamma1-1), gamma1] by squeezing SHAKE256 output
/// and unpacking via `bit_unpack`. Each polynomial uses a distinct nonce
/// derived from `kappa`.
///
/// - `rho_double_prime`: 64-byte seed derived from the secret key and randomness.
/// - `kappa`: counter incremented by l on each rejection loop iteration.
///
/// Returns a fixed array of MAX_L polynomials (valid entries 0..l).
pub fn expand_mask<P: Params>(rho_double_prime: &[u8; 64], kappa: u16) -> [[i32; N]; MAX_L] {
    let l = P::L;
    let gamma1 = P::GAMMA1 as u32;
    let c = P::BITLEN_GAMMA1_MINUS1 + 1;
    let poly_bytes = 32 * c;
    let mut y = [[0i32; N]; MAX_L];

    for r in 0..l {
        let nonce = kappa + r as u16;
        let mut state = sha3::shake256();
        state.absorb(rho_double_prime);
        state.absorb(&nonce.to_le_bytes());
        // Max poly_bytes = 32 * 20 = 640
        let mut buf = [0u8; 640];
        state.squeeze(&mut buf[..poly_bytes]);

        super::encode::bit_unpack(&buf[..poly_bytes], gamma1 - 1, gamma1, &mut y[r]);
    }
    y
}