kopis/arithmetic/

ntt.rs

//! Multiplication in ℤ[X]/(X^256 + 1) via a negacyclic NTT over an auxiliary prime field.
//!
//! Every ring multiplication in Kopis has one operand with small coefficients: a CBD secret
//! with coefficients in [-μ/2, μ/2]. The other operand has coefficients in [0, 2^13). The
//! coefficients of the *exact integer* product of such polynomials — even accumulated over an
//! ℓ-term matrix-vector product — are bounded in magnitude by
//!     ℓ · 256 · (2^13 - 1) · (μ/2) ≤ 3 · 256 · 8191 · 4 = 25_162_752,
//! (the maximum over all three parameter sets). This is less than p/2 for the prime
//! p = 50330113. So we can compute the product exactly: do the arithmetic mod p, lift the
//! result to the centered representative in (-p/2, p/2), and reduce mod 2^16. The result is
//! bit-identical to schoolbook multiplication of the wrapping-u16 ring elements.
//!
//! p ≡ 1 (mod 512), so ℤ/p has a primitive 512th root of unity ψ and X^256 + 1 splits
//! completely: a full 8-level negacyclic NTT applies, and products of transformed elements are
//! plain pointwise products. This makes matrix products cheap: each entry of a matrix-vector
//! product costs one pointwise multiply-accumulate instead of a full ring multiplication, and
//! the transforms themselves are shared across rows/columns. Callers cache fixed operands
//! (the matrix A, the public vector b, the secret s) in NTT form.
//!
//! The implementation follows the reference ML-DSA (Dilithium) NTT structure: an in-place
//! Cooley-Tukey forward pass and Gentleman-Sande inverse pass over a table of ψ powers in
//! bit-reversed order, with signed Montgomery reduction (R = 2^32). Unlike Dilithium's q,
//! our p is large enough that Gentleman-Sande sums could overflow an i32 after a few levels,
//! so the inverse butterfly Barrett-reduces the sum path every level. All arithmetic is
//! branch-free and all loop bounds are public, so the code is constant-time.

// The explicit `for i in 0..N` index loops that trigger this lint are deliberate: aeneas (the
// Lean extractor) handles them better than the iterator patterns clippy suggests
#![allow(clippy::needless_range_loop)]

use crate::{
    arithmetic::{Matrix, RingElem},
    consts::RING_DEG,
};

use zeroize::Zeroize;

/// The NTT prime modulus
const P: i32 = 50330113;
/// p^-1 mod 2^32 (as a wrapping u32)
const P_INV: u32 = 3575907841;
/// round(2^48 / p), for Barrett reduction
const BARRETT_M: i64 = 5592576;
/// The rounding addend for Barrett reduction, 2^47. Spelled as a literal rather than `1 << 47`
/// so that the extracted Lean contains no shift operation to discharge.
const BARRETT_ROUND: i64 = 140737488355328;
/// ⌊p/2⌋, the centering threshold. A literal for the same reason (no division to discharge).
const P_HALF: i32 = 25165056;
/// 256^-1 · 2^64 mod p: the inverse-NTT output scale. One Montgomery reduction by this value
/// undoes both the 1/256 of the inverse transform and the 2^-32 introduced by the Montgomery
/// reduction in the pointwise multiplication step.
const INVNTT_SCALE: i32 = 44652572;

