krypteia-quantica 0.1.0

/// Number-Theoretic Transform for ML-KEM (FIPS 203 Section 4.3).
///
/// Full Montgomery arithmetic: all NTT butterflies, basemul, and iNTT use
/// Montgomery multiplication (shifts instead of divisions).
///
/// q = 3329, R = 2^16. Coefficients stored as i16 with lazy reduction.
/// Bounds: after NTT |c| <= 4q, after basemul |c| <= q, after iNTT |c| <= 8q.
/// All bounds fit in i16 range [-32768, 32767] since 8q = 26632 < 32768.
use super::params::N;

// ---- Constants ----

pub const Q: i16 = 3329;
const Q32: i32 = 3329;
const QINV: i32 = -3327; // q^{-1} mod 2^16 (signed)
const R_SQ_MOD_Q: i32 = 1353; // R² mod q

/// Zetas: ζ^{BitRev_7(i)} mod q (FIPS 203 Appendix A)
const ZETAS: [i16; 128] = [
    1, 1729, 2580, 3289, 2642, 630, 1897, 848, 1062, 1919, 193, 797, 2786, 3260, 569, 1746, 296, 2447, 1339, 1476,
    3046, 56, 2240, 1333, 1426, 2094, 535, 2882, 2393, 2879, 1974, 821, 289, 331, 3253, 1756, 1197, 2304, 2277, 2055,
    650, 1977, 2513, 632, 2865, 33, 1320, 1915, 2319, 1435, 807, 452, 1438, 2868, 1534, 2402, 2647, 2617, 1481, 648,
    2474, 3110, 1227, 910, 17, 2761, 583, 2649, 1637, 723, 2288, 1100, 1409, 2662, 3281, 233, 756, 2156, 3015, 3050,
    1703, 1651, 2789, 1789, 1847, 952, 1461, 2687, 939, 2308, 2437, 2388, 733, 2337, 268, 641, 1584, 2298, 2037, 3220,
    375, 2549, 2090, 1645, 1063, 319, 2773, 757, 2099, 561, 2466, 2594, 2804, 1092, 403, 1026, 1143, 2150, 2775, 886,
    1722, 1212, 1874, 1029, 2110, 2935, 885, 2154,
];

/// Gammas: γᵢ = ζ^{2·BitRev_7(i)+1} mod q (FIPS 203 Appendix A)
const GAMMAS: [i16; 128] = [
    17, 3312, 2761, 568, 583, 2746, 2649, 680, 1637, 1692, 723, 2606, 2288, 1041, 1100, 2229, 1409, 1920, 2662, 667,
    3281, 48, 233, 3096, 756, 2573, 2156, 1173, 3015, 314, 3050, 279, 1703, 1626, 1651, 1678, 2789, 540, 1789, 1540,
    1847, 1482, 952, 2377, 1461, 1868, 2687, 642, 939, 2390, 2308, 1021, 2437, 892, 2388, 941, 733, 2596, 2337, 992,
    268, 3061, 641, 2688, 1584, 1745, 2298, 1031, 2037, 1292, 3220, 109, 375, 2954, 2549, 780, 2090, 1239, 1645, 1684,
    1063, 2266, 319, 3010, 2773, 556, 757, 2572, 2099, 1230, 561, 2768, 2466, 863, 2594, 735, 2804, 525, 1092, 2237,
    403, 2926, 1026, 2303, 1143, 2186, 2150, 1179, 2775, 554, 886, 2443, 1722, 1607, 1212, 2117, 1874, 1455, 1029,
    2300, 2110, 1219, 2935, 394, 885, 2444, 2154, 1175,
];

// ---- Compile-time precomputation ----

const fn mont_reduce_const(a: i32) -> i16 {
    let t = (a as i16).wrapping_mul(QINV as i16);
    ((a - t as i32 * Q32) >> 16) as i16
}

const fn to_mont_const(a: i16) -> i16 {
    mont_reduce_const(a as i32 * R_SQ_MOD_Q)
}

const fn compute_table_mont(table: &[i16; 128]) -> [i16; 128] {
    let mut r = [0i16; 128];
    let mut i = 0;
    while i < 128 {
        r[i] = to_mont_const(table[i]);
        i += 1;
    }
    r
}

const ZETAS_MONT: [i16; 128] = compute_table_mont(&ZETAS);
const GAMMAS_MONT: [i16; 128] = compute_table_mont(&GAMMAS);

/// iNTT scaling factor: R² · 128⁻¹ mod q = 1441.
/// Compensates: 128⁻¹ (NTT normalization) + R⁻¹ (from basemul via mont_mul).
const F_SCALE: i16 = 1441;

