quantica/ml_dsa/

ntt.rs

//! Number Theoretic Transform for ML-DSA (FIPS 204, Algorithms 41-45).
//!
//! Implements the forward and inverse NTT over `Z_q[X]/(X^256 + 1)` with
//! q = 8380417 and zeta = 1753 (a primitive 512th root of unity mod q).
//! The NTT uses 8-bit reversal (BitRev_8) and goes all the way down to
//! length-1 butterflies, so pointwise multiplication is simple element-wise
//! multiplication.

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

const Q64: i64 = Q as i64;

/// Montgomery constant: q^{-1} mod 2^32.
const QINV: i32 = 58728449;
/// R^2 mod q where R = 2^32.
const R_SQ_MOD_Q: i64 = 2365951;

/// Reduce `a` modulo q to the range [0, q-1].
#[inline(always)]
pub fn mod_q(a: i32) -> i32 {
    let mut r = a % Q;
    r += Q & (r >> 31);
    r
}

/// Modular multiplication: `(a * b) mod q`.
#[inline(always)]
pub fn mul_mod_q(a: i32, b: i32) -> i32 {
    mod_q(((a as i64 * b as i64) % Q64) as i32)
}

/// Montgomery reduction: a * R^{-1} mod q, where R = 2^32.
/// Input: |a| < q * R. Output: roughly [-q, q].
#[inline(always)]
fn montgomery_reduce(a: i64) -> i32 {
    let t = (a as i32).wrapping_mul(QINV) as i64;
    ((a - t * Q64) >> 32) as i32
}

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

/// Convert to Montgomery domain: a * R mod q.
#[inline(always)]
fn to_mont(a: i32) -> i32 {
    montgomery_reduce(a as i64 * R_SQ_MOD_Q)
}

/// Precomputed zetas in Montgomery domain (compile-time).
const fn compute_zetas_mont() -> [i32; N] {
    let mut table = [0i32; N];
    let mut k = 0;
    while k < N {
        let rev = bitrev8(k as u8) as u64;
        let z = pow_mod(1753, rev, Q64);
        // to_mont(z): montgomery_reduce(z * R_SQ_MOD_Q)
        let a = z as i64 * R_SQ_MOD_Q;
        let t = (a as i32).wrapping_mul(QINV) as i64;
        table[k] = ((a - t * Q64) >> 32) as i32;
        k += 1;
    }
    table
}

const ZETAS_MONT: [i32; N] = compute_zetas_mont();

/// Reverse 8 bits of k.
const fn bitrev8(k: u8) -> u8 {
    let mut r = 0u8;
    let mut v = k;
    let mut i = 0;
    while i < 8 {
        r = (r << 1) | (v & 1);
        v >>= 1;
        i += 1;
    }
    r
}

/// Modular exponentiation: base^exp mod m (const fn).
const fn pow_mod(mut base: i64, mut exp: u64, m: i64) -> i64 {
    let mut result = 1i64;
    base %= m;
    while exp > 0 {
        if exp & 1 == 1 {
            result = (result * base) % m;
        }
        exp >>= 1;
        base = (base * base) % m;
    }
    result
}

/// Precomputed zetas: zetas[k] = zeta^{BitRev_8(k)} mod q for k = 0..255.
const fn compute_zetas() -> [i32; N] {
    let mut table = [0i32; N];
    let mut k = 0usize;
    while k < N {
        let rev = bitrev8(k as u8) as u64;
        let val = pow_mod(1753, rev, Q as i64);
        table[k] = val as i32;
        k += 1;
    }
    table
}

/// Precomputed zeta table: `ZETAS[k] = zeta^{BitRev_8(k)} mod q` for k in 0..256.
///
/// This table is computed at compile time via `compute_zetas` (private to this module). Entry 0 is
/// always 1 (since `zeta^0 = 1`). The NTT and inverse NTT index into this
/// table sequentially during their butterfly passes.
pub const ZETAS: [i32; N] = compute_zetas();

/// Forward NTT (Algorithm 41 of FIPS 204).
///
/// Transforms a polynomial `f` from the standard domain to the NTT domain
/// in place. Input coefficients should be in [0, q-1]; output coefficients
/// are also in [0, q-1].
///
/// After this call, `f` represents the evaluation of the original polynomial
/// at the 256 roots of unity used by ML-DSA.
pub fn ntt(f: &mut [i32; N]) {
    let mut m = 0usize;
    let mut len = 128;
    while len >= 1 {
        let mut start = 0;
        while start < N {
            m += 1;
            let zeta = ZETAS_MONT[m];
            let mut j = start;
            while j < start + len {
                // Montgomery butterfly: mont_mul(zeta_R, f) = zeta*f (R cancels)
                let t = mont_mul(zeta, f[j + len]);
                f[j + len] = f[j] - t;
                f[j] = f[j] + t;
                j += 1;
            }
            start += 2 * len;
        }
        len /= 2;
    }
    // Reduce to [0, q-1]
    for c in f.iter_mut() {
        *c = mod_q(*c);
    }
}