/// Powers of ψ = 49118445 (a primitive 512th root of unity mod p) in bit-reversed order and
/// Montgomery form: ZETAS[k] = ψ^brv8(k) · 2^32 mod p. Generated by gen_ntt_consts.py; the
/// `zetas_table_is_correct` test recomputes the table from ψ and checks every entry.
#[rustfmt::skip]
const ZETAS: [i32; 256] = [
    16907691, 2667794, 25435945, 38041794, 4526213, 35360986, 30701667, 21078350,
    36956843, 13677930, 22323069, 13735052, 16437056, 5691837, 29800035, 15764652,
    45309494, 32237400, 24236486, 2001600, 1321324, 27016828, 15167304, 4366839,
    28120194, 32421317, 30261273, 12292806, 22968316, 13702742, 937110, 2965954,
    34335133, 44939767, 6512044, 26419553, 19973946, 49111849, 40847418, 20023056,
    43261597, 41129446, 12627709, 45952874, 30150184, 36681813, 45532529, 39191671,
    45414311, 2081056, 808565, 49693071, 31006016, 28961757, 15337390, 31683976,
    4955339, 27911199, 29647313, 30189282, 40031443, 49298885, 46092656, 18408488,
    245475, 31889855, 39534436, 41714666, 8292132, 21449251, 22309987, 8205570,
    31096891, 47405297, 930348, 6971348, 10106150, 31035883, 46359140, 35263708,
    25740623, 9884444, 542513, 16415589, 9876654, 44372807, 21881921, 19165831,
    9227072, 2984297, 36824748, 43900147, 43917176, 9030513, 26543409, 34836022,
    45307234, 35198677, 14645063, 2098347, 37061938, 3244828, 41096247, 32062813,
    8918713, 47344810, 46288223, 21000897, 41046191, 20508021, 35242638, 21557502,
    49459492, 3469647, 1704205, 29274982, 27408601, 6834376, 3860969, 50170866,
    7253270, 26028155, 34019372, 9518412, 22713734, 14612739, 26018446, 43305093,
    12047271, 21217146, 26069968, 4256850, 3639874, 48691800, 17320569, 3323937,
    32443671, 42492417, 9308803, 33498296, 25236161, 17970048, 30529967, 11776789,
    23284408, 23725761, 39814592, 32816444, 45216211, 34497791, 1878200, 45702238,
    40014922, 39122569, 28779735, 44971151, 12142149, 25663462, 31479913, 19837380,
    39143843, 21323231, 17806070, 36807664, 47595478, 49950288, 16588490, 32773564,
    19315478, 33753056, 15093953, 23446825, 49044025, 41815538, 1288725, 23610193,
    2754751, 40030005, 13643238, 20793088, 16114388, 33147218, 15927487, 41781494,
    5795109, 11439966, 15064762, 35884894, 22645018, 10582966, 14900694, 21660678,
    16765530, 20825350, 30320649, 42368153, 39131901, 42360159, 6363984, 35581825,
    28072132, 11706419, 20970510, 44236552, 17934700, 14577479, 42770390, 44206858,
    28865919, 40187227, 15964209, 15444696, 17019703, 32900174, 28423607, 42332883,
    32435385, 40121202, 37968791, 33541227, 35393185, 34591081, 48852822, 8473545,
    43827986, 23231608, 34405552, 26405625, 27381101, 33514230, 46158701, 27741938,
    10409285, 21224207, 6796942, 35314909, 49625944, 20271219, 6354214, 38361969,
    34767461, 14097694, 13212224, 47820051, 49316243, 20113174, 21747171, 39170867,
    38847587, 20584729, 27124165, 40492847, 7742348, 43534070, 7422899, 12232461,
];

// Every arithmetic operation below is written with an explicit `wrapping_*` / `wrapping_shr`
// even where the value bounds guarantee no overflow. This is deliberate: aeneas extracts
// checked `+`/`-`/`*`/`>>` into the `Result` monad, so each one becomes a separate
// panic-freedom obligation whose discharge needs the full magnitude analysis. The wrapping
// forms extract as total functions, which keeps the generated Lean in plain (non-monadic)
// arithmetic and leaves exactly one thing to prove: that the mathematical values are in range,
// so the wrapping never actually wraps. That argument is made once, in the correctness proof,
// instead of 36 times inline. Runtime behaviour is unchanged (the KATs pin it down).

/// Signed Montgomery reduction: for |a| < 2^31 · p, returns t ≡ a · 2^-32 (mod p) with |t| < p
#[inline(always)]
fn mont_reduce(a: i64) -> i32 {
    let t = (a as u32).wrapping_mul(P_INV) as i32 as i64;
    a.wrapping_sub(t.wrapping_mul(P as i64)).wrapping_shr(32) as i32
}