// ---- Montgomery arithmetic ----

/// Montgomery reduction: a·R⁻¹ mod q. Result bounded by |q|.
#[inline(always)]
fn montgomery_reduce(a: i32) -> i16 {
    let t = (a as i16).wrapping_mul(QINV as i16);
    ((a - t as i32 * Q32) >> 16) as i16
}

/// Montgomery multiply: a·b·R⁻¹ mod q.
#[inline(always)]
fn mont_mul(a: i16, b: i16) -> i16 {
    montgomery_reduce(a as i32 * b as i32)
}

/// Barrett reduction to [0, q-1]. Constant-time.
#[inline(always)]
pub fn barrett_reduce(a: i16) -> i16 {
    let t = ((20159i32 * a as i32 + (1 << 25)) >> 26) as i16;
    let mut r = a - t.wrapping_mul(Q);
    r += (r >> 15) & Q;
    r
}

// ---- NTT ----

/// Forward NTT (Algorithm 9) — full Montgomery.
///
/// Input: coefficients in [0, q-1].
/// Output: NTT-domain coefficients, |c| <= 4q (lazy reduced).
pub fn ntt(f: &mut [i16; N]) {
    let mut k = 1usize;
    let mut len = 128;
    while len >= 2 {
        let mut start = 0;
        while start < N {
            let zeta = ZETAS_MONT[k];
            k += 1;
            for j in start..start + len {
                let t = mont_mul(zeta, f[j + len]);
                f[j + len] = f[j] - t;
                f[j] = f[j] + t;
            }
            start += 2 * len;
        }
        len >>= 1;
    }
    // Reduce output to [0,q-1] for basemul
    for c in f.iter_mut() {
        *c = barrett_reduce(*c);
    }
}

/// Inverse NTT (Algorithm 10) — full Montgomery.
///
/// Input: NTT-domain (from basemul, in /R domain), |c| <= q.
/// Output: standard-domain, |c| <= q after final scaling.
///
/// Uses **negative** Montgomery zetas: `-ZETAS_MONT[k]`.
/// GS butterfly: `f[j] = t + f[j+l]`, `f[j+l] = fqmul(-z, t - f[j+l])`.
/// Final scaling by F_SCALE = 1441 = R²·128⁻¹ mod q.
pub fn ntt_inv(f: &mut [i16; N]) {
    // Use i32 working array to avoid i16 overflow in butterfly additions.
    // The accumulation through 7 GS levels can reach ~8q ≈ 26632,
    // which fits i16 but leaves no margin. Using i32 eliminates this concern.
    let mut w = [0i32; N];
    for i in 0..N {
        w[i] = f[i] as i32;
    }

    let mut k = 127usize;
    let mut len = 2;
    while len <= 128 {
        let mut start = 0;
        while start < N {
            let neg_zeta = ZETAS_MONT[k].wrapping_neg();
            k = k.wrapping_sub(1);
            for j in start..start + len {
                let t = w[j];
                let u = w[j + len];
                w[j] = t + u;
                // mont_reduce(neg_zeta * (t-u)) with i32 diff
                w[j + len] = montgomery_reduce(neg_zeta as i32 * ((t - u) as i32)) as i32;
            }
            start += 2 * len;
        }
        len <<= 1;
    }
    for i in 0..N {
        // w[i] can exceed i16 range; use montgomery_reduce with i32 multiply
        f[i] = montgomery_reduce(F_SCALE as i32 * w[i]) as i16;
    }
}

/// Pointwise basemul in NTT domain (Algorithms 11+12) — full Montgomery.
///
/// Each pair `(h[2i], h[2i+1])` is the product of input pairs modulo `(X²-γᵢ)`.
/// Output is in **/R domain** (one Montgomery division per coefficient).
///
/// For the encrypt/decrypt pipeline (basemul → iNTT), the /R is compensated
/// by the F_SCALE factor in iNTT.
///
/// For KeyGen (basemul accumulate in NTT domain), call [`to_mont_poly`] to
/// convert from /R to normal domain before adding NTT(e).
pub fn multiply_ntts(f: &[i16; N], g: &[i16; N], h: &mut [i16; N]) {
    for i in 0..128 {
        let gamma = GAMMAS_MONT[i];
        let a0 = f[2 * i];
        let a1 = f[2 * i + 1];
        let b0 = g[2 * i];
        let b1 = g[2 * i + 1];
        // c0 = a0·b0·R⁻¹ + a1·b1·R⁻¹·γ·R·R⁻¹ = (a0·b0 + a1·b1·γ)·R⁻¹
        let t = mont_mul(mont_mul(a1, b1), gamma);
        h[2 * i] = mont_mul(a0, b0) + t;
        // c1 = (a0·b1 + a1·b0)·R⁻¹
        h[2 * i + 1] = mont_mul(a0, b1) + mont_mul(a1, b0);
    }
}

