krypteia-quantica 0.1.0

//! Decomposition algorithms for ML-DSA (FIPS 204, Algorithms 35-40).
//!
//! Implements the Decompose, HighBits, LowBits, MakeHint, and UseHint
//! functions used during signing and verification to split polynomial
//! coefficients into high-order and low-order parts with respect to
//! `alpha = 2 * gamma2`.
//!
//! All vector functions operate on slices and fixed arrays -- no heap
//! allocations.

use super::params::{MAX_K, N, Q};

/// Decompose a coefficient into high and low parts.
///
/// Implements Algorithm 36 of FIPS 204 (Decompose). Given `r` in [0, q-1]
/// and `gamma2`, computes `(r1, r0)` such that `r = r1 * alpha + r0` where
/// `alpha = 2 * gamma2` and `r0` lies in the centered range
/// `(-alpha/2, alpha/2]`. A special case handles `r1 * alpha == q - 1` to
/// keep r1 in the valid range.
///
/// - `r`: coefficient to decompose (should be in [0, q-1]).
/// - `gamma2`: half the decomposition modulus (parameter-set dependent).
///
/// Returns `(r1, r0)`.
pub fn decompose(r: i32, gamma2: i32) -> (i32, i32) {
    let alpha = 2 * gamma2;
    let rp = r % Q + ((r >> 31) & Q); // ensure positive
    let mut r0 = rp % alpha; // [0, alpha-1]
    if r0 > alpha / 2 {
        r0 -= alpha;
    }
    if rp - r0 == Q - 1 {
        (0, r0 - 1)
    } else {
        ((rp - r0) / alpha, r0)
    }
}

/// Extract the high-order representative of a coefficient.
///
/// Implements Algorithm 37 of FIPS 204 (HighBits). Returns the r1 component
/// from [`decompose`]. This is used to compute the commitment w1 during
/// signing and verification.
pub fn high_bits(r: i32, gamma2: i32) -> i32 {
    decompose(r, gamma2).0
}

/// Extract the low-order representative of a coefficient.
///
/// Implements Algorithm 38 of FIPS 204 (LowBits). Returns the r0 component
/// from [`decompose`]. The norm of r0 is checked during signing to decide
/// whether to accept or reject a candidate signature.
pub fn low_bits(r: i32, gamma2: i32) -> i32 {
    decompose(r, gamma2).1
}

/// Compute a single hint bit.
///
/// Implements Algorithm 39 of FIPS 204 (MakeHint). Returns 1 if adding `z`
/// to `r` changes the high bits (i.e., `HighBits(r + z) != HighBits(r)`),
/// and 0 otherwise. The hint allows the verifier to recover w1 without
/// knowing the secret low-order information.
///
/// - `z`: perturbation value (typically `-ct0[i][j]`).
/// - `r`: base value (typically `(w - cs2 + ct0)[i][j]`).
/// - `gamma2`: decomposition parameter.
pub fn make_hint(z: i32, r: i32, gamma2: i32) -> i32 {
    let r1 = high_bits(r, gamma2);
    let v1 = high_bits((r + z) % Q + (((r + z) >> 31) & Q), gamma2);
    if r1 != v1 { 1 } else { 0 }
}

/// Use a hint bit to recover the correct high bits.
///
/// Implements Algorithm 40 of FIPS 204 (UseHint). When `h == 0`, returns
/// `HighBits(r)` directly. When `h == 1`, adjusts the high-bits value by
/// +/-1 (modulo m = (q-1)/alpha) depending on the sign of the low bits.
///
/// - `h`: hint bit (0 or 1).
/// - `r`: the approximate value to round.
/// - `gamma2`: decomposition parameter.
///
/// Returns the corrected high-order representative.
pub fn use_hint(h: i32, r: i32, gamma2: i32) -> i32 {
    let alpha = 2 * gamma2;
    let (r1, r0) = decompose(r, gamma2);
    if h == 0 {
        return r1;
    }
    // m = (q-1) / alpha
    let m = (Q - 1) / alpha;
    if r0 > 0 { (r1 + 1) % m } else { (r1 - 1 + m) % m }
}

/// Apply [`high_bits`] to every coefficient of a polynomial vector.
///
/// Writes the high-order parts into `out`. Only the first `len` polynomials
/// are processed.
pub fn high_bits_vec(w: &[[i32; N]], gamma2: i32, out: &mut [[i32; N]], len: usize) {
    for i in 0..len {
        for j in 0..N {
            out[i][j] = high_bits(w[i][j], gamma2);
        }
    }
}

/// Apply [`low_bits`] to every coefficient of a polynomial vector.
///
/// Writes the low-order parts into `out`. Only the first `len` polynomials
/// are processed.
pub fn low_bits_vec(w: &[[i32; N]], gamma2: i32, out: &mut [[i32; N]], len: usize) {
    for i in 0..len {
        for j in 0..N {
            out[i][j] = low_bits(w[i][j], gamma2);
        }
    }
}

/// Apply [`make_hint`] to every coefficient of two polynomial vectors.
///
/// Returns `(h, count)` where `h` is the fixed array of hint polynomials and
/// `count` is the total number of 1-bits across all polynomials. The count
/// is compared against `omega` to decide whether to accept the signature.
pub fn make_hint_vec(z: &[[i32; N]], r: &[[i32; N]], gamma2: i32, len: usize) -> ([[i32; N]; MAX_K], usize) {
    let mut h = [[0i32; N]; MAX_K];
    let mut count = 0usize;
    for i in 0..len {
        for j in 0..N {
            h[i][j] = make_hint(z[i][j], r[i][j], gamma2);
            count += h[i][j] as usize;
        }
    }
    (h, count)
}

/// Apply [`use_hint`] to every coefficient of a hint vector and an approximate vector.
///
/// Returns the corrected w1 vector used during verification.
pub fn use_hint_vec(h: &[[i32; N]], r: &[[i32; N]], gamma2: i32, len: usize) -> [[i32; N]; MAX_K] {
    let mut w1 = [[0i32; N]; MAX_K];
    for i in 0..len {
        for j in 0..N {
            w1[i][j] = use_hint(h[i][j], r[i][j], gamma2);
        }
    }
    w1
}

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

    #[test]
    fn test_decompose_roundtrip() {
        let gamma2 = (Q - 1) / 88;
        let alpha = 2 * gamma2;
        for r in [0, 1, 100000, Q / 2, Q - 1] {
            let (r1, r0) = decompose(r, gamma2);
            let reconstructed = (r1 * alpha + r0) % Q + (((r1 * alpha + r0) >> 31) & Q);
            assert_eq!(reconstructed, r, "decompose roundtrip failed for r={}", r);
        }
    }

    #[test]
    fn test_use_hint_consistency() {
        let gamma2 = (Q - 1) / 88;
        let r = 1234567;
        let r1 = high_bits(r, gamma2);
        // With hint=0, use_hint should return r1
        assert_eq!(use_hint(0, r, gamma2), r1);
    }
}