/// Centered Barrett reduction: for any i32 input, returns r ≡ x (mod p) with |r| ≤ p/2 + 1
#[inline(always)]
fn barrett_reduce(x: i32) -> i32 {
    let q = (x as i64)
        .wrapping_mul(BARRETT_M)
        .wrapping_add(BARRETT_ROUND)
        .wrapping_shr(48) as i32;
    x.wrapping_sub(q.wrapping_mul(P))
}

/// Reduces a value in (-p, p) to the canonical range [0, p), branch-free
#[inline(always)]
fn to_canonical(x: i32) -> i32 {
    x.wrapping_add(x.wrapping_shr(31) & P)
}

/// Lifts a value in (-p, p) to its centered representative in (-p/2, p/2] and reduces it
/// mod 2^16, branch-free. This is exact whenever the true integer value it represents lies in
/// (-p/2, p/2], which the coefficient bound in the module docs guarantees.
#[inline(always)]
fn to_wrapping_u16(x: i32) -> u16 {
    let x = to_canonical(x); // [0, p)
    let x = x.wrapping_sub(P & P_HALF.wrapping_sub(x).wrapping_shr(31)); // (-p/2, p/2]
    x as u16
}

/// In-place forward negacyclic NTT (Cooley-Tukey), then centering to |a[i]| ≤ p/2 + 1.
///
/// Input coefficients must satisfy |a[i]| < 2^13. Butterflies add at most p in magnitude per
/// level, so intermediate values stay below 8p + 2^13 < 2^29 and never overflow. The output
/// is the evaluation of the input polynomial at the odd powers of ψ, in bit-reversed order.
fn ntt(a: &mut [i32; RING_DEG]) {
    let mut k = 0;
    let mut len = RING_DEG / 2;
    while len > 0 {
        let mut start = 0;
        while start < RING_DEG {
            k += 1;
            let zeta = ZETAS[k] as i64;
            for j in start..start + len {
                let t = mont_reduce(zeta.wrapping_mul(a[j + len] as i64));
                a[j + len] = a[j].wrapping_sub(t);
                a[j] = a[j].wrapping_add(t);
            }
            start += 2 * len;
        }
        len /= 2;
    }

    for coeff in a.iter_mut() {
        *coeff = barrett_reduce(*coeff);
    }
}

/// In-place inverse negacyclic NTT (Gentleman-Sande). Input coefficients must satisfy
/// |a[i]| < p. The final Montgomery multiplication by INVNTT_SCALE leaves plain
/// (non-Montgomery) centered values with |a[i]| < p, *assuming* the pointwise step that
/// produced the input performed exactly one Montgomery reduction (see module docs).
///
/// p is too large for Dilithium-style lazy growth to fit in an i32 across all 8 levels: the
/// un-reduced sum path doubles per level. Starting below p, four levels reach 16p < 2^31;
/// one Barrett pass then re-centers everything to ≤ p/2 + 1, and the remaining four levels
/// reach at most 16(p/2 + 1) < 2^31 again. The Montgomery inputs also stay in range: the
/// largest is ζ·(t - a) with |ζ| < p and |t - a| < 16p, and 16p² < 2^31 · p.
fn invntt(a: &mut [i32; RING_DEG]) {
    let mut k = RING_DEG;
    let mut len = 1;
    while len < RING_DEG {
        let mut start = 0;
        while start < RING_DEG {
            k -= 1;
            let neg_zeta = (ZETAS[k] as i64).wrapping_neg();
            for j in start..start + len {
                let t = a[j];
                a[j] = t.wrapping_add(a[j + len]);
                a[j + len] = mont_reduce(neg_zeta.wrapping_mul(t.wrapping_sub(a[j + len]) as i64));
            }
            start += 2 * len;
        }
        len *= 2;

        // Halfway (after the len=8 level), re-center to keep sums within i32
        if len == 16 {
            for coeff in a.iter_mut() {
                *coeff = barrett_reduce(*coeff);
            }
        }
    }

    for coeff in a.iter_mut() {
        *coeff = mont_reduce((*coeff as i64).wrapping_mul(INVNTT_SCALE as i64));
    }
}