/// Inverse NTT (Algorithm 42 of FIPS 204).
///
/// Transforms a polynomial `f` from the NTT domain back to the standard
/// domain in place, including the final scaling by N^{-1} mod q.
///
/// If `ntt(f)` was called first, then `ntt_inv(f)` recovers the original
/// polynomial exactly.
/// iNTT scaling factor: R² · 256⁻¹ mod q = 41978.
/// Compensates both the 256⁻¹ normalization and the /R from pointwise_mul.
const F_SCALE_DSA: i32 = 41978;

pub fn ntt_inv(f: &mut [i32; N]) {
    let mut m = N; // 256
    let mut len = 1;
    while len <= 128 {
        let mut start = 0;
        while start < N {
            m -= 1;
            let neg_zeta = mod_q(-ZETAS_MONT[m]);
            let mut j = start;
            while j < start + len {
                let t = f[j];
                f[j] = t + f[j + len];
                f[j + len] = montgomery_reduce(neg_zeta as i64 * (t - f[j + len]) as i64);
                j += 1;
            }
            start += 2 * len;
        }
        len *= 2;
    }
    for coeff in f.iter_mut() {
        *coeff = montgomery_reduce(F_SCALE_DSA as i64 * *coeff as i64);
    }
}

/// Pointwise multiplication of two NTT-domain polynomials.
///
/// Implements Algorithm 45 of FIPS 204. Because the ML-DSA NTT decomposes
/// all the way down to length-1 components, this is a simple element-wise
/// modular multiplication (no base-case Karatsuba needed).
/// Pointwise multiplication — full Montgomery.
/// Output is in /R domain. Use iNTT(F_SCALE_DSA) to compensate.
/// For accumulation in NTT domain (KeyGen), call to_mont_poly after.
pub fn pointwise_mul(a: &[i32; N], b: &[i32; N]) -> [i32; N] {
    let mut c = [0i32; N];
    for i in 0..N {
        c[i] = mont_mul(a[i], b[i]);
    }
    c
}

/// Convert polynomial from /R to normal domain (multiply each coeff by R).
pub fn to_mont_poly(f: &mut [i32; N]) {
    for c in f.iter_mut() {
        *c = montgomery_reduce(*c as i64 * R_SQ_MOD_Q);
    }
}

/// Add two polynomials coefficient-wise, reducing each result modulo q.
///
/// Returns a new polynomial `c` where `c[i] = (a[i] + b[i]) mod q`.
pub fn poly_add(a: &[i32; N], b: &[i32; N]) -> [i32; N] {
    let mut c = [0i32; N];
    for i in 0..N {
        c[i] = mod_q(a[i] + b[i]);
    }
    c
}

/// Subtract two polynomials coefficient-wise, reducing each result modulo q.
///
/// Returns a new polynomial `c` where `c[i] = (a[i] - b[i]) mod q`.
pub fn poly_sub(a: &[i32; N], b: &[i32; N]) -> [i32; N] {
    let mut c = [0i32; N];
    for i in 0..N {
        c[i] = mod_q(a[i] - b[i]);
    }
    c
}

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

    #[test]
    fn test_zetas_0() {
        // zetas[0] = zeta^{BitRev_8(0)} = zeta^0 = 1
        assert_eq!(ZETAS[0], 1);
    }

    #[test]
    fn test_ntt_roundtrip() {
        // Full pipeline: NTT → pointwise_mul (identity) → iNTT
        let mut f = [0i32; N];
        for i in 0..N {
            f[i] = (i as i32 * 17 + 3) % Q;
        }
        let orig = f;
        ntt(&mut f);
        // Multiply by NTT(1) = [1, 0, 0, ..., 0] in NTT domain
        let mut one = [0i32; N];
        one[0] = 1;
        ntt(&mut one);
        let h = pointwise_mul(&one, &f);
        let mut result = h;
        ntt_inv(&mut result);
        for i in 0..N {
            let r = mod_q(result[i]);
            assert_eq!(r, orig[i], "mismatch at index {}: got {} expected {}", i, r, orig[i]);
        }
    }

    #[test]
    fn test_mod_q_negative() {
        assert_eq!(mod_q(-1), Q - 1);
        assert_eq!(mod_q(0), 0);
        assert_eq!(mod_q(Q), 0);
    }
}