/// Convert polynomial from /R domain to normal domain (multiply by R).
///
/// Used in KeyGen after accumulating basemul results, before adding NTT(e).
/// `to_mont(c) = c·R mod q` via Montgomery: `mont_reduce(c · R²_mod_q)`.
pub fn to_mont_poly(f: &mut [i16; N]) {
    for c in f.iter_mut() {
        *c = montgomery_reduce(*c as i32 * R_SQ_MOD_Q);
    }
}

// ---- Polynomial helpers ----

pub fn poly_add(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
    for i in 0..N {
        c[i] = a[i] + b[i];
    }
}

pub fn poly_sub(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
    for i in 0..N {
        c[i] = a[i] - b[i];
    }
}

/// Reduce all coefficients to [0, q-1].
pub fn reduce(f: &mut [i16; N]) {
    for c in f.iter_mut() {
        *c = barrett_reduce(*c);
    }
}

// ---- Zeroize ----

#[inline(never)]
pub fn zeroize_poly(f: &mut [i16; N]) {
    for c in f.iter_mut() {
        unsafe { core::ptr::write_volatile(c, 0) };
    }
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

#[inline(never)]
pub fn zeroize_bytes(b: &mut [u8]) {
    for byte in b.iter_mut() {
        unsafe { core::ptr::write_volatile(byte, 0) };
    }
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
}

// ---- Tests ----

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

    #[test]
    fn test_ntt_basemul_intt_simple() {
        // f(X) = 42, g(X) = 1 + X → product = 42 + 42X
        let mut f = [0i16; N];
        f[0] = 42;
        let mut g = [0i16; N];
        g[0] = 1;
        g[1] = 1;

        ntt(&mut f);
        ntt(&mut g);

        let mut h = [0i16; N];
        multiply_ntts(&f, &g, &mut h);
        ntt_inv(&mut h);
        reduce(&mut h);

        assert_eq!(h[0], 42, "h[0]={} expected 42", h[0]);
        assert_eq!(h[1], 42, "h[1]={} expected 42", h[1]);
        for i in 2..N {
            assert_eq!(h[i], 0, "h[{}]={} expected 0", i, h[i]);
        }
    }

    #[test]
    fn test_ntt_basemul_intt_identity() {
        // NTT(1) × NTT(b) → iNTT should recover b
        let mut one = [0i16; N];
        one[0] = 1;
        let mut b = [0i16; N];
        for i in 0..N {
            b[i] = (i as i16 * 7 + 13) % Q;
        }
        let orig = b;

        ntt(&mut one);
        ntt(&mut b);
        let mut c = [0i16; N];
        multiply_ntts(&one, &b, &mut c);
        ntt_inv(&mut c);
        reduce(&mut c);

        for i in 0..N {
            assert_eq!(c[i], orig[i], "mismatch at {}: got {} expected {}", i, c[i], orig[i]);
        }
    }

    #[test]
    fn test_to_mont_poly_keygen_pattern() {
        // KeyGen: basemul accumulate → to_mont → add e
        let mut a = [0i16; N];
        a[0] = 100;
        let mut s = [0i16; N];
        s[0] = 1;
        let mut e = [0i16; N];
        e[0] = 5;

        ntt(&mut a);
        ntt(&mut s);
        ntt(&mut e);

        let mut t = [0i16; N];
        multiply_ntts(&a, &s, &mut t);
        // t is in /R domain
        to_mont_poly(&mut t);
        // t is now in normal domain, safe to add e
        for i in 0..N {
            t[i] = t[i] + e[i];
        }

        // Verify: decode via exact iNTT of t should give [105, 0, ...]
        // But we can't directly iNTT (it expects /R input).
        // Instead verify the encoded t_hat matches what exact arithmetic gives.
        reduce(&mut t);
        // t[0] should be the NTT-domain value corresponding to polynomial 105
        // The exact test is done by the ACVP KeyGen vectors.
    }

    #[test]
    fn test_montgomery_reduce_basic() {
        let a: i16 = 1729;
        let in_mont = to_mont_const(a);
        let back = montgomery_reduce(in_mont as i32);
        let back_reduced = barrett_reduce(back);
        assert_eq!(back_reduced, a);
    }
}