/// A ring element in the NTT domain. Coefficients are centered mod-p values, |·| ≤ p/2 + 1.
// The NTT is an invertible linear map, so a transformed secret is exactly as sensitive as the
// coefficient-domain one. Zeroize accordingly, matching RingElem.
#[derive(Clone, Copy, Zeroize)]
pub(crate) struct NttElem(pub(crate) [i32; RING_DEG]);

impl Default for NttElem {
    fn default() -> Self {
        NttElem([0i32; RING_DEG])
    }
}

impl NttElem {
    /// Forward-transforms a ring element whose coefficients are plain values in [0, 2^13),
    /// e.g. an element of the uniform matrix A (13 bits) or a rounded vector (10 bits).
    ///
    /// The `< 2^13` bound on every coefficient is a precondition, not a checked one: it is
    /// what keeps the forward transform's intermediate values inside an i32 and the final
    /// product inside the exactness bound. It is discharged at each call site by the Lean
    /// correspondence proof rather than by a per-coefficient runtime check.
    pub(crate) fn from_uniform(elem: &RingElem) -> Self {
        let mut a = [0i32; RING_DEG];
        for i in 0..RING_DEG {
            a[i] = elem.0[i] as i32;
        }
        ntt(&mut a);
        NttElem(a)
    }

    /// Forward-transforms a CBD secret, interpreting each wrapping-u16 coefficient as the
    /// signed value it represents (e.g. 0xFFFB is -5).
    ///
    /// As with [`NttElem::from_uniform`], the bound on the coefficients — here |·| ≤ μ/2 — is
    /// a proof-side precondition rather than a runtime check.
    pub(crate) fn from_secret(elem: &RingElem) -> Self {
        let mut a = [0i32; RING_DEG];
        for i in 0..RING_DEG {
            a[i] = elem.0[i] as i16 as i32;
        }
        ntt(&mut a);
        NttElem(a)
    }
}

/// A matrix of NTT-domain ring elements, stored in row-major order like [`Matrix`]
#[derive(Clone, Zeroize)]
pub(crate) struct NttMatrix<const X: usize, const Y: usize>(pub(crate) [[NttElem; Y]; X]);

impl<const X: usize, const Y: usize> Default for NttMatrix<X, Y> {
    fn default() -> Self {
        NttMatrix([[NttElem::default(); Y]; X])
    }
}

/// Adds the pointwise product lhs ∘ rhs into the i64 accumulator, without reducing.
///
/// Products of centered values are below (p/2 + 1)² and callers accumulate at most 4 (= MAX_L)
/// of them, so the accumulator stays (just) below p² < 2^52, within Montgomery reduction's
/// valid input range.
fn pointwise_mul_acc(acc: &mut [i64; RING_DEG], lhs: &NttElem, rhs: &NttElem) {
    for i in 0..RING_DEG {
        acc[i] = acc[i].wrapping_add((lhs.0[i] as i64).wrapping_mul(rhs.0[i] as i64));
    }
}

/// Montgomery-reduces an accumulator of pointwise products, inverse-transforms it, and returns
/// the coefficient-domain result as a wrapping-u16 ring element.
fn reduce_invntt_to_ring_elem(acc: &[i64; RING_DEG]) -> RingElem {
    let mut v = [0i32; RING_DEG];
    for i in 0..RING_DEG {
        v[i] = mont_reduce(acc[i]);
    }
    invntt(&mut v);

    let mut out = RingElem::default();
    for i in 0..RING_DEG {
        out.0[i] = to_wrapping_u16(v[i]);
    }
    out
}

impl<const X: usize, const Y: usize> NttMatrix<X, Y> {
    /// Forward-transforms a matrix of uniform (≤ 13-bit) ring elements, entry by entry
    pub(crate) fn from_uniform_matrix(mat: &Matrix<X, Y>) -> Self {
        let mut ret = NttMatrix::default();
        for i in 0..X {
            for j in 0..Y {
                ret.0[i][j] = NttElem::from_uniform(&mat.0[i][j]);
            }
        }
        ret
    }

    /// Forward-transforms a matrix of CBD secrets, entry by entry
    pub(crate) fn from_secret_matrix(mat: &Matrix<X, Y>) -> Self {
        let mut ret = NttMatrix::default();
        for i in 0..X {
            for j in 0..Y {
                ret.0[i][j] = NttElem::from_secret(&mat.0[i][j]);
            }
        }
        ret
    }

    /// Multiplies two NTT-domain matrices, returning the result in the coefficient domain.
    /// Equivalent to [`Matrix::mul`] on the corresponding coefficient-domain matrices.
    pub(crate) fn mul<const Z: usize>(&self, other: &NttMatrix<Y, Z>) -> Matrix<X, Z> {
        // The inner dimension bounds the pointwise accumulator (see pointwise_mul_acc)
        debug_assert!(Y <= crate::consts::MAX_L);

        let mut result = Matrix::default();
        for i in 0..X {
            for k in 0..Z {
                let mut acc = [0i64; RING_DEG];
                for j in 0..Y {
                    pointwise_mul_acc(&mut acc, &self.0[i][j], &other.0[j][k]);
                }
                result.0[i][k] = reduce_invntt_to_ring_elem(&acc);
            }
        }
        result
    }

    /// Multiplies the transpose of this NTT-domain matrix by another, returning the result in
    /// the coefficient domain. Equivalent to [`Matrix::mul_transpose`].
    pub(crate) fn mul_transpose<const Z: usize>(&self, other: &NttMatrix<X, Z>) -> Matrix<Y, Z> {
        // The inner dimension bounds the pointwise accumulator (see pointwise_mul_acc)
        debug_assert!(X <= crate::consts::MAX_L);

        let mut result = Matrix::default();
        for j in 0..Y {
            for k in 0..Z {
                let mut acc = [0i64; RING_DEG];
                for i in 0..X {
                    pointwise_mul_acc(&mut acc, &self.0[i][j], &other.0[i][k]);
                }
                result.0[j][k] = reduce_invntt_to_ring_elem(&acc);
            }
        }
        result
    }
}

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

    use rand::{Rng, rng};

    /// ψ, the primitive 512th root of unity underlying ZETAS
    const PSI: u64 = 49118445;

    fn pow_mod(mut base: u64, mut exp: u64, modulus: u64) -> u64 {
        let mut acc = 1u64;
        base %= modulus;
        while exp > 0 {
            if exp & 1 == 1 {
                acc = acc * base % modulus;
            }
            base = base * base % modulus;
            exp >>= 1;
        }
        acc
    }

    // The constants spelled as literals (to keep shifts and divisions out of the extracted
    // Lean) must equal the expressions they stand for
    #[test]
    fn literal_constants_are_correct() {
        assert_eq!(BARRETT_ROUND, 1i64 << 47);
        assert_eq!(P_HALF, P / 2);
        assert_eq!(BARRETT_M, ((1i64 << 48) + (P as i64) / 2) / (P as i64));
        // P_INV is p^-1 mod 2^32, i.e. p · P_INV ≡ 1
        assert_eq!((P as u32).wrapping_mul(P_INV), 1);
    }

    // Recompute the ZETAS table from ψ and check every entry, plus ψ's defining properties
    #[test]
    fn zetas_table_is_correct() {
        let p = P as u64;
        // ψ is a 512th root of unity and ψ^256 = -1 (so the transform is negacyclic)
        assert_eq!(pow_mod(PSI, 256, p), p - 1);
        assert_eq!(pow_mod(PSI, 512, p), 1);

        for (k, &z) in ZETAS.iter().enumerate() {
            let brv = (k as u8).reverse_bits() as u64;
            let expected = pow_mod(PSI, brv, p) * (1u64 << 32) % p;
            assert_eq!(z as u64, expected, "ZETAS[{k}]");
        }
    }

    /// A random CBD-like secret: wrapping-u16 coefficients in [-mu/2, mu/2]
    fn rand_secret(rng: &mut impl Rng, half_mu: u16) -> RingElem {
        let mut ret = RingElem::default();
        for coeff in ret.0.iter_mut() {
            *coeff = rng.random_range(0..=2 * half_mu).wrapping_sub(half_mu);
        }
        ret
    }

    /// A random uniform ring element with `bits`-bit coefficients
    fn rand_uniform(rng: &mut impl Rng, bits: usize) -> RingElem {
        let mut ret = RingElem::default();
        for coeff in ret.0.iter_mut() {
            *coeff = rng.random_range(0..(1u16 << bits));
        }
        ret
    }

    // Multiplying by the polynomial "1" via the NTT pipeline must return the input unchanged
    #[test]
    fn mul_by_one_roundtrip() {
        let mut rng = rng();
        for _ in 0..50 {
            let a = rand_uniform(&mut rng, 13);
            let mut one = RingElem::default();
            one.0[0] = 1;

            let a_ntt = NttMatrix::<1, 1>([[NttElem::from_uniform(&a)]]);
            let one_ntt = NttMatrix::<1, 1>([[NttElem::from_secret(&one)]]);
            let prod = a_ntt.mul(&one_ntt);
            assert_eq!(prod.0[0][0], a);
        }
    }

    // NTT-domain matrix products must match the schoolbook Matrix products exactly, for the
    // shapes and operand ranges of all three parameter sets
    #[test]
    fn matches_schoolbook() {
        fn check<const L: usize>(half_mu: u16) {
            let mut rng = rng();
            for _ in 0..20 {
                // A is L×L uniform 13-bit, b is L×1 uniform 10-bit, s is an L×1 secret
                let mut mat_a = Matrix::<L, L>::default();
                for row in mat_a.0.iter_mut() {
                    for entry in row.iter_mut() {
                        *entry = rand_uniform(&mut rng, 13);
                    }
                }
                let mut vec_b = Matrix::<L, 1>::default();
                let mut vec_s = Matrix::<L, 1>::default();
                for i in 0..L {
                    vec_b.0[i][0] = rand_uniform(&mut rng, 10);
                    vec_s.0[i][0] = rand_secret(&mut rng, half_mu);
                }

                let a_ntt = NttMatrix::from_uniform_matrix(&mat_a);
                let b_ntt = NttMatrix::from_uniform_matrix(&vec_b);
                let s_ntt = NttMatrix::from_secret_matrix(&vec_s);

                assert_eq!(a_ntt.mul(&s_ntt), mat_a.mul(&vec_s));
                assert_eq!(a_ntt.mul_transpose(&s_ntt), mat_a.mul_transpose(&vec_s));
                assert_eq!(b_ntt.mul_transpose(&s_ntt), vec_b.mul_transpose(&vec_s));
            }
        }

        check::<2>(5); // kopis512: L=2, mu=10
        check::<3>(4); // kopis768: L=3, mu=8
        check::<4>(3); // kopis1024: L=4, mu=6
    }

    // The worst-case accumulated product coefficient over all parameter sets is exactly
    // L·256·8191·(μ/2) = 25_162_752 (L=3, μ=8), just under p/2. Hit it exactly: with
    // a_i = 8191 for all i, s_0 = μ/2, and s_j = -μ/2 for j > 0, coefficient 0 of a·s is
    // Σ_i 8191·(μ/2) = 256·8191·(μ/2), and the matrix product accumulates it L times.
    #[test]
    fn extremal_coefficients() {
        const L: usize = 3;
        const HALF_MU: u16 = 4;

        let mut mat_a = Matrix::<L, L>::default();
        for row in mat_a.0.iter_mut() {
            for entry in row.iter_mut() {
                *entry = RingElem([8191u16; RING_DEG]);
            }
        }
        let mut vec_s = Matrix::<L, 1>::default();
        for i in 0..L {
            let mut s = RingElem([0u16.wrapping_sub(HALF_MU); RING_DEG]);
            s.0[0] = HALF_MU;
            vec_s.0[i][0] = s;
        }

        let a_ntt = NttMatrix::from_uniform_matrix(&mat_a);
        let s_ntt = NttMatrix::from_secret_matrix(&vec_s);

        assert_eq!(a_ntt.mul(&s_ntt), mat_a.mul(&vec_s));
        assert_eq!(a_ntt.mul_transpose(&s_ntt), mat_a.mul_transpose(&vec_s));
    }
}