quantica/ml_dsa/

dsa.rs

//! Core ML-DSA algorithms (FIPS 204, Algorithms 1-8).
//!
//! Contains the key generation, signing, and verification routines at both
//! the public API level (Algorithms 1-3) and the internal/deterministic level
//! (Algorithms 6-8).
//!
//! All internal polynomial vector operations use fixed-size stack arrays
//! (`[[i32; N]; MAX_K]` / `[[i32; N]; MAX_L]`) to avoid heap allocations.
//!
//! # Side-channel countermeasures (`sca-protected` feature)
//!
//! When the `sca-protected` Cargo feature is enabled (on by default),
//! `sign_internal` runs an additional layer of defences on the secret-key
//! material:
//!
//! | Countermeasure        | Module                       | Threat addressed                              |
//! |-----------------------|------------------------------|-----------------------------------------------|
//! | Constant-time arith   | always-on                    | Cache- / branch-based timing attacks          |
//! | Zeroization           | always-on                    | Cold-boot dumps, use-after-free               |
//! | Hedged signing        | always-on                    | Fault-induced nonce reuse (`rnd ≠ 0`)         |
//! | Shuffled NTT          | `super::shuffle` (sca)     | SPA, trace-alignment for DPA                  |
//! | First-order masking   | `super::masked` (sca)      | First-order DPA, template attacks             |
//! | Mask refresh / hop    | `super::masked` (sca)      | Inter-iteration share correlation             |
//!
//! The masking + shuffling layer is deliberately confined to `sign_internal`,
//! because that is where the secret key `(s1, s2, t0)` is consumed in
//! polynomial multiplications with values an attacker can influence:
//!
//! ```text
//!   ŝ1, ŝ2, t̂0  ←  NTT(s1), NTT(s2), NTT(t0)               // SPA + DPA target
//!   loop:
//!       ĉ ← NTT(SampleInBall(c̃))
//!       cs1[i]  ← ĉ · ŝ1[i]                                 // ×L  — DPA target
//!       cs2[i]  ← ĉ · ŝ2[i]                                 // ×K  — DPA target
//!       ct0[i]  ← ĉ · t̂0[i]                                 // ×K  — DPA target
//! ```
//!
//! The challenge polynomial `ĉ` is **public** (the verifier recomputes it),
//! so every secret×public multiplication only needs first-order masking:
//! `(s₀ + s₁) · ĉ = s₀·ĉ + s₁·ĉ`. There is no secret×secret operation in
//! Sign that would require second-order shares.
//!
//! Mask randomness is drawn from a SHAKE256-based deterministic
//! `ScaRng` seeded with `(K ‖ rnd ‖ tr ‖ M')`, so that:
//!
//! * `sign_internal` keeps a deterministic signature (no `&mut dyn CryptoRng`
//!   parameter), and the NIST ACVP fixed-`rnd = 0` test vectors still match
//!   bit-for-bit;
//! * different `rnd` values produce independent share streams (hedged
//!   signing entropy is preserved through the SCA layer).
//!
//! The standard build (without `sca-protected`) still benefits from the
//! always-on countermeasures listed above; only the masking + shuffling
//! defences are conditionally compiled out.

use super::MlDsaError;
use super::decompose;
use super::encode;
use super::ntt::{self, mod_q};
use super::params::{D, MAX_K, MAX_L, N, Params, Q};
use super::rng::CryptoRng;
use super::sample;
use super::sha3;
use alloc::vec::Vec;

#[cfg(any(feature = "compressed-poly", feature = "compressed-challenge"))]
use super::compressed;
#[cfg(feature = "sca-protected")]
use super::masked::{self, MaskedPoly};

#[cfg(all(feature = "sca-protected", feature = "compressed-challenge"))]
compile_error!(
    "features `sca-protected` and `compressed-challenge` are mutually exclusive: masking requires NTT-domain multiplication, schoolbook operates in time domain"
);

#[cfg(all(feature = "sca-protected", feature = "small-secret"))]
compile_error!("features `sca-protected` and `small-secret` are mutually exclusive: masking operates in i32 domain");

#[cfg(all(feature = "sca-protected", feature = "union-buffer"))]
compile_error!("features `sca-protected` and `union-buffer` are mutually exclusive");

#[cfg(feature = "sca-protected")]
use super::sha3::KeccakState;
#[cfg(feature = "sca-protected")]
use super::shuffle;
#[cfg(feature = "small-secret")]
use super::smallpoly::{self, SmallPoly};

/// Deterministic SHAKE256-based randomness source for the SCA layer.
///
/// `sign_internal` does not take a `&mut dyn CryptoRng` parameter
/// (it must stay fully deterministic so that the NIST ACVP fixed-`rnd`
/// vectors still match bit-for-bit), so the masking and shuffling
/// modules cannot reach for [`super::OsRng`] either. Instead they
/// share a per-call `ScaRng` whose seed is derived from
/// `(K ‖ rnd ‖ tr ‖ M')` via SHAKE256:
///
/// * `K` is the secret-key field used by FIPS 204 hedged signing.
/// * `rnd` is the 32-byte hedged-signing randomness — all-zero in
///   deterministic / ACVP test mode, fresh entropy otherwise.
/// * `tr` and `M'` make the seed bind to the public key + message,
///   so two signatures over different messages produce uncorrelated
///   share streams even when `rnd = 0`.
///
/// A short domain-separation tag (`b"quantica-mldsa-sca-v1"`) is
/// absorbed first to keep the SHAKE squeeze stream disjoint from
/// any other SHAKE use elsewhere in the algorithm.
///
/// All ML-DSA share / shuffle randomness for one signature flows
/// from one `ScaRng` instance: the initial mask of `(s1, s2, t0)`,
/// the shuffled-NTT permutations, and the per-rejection-iteration
/// `MaskedPoly::refresh()` calls. This guarantees a single coherent
/// stream that the test vectors can reproduce.
///
/// The PRG itself is **not cryptographic in the standard sense** —
/// it is not seeded from system entropy. Its job is purely to make
/// internal mask shares unpredictable to a passive side-channel
/// observer. The actual signature security still derives from FIPS
/// 204's own randomness (`rnd`), which is mixed into the seed.
#[cfg(feature = "sca-protected")]
struct ScaRng {
    state: KeccakState,
}

#[cfg(feature = "sca-protected")]
impl ScaRng {
    /// Initialize a fresh SCA RNG state from a caller-supplied seed.
    /// The domain-separation tag is absorbed first so the squeeze
    /// stream is disjoint from any other SHAKE256 usage.
    fn from_seed(seed: &[u8]) -> Self {
        let mut s = sha3::shake256();
        s.absorb(b"quantica-mldsa-sca-v1");
        s.absorb(seed);
        Self { state: s }
    }
}

#[cfg(feature = "sca-protected")]
impl CryptoRng for ScaRng {
    fn fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), MlDsaError> {
        self.state.squeeze(dest);
        Ok(())
    }
}

/// Compute the matrix-vector product A_hat * s in the NTT domain.
///
/// `a_hat` is a k-by-l matrix of NTT-domain polynomials and `s` is a vector
/// of l NTT-domain polynomials. The result is a vector of k polynomials.
///
/// Used when `low-mem` is **disabled** (default). The full matrix is
/// pre-expanded in RAM for maximum throughput.
#[cfg(not(feature = "low-mem"))]
fn mat_vec_mul(a_hat: &[[[i32; N]; MAX_L]; MAX_K], s: &[[i32; N]], k: usize, l: usize, result: &mut [[i32; N]]) {
    for i in 0..k {
        result[i] = [0i32; N];
        for j in 0..l {
            let prod = ntt::pointwise_mul(&a_hat[i][j], &s[j]);
            result[i] = ntt::poly_add(&result[i], &prod);
        }
    }
}

/// Low-memory matrix-vector product: recomputes each `a_hat[i][j]`
/// polynomial on-the-fly from `rho` via SHAKE128, instead of holding
/// the full k×l matrix (57 KB for ML-DSA-87) on the stack.
///
/// Trade-off: saves **57 KB of stack** at the cost of re-running
/// SHAKE128 for each polynomial element each time this function is
/// called. In `sign_internal`, the rejection loop calls this once per
/// iteration, so the SHAKE overhead is multiplied by the average
/// rejection count (~4 for ML-DSA-65).
///
/// Used when the `low-mem` feature is **enabled**.
#[cfg(feature = "low-mem")]
fn mat_vec_mul_lazy(rho: &[u8; 32], s: &[[i32; N]], k: usize, l: usize, result: &mut [[i32; N]]) {
    for i in 0..k {
        result[i] = [0i32; N];
        for j in 0..l {
            // Recompute a_hat[i][j] from rho (same as expand_a).
            let a_ij = sample::rej_ntt_poly(rho, j as u8, i as u8);
            let prod = ntt::pointwise_mul(&a_ij, &s[j]);
            result[i] = ntt::poly_add(&result[i], &prod);
        }
    }
}

/// NTT of a polynomial vector (in-place, first `len` elements).
fn ntt_vec(v: &mut [[i32; N]], len: usize) {
    for poly in v[..len].iter_mut() {
        ntt::ntt(poly);
    }
}

/// Inverse NTT of a polynomial vector (in-place, first `len` elements).
fn ntt_inv_vec(v: &mut [[i32; N]], len: usize) {
    for poly in v[..len].iter_mut() {
        ntt::ntt_inv(poly);
    }
}

/// Add two polynomial vectors into `out`. Only processes `len` elements.
fn vec_add(a: &[[i32; N]], b: &[[i32; N]], out: &mut [[i32; N]], len: usize) {
    for i in 0..len {
        out[i] = ntt::poly_add(&a[i], &b[i]);
    }
}

/// Subtract two polynomial vectors into `out`. Only processes `len` elements.
fn vec_sub(a: &[[i32; N]], b: &[[i32; N]], out: &mut [[i32; N]], len: usize) {
    for i in 0..len {
        out[i] = ntt::poly_sub(&a[i], &b[i]);
    }
}

// =====================================================================
// low-stack helpers: heap-allocated polynomial vectors
// =====================================================================
//
// When `low-stack` is enabled, the rejection-loop temporaries in
// sign_internal are allocated on the heap (Vec) with scoped lifetimes
// and explicit drop() calls to keep the high-water mark low (~23 KB).
// The downstream functions (vec_add, check_norm_vec, decompose::*)
// take &[[i32; N]] slices and work unchanged with either stack or heap.

/// Allocate a zero-initialized polynomial vector of `len` polynomials.
#[cfg(feature = "low-stack")]
fn poly_vec(len: usize) -> Vec<[i32; N]> {
    vec![[0i32; N]; len]
}

/// Check infinity norm of polynomial: all coefficients strictly below `bound`.
///
/// Returns `true` iff every coefficient `c` satisfies `|c| < bound`
/// (i.e., `c ∈ (−bound, bound)`). This implements the **strict**
/// inequality `||v||∞ < bound` required by FIPS 204 Algorithm 7
/// step 25 / Algorithm 8 step 15: "if `||z||∞ ≥ γ₁ − β` then
/// return ⊥".
fn check_norm(v: &[i32; N], bound: i32) -> bool {
    for &c in v.iter() {
        // Bring to centered representation
        let mut val = mod_q(c);
        if val > Q / 2 {
            val -= Q;
        }
        // Strict: reject if |val| >= bound
        if val >= bound || val <= -bound {
            return false;
        }
    }
    true
}

/// Check infinity norm of polynomial vector. Only checks first `len` elements.
fn check_norm_vec(v: &[[i32; N]], bound: i32, len: usize) -> bool {
    for poly in v[..len].iter() {
        if !check_norm(poly, bound) {
            return false;
        }
    }
    true
}

/// Deterministic key generation from a 32-byte seed.
///
/// Implements Algorithm 6 of FIPS 204 (ML-DSA.KeyGen_internal).
///
/// Given a 32-byte seed `xi`, derives the public matrix A (via ExpandA),
/// secret vectors s1 and s2 (via ExpandS), and computes the public key
/// `pk = (rho, t1)` and secret key `sk = (rho, K, tr, s1, s2, t0)`.
///
/// - `xi`: 32-byte random seed.
///
/// Returns `(pk, sk)` as byte vectors.
pub fn keygen_internal<P: Params>(xi: &[u8; 32]) -> (Vec<u8>, Vec<u8>) {
    let k = P::K;
    let l = P::L;

    // (rho, rho', K) = H(xi || k || l)
    let mut h_input = [0u8; 34];
    h_input[..32].copy_from_slice(xi);
    h_input[32] = k as u8;
    h_input[33] = l as u8;

    let mut hash_out = [0u8; 128]; // need 32 + 64 + 32 = 128 bytes
    let mut state = sha3::shake256();
    state.absorb(&h_input);
    state.squeeze(&mut hash_out);

    let mut rho = [0u8; 32];
    rho.copy_from_slice(&hash_out[..32]);
    let mut rho_prime = [0u8; 64];
    rho_prime.copy_from_slice(&hash_out[32..96]);
    let mut k_seed = [0u8; 32];
    k_seed.copy_from_slice(&hash_out[96..128]);

    // Generate s1, s2 from rho'
    let (mut s1, mut s2) = sample::expand_s::<P>(&rho_prime);

    // s1_hat = NTT(s1)
    let mut s1_hat = s1;
    ntt_vec(&mut s1_hat, l);

    // t = NTT^{-1}(A-hat * s1_hat) + s2
    let mut t = [[0i32; N]; MAX_K];
    #[cfg(not(feature = "low-mem"))]
    {
        let a_hat = sample::expand_a::<P>(&rho);
        mat_vec_mul(&a_hat, &s1_hat, k, l, &mut t);
    }
    #[cfg(feature = "low-mem")]
    mat_vec_mul_lazy(&rho, &s1_hat, k, l, &mut t);
    ntt_inv_vec(&mut t, k);
    // t = t + s2 (in-place into t)
    {
        let mut tmp = [[0i32; N]; MAX_K];
        vec_add(&t, &s2, &mut tmp, k);
        t = tmp;
    }

    // (t1, t0) = Power2Round(t)
    let (t1, t0) = encode::power2round_vec(&t, k);

    // Encode public key
    let pk = encode::pk_encode::<P>(&rho, &t1);

    // tr = H(pk)  (SHAKE256, 64 bytes)
    let mut tr = [0u8; 64];
    sha3::shake256_digest(&pk, &mut tr);

    // Encode secret key
    let sk = encode::sk_encode::<P>(&rho, &k_seed, &tr, &s1, &s2, &t0);

    // Zeroize sensitive data
    for poly in s1[..l].iter_mut() {
        for c in poly.iter_mut() {
            *c = 0;
        }
    }
    for poly in s2[..k].iter_mut() {
        for c in poly.iter_mut() {
            *c = 0;
        }
    }
    for byte in rho_prime.iter_mut() {
        *byte = 0;
    }
    for byte in k_seed.iter_mut() {
        *byte = 0;
    }

    (pk, sk)
}

/// Sign a pre-formatted message (deterministic or hedged).
///
/// Implements Algorithm 7 of FIPS 204 (ML-DSA.Sign_internal).
///
/// This function contains the core rejection sampling loop: candidate
/// signatures `(z, h)` are generated from a masking vector `y` and the
/// challenge polynomial `c`, then tested against the norm bounds
/// `||z||_inf < gamma1 - beta` and `||r0||_inf < gamma2 - beta`. If any
/// check fails, the counter `kappa` is incremented and a new attempt begins.
///
/// - `sk`: encoded secret key bytes.
/// - `m_prime`: pre-formatted message (e.g., `0x00 || len(ctx) || ctx || msg`).
/// - `rnd`: 32-byte randomness. Use random bytes for hedged signing or
///   all-zeros for fully deterministic signing.
///
/// # Side-channel countermeasures
///
/// With the `sca-protected` Cargo feature enabled (default), this
/// function activates the additional defences described in the
/// crate-level documentation:
///
/// 1. **Shuffled NTT** on `s1`, `s2`, `t0` — runs once at entry,
///    via [`super::shuffle::ntt_shuffled`]. Defends against SPA on
///    the secret-key NTT and disrupts trace alignment for any
///    later DPA campaign that tries to average aligned traces.
/// 2. **First-order additive masking** of the NTT-domain secrets,
///    via [`super::masked::MaskedPoly`]. Each polynomial is split
///    into two shares mod `q = 8 380 417`; no single intermediate
///    value reveals the secret to a first-order observer.
/// 3. **Per-iteration `c·sₓ` multiplications** go through
///    [`super::masked::masked_pointwise_mul_public`], which
///    multiplies each share independently by the public challenge
///    `ĉ`. Because `ĉ` is public, first-order shares are sufficient
///    — no secret×secret operation is performed.
/// 4. **Mask refresh after every use**: the share pair is
///    re-randomized via `MaskedPoly::refresh()` between rejection
///    iterations, so the same secret never multiplies the same
///    share twice — defeating higher-order correlation attacks
///    that would otherwise become available across many rejection
///    retries on the same key.
///
/// All randomness for the SCA layer comes from a deterministic
/// SHAKE256-based `ScaRng` seeded with `(K ‖ rnd ‖ tr ‖ M')`, so
/// the function remains deterministic for fixed `rnd`. The masked
/// path produces signatures **bit-identical** to the unmasked path
/// — proven by the NIST ACVP siggen vectors, which the SCA build
/// passes unchanged.
///
/// # Errors
///
/// Returns [`MlDsaError::InvalidSecretKey`] if `sk` has incorrect length
/// (checked by the caller in the public API).
pub fn sign_internal<P: Params>(sk: &[u8], m_prime: &[u8], rnd: &[u8; 32]) -> Result<Vec<u8>, MlDsaError> {
    let k = P::K;
    let l = P::L;
    let gamma1 = P::GAMMA1;
    let gamma2 = P::GAMMA2;
    let beta = P::BETA;
    let omega = P::OMEGA;
    let c_tilde_len = P::LAMBDA / 4;

    // Decode secret key seeds (128 bytes on stack).
    //
    // indexed-sk: decode only rho/K/tr here; the polynomial vectors
    // s1/s2/t0 are decoded one-at-a-time below, directly into the
    // NTT-domain destination arrays, avoiding the 23 KB intermediate
    // tuple that sk_decode() would put on the stack.
    //
    // Default: sk_decode() returns the full tuple at once (simpler,
    // but 23 KB of stack for the return value alone).
    #[cfg(feature = "indexed-sk")]
    let (rho, k_seed, tr) = encode::sk_decode_seeds::<P>(sk);
    #[cfg(not(feature = "indexed-sk"))]
    let (rho, k_seed, tr, s1, s2, t0) = encode::sk_decode::<P>(sk);

    // ----- ŝ1, ŝ2, t̂0 = NTT(s1, s2, t0) ------------------------------
    //
    // This is the most leakage-prone step in Sign.
    //
    // indexed-sk: decode each polynomial from the packed sk directly
    // into the destination slot, then NTT in-place. Only one decoded
    // polynomial (1 KB) is live at a time instead of the full 23 KB
    // tuple from sk_decode.
    //
    // SCA-protected build: after NTT, each polynomial is additionally
    // split into masked shares (see below).
    //
    // Standard build: straight in-place Montgomery NTT.
    #[cfg(not(feature = "low-stack"))]
    let mut s1_hat = {
        #[cfg(feature = "indexed-sk")]
        {
            let mut v = [[0i32; N]; MAX_L];
            for i in 0..l {
                encode::sk_decode_s1::<P>(sk, i, &mut v[i]);
            }
            v
        }
        #[cfg(not(feature = "indexed-sk"))]
        {
            s1
        }
    };
    #[cfg(feature = "low-stack")]
    let mut s1_hat = {
        let mut v = poly_vec(l);
        #[cfg(feature = "indexed-sk")]
        for i in 0..l {
            encode::sk_decode_s1::<P>(sk, i, &mut v[i]);
        }
        #[cfg(not(feature = "indexed-sk"))]
        for i in 0..l {
            v[i] = s1[i];
        }
        v
    };

    #[cfg(not(feature = "low-stack"))]
    let mut s2_hat = {
        #[cfg(feature = "indexed-sk")]
        {
            let mut v = [[0i32; N]; MAX_K];
            for i in 0..k {
                encode::sk_decode_s2::<P>(sk, i, &mut v[i]);
            }
            v
        }
        #[cfg(not(feature = "indexed-sk"))]
        {
            s2
        }
    };
    #[cfg(feature = "low-stack")]
    let mut s2_hat = {
        let mut v = poly_vec(k);
        #[cfg(feature = "indexed-sk")]
        for i in 0..k {
            encode::sk_decode_s2::<P>(sk, i, &mut v[i]);
        }
        #[cfg(not(feature = "indexed-sk"))]
        for i in 0..k {
            v[i] = s2[i];
        }
        v
    };

    #[cfg(not(feature = "low-stack"))]
    let mut t0_hat = {
        #[cfg(feature = "indexed-sk")]
        {
            let mut v = [[0i32; N]; MAX_K];
            for i in 0..k {
                encode::sk_decode_t0::<P>(sk, i, &mut v[i]);
            }
            v
        }
        #[cfg(not(feature = "indexed-sk"))]
        {
            t0
        }
    };
    #[cfg(feature = "low-stack")]
    let mut t0_hat = {
        let mut v = poly_vec(k);
        #[cfg(feature = "indexed-sk")]
        for i in 0..k {
            encode::sk_decode_t0::<P>(sk, i, &mut v[i]);
        }
        #[cfg(not(feature = "indexed-sk"))]
        for i in 0..k {
            v[i] = t0[i];
        }
        v
    };
    #[cfg(feature = "sca-protected")]
    let (mut s1_hat_m, mut s2_hat_m, mut t0_hat_m, mut sca_rng) = {
        // Seed the SCA RNG from K ‖ rnd ‖ tr ‖ M'. K and rnd give us
        // the FIPS 204 hedged-signing entropy; tr and M' bind the
        // share stream to this particular (key, message) pair so two
        // signatures over different inputs use uncorrelated shares
        // even when rnd = 0 (deterministic / ACVP test mode).
        let mut sca_seed = [0u8; 64];
        {
            let mut h = sha3::shake256();
            h.absorb(b"quantica-mldsa-sca-seed-v1");
            h.absorb(&k_seed);
            h.absorb(rnd);
            h.absorb(&tr);
            h.absorb(m_prime);
            h.squeeze(&mut sca_seed);
        }
        let mut rng = ScaRng::from_seed(&sca_seed);

        // Step 1 — SPA defence: NTT each secret polynomial through
        // the Fisher-Yates shuffled NTT, drawing fresh per-level and
        // per-group permutations from the SCA RNG.
        for i in 0..l {
            shuffle::ntt_shuffled(&mut s1_hat[i], &mut rng)?;
        }
        for i in 0..k {
            shuffle::ntt_shuffled(&mut s2_hat[i], &mut rng)?;
        }
        for i in 0..k {
            shuffle::ntt_shuffled(&mut t0_hat[i], &mut rng)?;
        }

        // Step 2 — DPA defence: split each NTT-domain secret into two
        // additive shares mod q. The MaskedPoly::zero() initializer
        // is a stack-resident no-allocation array fill; the real
        // shares are written immediately below by MaskedPoly::mask.
        let mut s1m: [MaskedPoly; MAX_L] = core::array::from_fn(|_| MaskedPoly::zero());
        let mut s2m: [MaskedPoly; MAX_K] = core::array::from_fn(|_| MaskedPoly::zero());
        let mut t0m: [MaskedPoly; MAX_K] = core::array::from_fn(|_| MaskedPoly::zero());
        for i in 0..l {
            s1m[i] = MaskedPoly::mask(&s1_hat[i], &mut rng)?;
        }
        for i in 0..k {
            s2m[i] = MaskedPoly::mask(&s2_hat[i], &mut rng)?;
        }
        for i in 0..k {
            t0m[i] = MaskedPoly::mask(&t0_hat[i], &mut rng)?;
        }

        // Step 3 — wipe the unmasked NTT-domain buffers. From this
        // point on the secret only exists as `(share0, share1)` pairs;
        // any side-channel observation of `s1_hat[i]` etc. yields zero
        // information about the underlying coefficients.
        for i in 0..l {
            s1_hat[i] = [0i32; N];
        }
        for i in 0..k {
            s2_hat[i] = [0i32; N];
        }
        for i in 0..k {
            t0_hat[i] = [0i32; N];
        }
        (s1m, s2m, t0m, rng)
    };
    #[cfg(not(feature = "sca-protected"))]
    {
        // compressed-challenge: secrets stay in time domain for
        // schoolbook multiplication. No NTT needed.
        // small-secret: s1/s2 are converted to SmallPoly and NTT'd
        // via the i16 Kyber NTT instead. t0 still uses i32 NTT
        // (coefficients too large for i16).
        #[cfg(not(any(feature = "compressed-challenge", feature = "small-secret")))]
        {
            ntt_vec(&mut s1_hat, l);
            ntt_vec(&mut s2_hat, k);
            ntt_vec(&mut t0_hat, k);
        }
        #[cfg(all(feature = "small-secret", not(feature = "compressed-challenge")))]
        {
            // t0 still needs i32 NTT (coefficients up to 4096).
            ntt_vec(&mut t0_hat, k);
            // s1/s2 are converted to SmallPoly below; we don't NTT the i32 versions.
        }
    }

    // small-secret: convert s1/s2 to i16 SmallPoly and NTT via Kyber NTT.
    // The i32 s1_hat/s2_hat arrays are kept for any non-small-secret
    // code paths but are effectively unused when small-secret is on.
    #[cfg(feature = "small-secret")]
    let (s1_small, s2_small) = {
        let mut s1s: [SmallPoly; MAX_L] = core::array::from_fn(|_| SmallPoly::zero());
        let mut s2s: [SmallPoly; MAX_K] = core::array::from_fn(|_| SmallPoly::zero());
        for i in 0..l {
            s1s[i] = SmallPoly::from_i32(&s1_hat[i]);
            smallpoly::small_ntt(&mut s1s[i]);
        }
        for i in 0..k {
            s2s[i] = SmallPoly::from_i32(&s2_hat[i]);
            smallpoly::small_ntt(&mut s2s[i]);
        }
        (s1s, s2s)
    };

    // A-hat = ExpandA(rho)
    // Default: full matrix on stack (57 KB). Low-mem: recomputed on-the-fly.
    #[cfg(not(feature = "low-mem"))]
    let a_hat = sample::expand_a::<P>(&rho);

    // mu = H(tr || M')
    let mut mu = [0u8; 64];
    {
        let mut state = sha3::shake256();
        state.absorb(&tr);
        state.absorb(m_prime);
        state.squeeze(&mut mu);
    }

    // rho'' = H(K || rnd || mu)
    let mut rho_double_prime = [0u8; 64];
    {
        let mut state = sha3::shake256();
        state.absorb(&k_seed);
        state.absorb(rnd);
        state.absorb(&mu);
        state.squeeze(&mut rho_double_prime);
    }

    let mut kappa: u16 = 0;

    loop {
        // T1-A — refresh the persistent masked-secret-poly shares at
        // the **start** of every rejection iteration, before any
        // operation on them (Hermelink-Ning-Petri 2025/276 §4).
        // `s1_hat_m`, `s2_hat_m`, `t0_hat_m` survive across all
        // iterations (declared at line 530); without per-iteration
        // refresh, higher-order DPA aggregating traces over multiple
        // iterations sees correlated share pairs. Cost is one
        // `MaskedPoly::refresh` per polynomial per iteration —
        // identical to the previous end-of-cs/ct refresh placement;
        // KAT output bytes are unchanged because the mask cancels in
        // every `unmask()`.
        #[cfg(feature = "sca-protected")]
        {
            for i in 0..l {
                s1_hat_m[i].refresh(&mut sca_rng)?;
            }
            for i in 0..k {
                s2_hat_m[i].refresh(&mut sca_rng)?;
            }
            for i in 0..k {
                t0_hat_m[i].refresh(&mut sca_rng)?;
            }
        }

        // y = ExpandMask(rho'', kappa)
        //
        // sca-masked-y: sample y directly as arithmetic shares from
        // SHAKE256 (MaskedPoly::sample_expand_mask), keep it masked
        // through NTT + mat_vec_mul + iNTT. Only unmask y and w at
        // the end of the linear ops: w is about to be published via
        // w1 in c_tilde, and y is recoverable from z = y + cs1 in
        // the final signature anyway.
        //
        // Default: sample y in clear via expand_mask.
        #[cfg(not(feature = "sca-masked-y"))]
        let y = sample::expand_mask::<P>(&rho_double_prime, kappa);

        #[cfg(feature = "sca-masked-y")]
        let (y, w_precomputed) = {
            // Full masking pipeline: y stays split into two arithmetic
            // shares from sampling through NTT, A·y, and iNTT. Only
            // once w reaches its "about-to-be-published" form do we
            // unmask y and w together.
            //
            // Countermeasure references:
            //   ePrint 2025/276 — Hermelink–Ning–Petri, DPA on y
            //   ePrint 2025/582 — Rejected-signature timing leak

            // 1. Sample y as arithmetic shares from SHAKE256. The
            //    unmasked coefficient value only transits through CPU
            //    registers — never written to RAM.
            let mut y_m: [masked::MaskedPoly; MAX_L] = core::array::from_fn(|_| masked::MaskedPoly::zero());
            for r in 0..l {
                y_m[r] = masked::MaskedPoly::sample_expand_mask(
                    &rho_double_prime,
                    kappa + r as u16,
                    gamma1,
                    P::BITLEN_GAMMA1_MINUS1,
                );
            }

            // 2. Masked NTT into y_hat_m — y_m is preserved for the
            //    later time-domain unmask (needed by z = y + cs1).
            let mut y_hat_m: [masked::MaskedPoly; MAX_L] = core::array::from_fn(|_| masked::MaskedPoly::zero());
            for r in 0..l {
                y_hat_m[r].share0 = y_m[r].share0;
                y_hat_m[r].share1 = y_m[r].share1;
                masked::masked_ntt(&mut y_hat_m[r]);
            }

            // 3. Masked A · y_hat → w_m (NTT domain, masked).
            //    A is public; the matrix multiplication touches each
            //    share independently.
            let mut w_m: [masked::MaskedPoly; MAX_K] = core::array::from_fn(|_| masked::MaskedPoly::zero());
            #[cfg(not(feature = "low-mem"))]
            masked::masked_mat_vec_mul(&a_hat, &y_hat_m, k, l, &mut w_m);
            #[cfg(feature = "low-mem")]
            masked::masked_mat_vec_mul_lazy(&rho, &y_hat_m, k, l, &mut w_m);

            for r in 0..l {
                y_hat_m[r].zeroize();
            }

            // 4. Masked iNTT on each share.
            for i in 0..k {
                masked::masked_ntt_inv(&mut w_m[i]);
            }

            // 5. Unmask w — w1 = HighBits(w) is public (it ends up in
            //    c_tilde). Output in [0, q-1]; `decompose` handles that.
            let mut w_tmp = [[0i32; N]; MAX_K];
            for i in 0..k {
                w_tmp[i] = w_m[i].unmask();
                w_m[i].zeroize();
            }

            // 6. Unmask y to centered (-gamma1, gamma1] time domain,
            //    matching the default `expand_mask` output range.
            let mut y_out = [[0i32; N]; MAX_L];
            for r in 0..l {
                let um = y_m[r].unmask();
                for n in 0..N {
                    let mut v = um[n];
                    if v > Q / 2 {
                        v -= Q;
                    }
                    y_out[r][n] = v;
                }
                y_m[r].zeroize();
            }
            (y_out, w_tmp)
        };

        // ============================================================
        // Rejection-loop body.
        //
        // low-stack build: temporary polynomial vectors (w, w1, cs1,
        // cs2, w_minus_cs2, r0, ct0, neg_ct0, w_cs2_ct0) are
        // heap-allocated via Vec with scoped lifetimes and explicit
        // drop() calls. Only ~23 KB of heap is live at peak instead
        // of ~96 KB of stack.
        //
        // Default build: everything on the stack as fixed arrays.
        // ============================================================

        // w = NTT^{-1}(A-hat * NTT(y))
        //
        // Default path: compute y_hat = NTT(y), then w_tmp = iNTT(A·y_hat).
        // sca-masked-y: w was already computed in the masked block
        // above (w_precomputed) — the unmasked y was never in RAM.

        // --- Compute w and w1, then derive c_tilde ---------------
        //
        // compressed-poly: after iNTT(w), pack w into 3-byte/coeff
        // compressed form (−25% RAM), then derive w1 and w_minus_cs2
        // from the compressed representation.
        #[cfg(not(feature = "compressed-poly"))]
        {
            #[cfg(not(feature = "low-stack"))]
            let mut _w_full = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut _w_full = poly_vec(k);
            // (assigned below, used via w_ref)
        }

        // Compute w into a temporary full-poly buffer, then either
        // keep it (default) or compress it (compressed-poly).
        #[cfg(not(feature = "sca-masked-y"))]
        let mut w_tmp = {
            let mut y_hat = y;
            ntt_vec(&mut y_hat, l);
            let mut wt = [[0i32; N]; MAX_K];
            #[cfg(not(feature = "low-mem"))]
            mat_vec_mul(&a_hat, &y_hat, k, l, &mut wt);
            #[cfg(feature = "low-mem")]
            mat_vec_mul_lazy(&rho, &y_hat, k, l, &mut wt);
            ntt_inv_vec(&mut wt, k);
            wt
        };
        #[cfg(feature = "sca-masked-y")]
        let mut w_tmp = w_precomputed;

        // compressed-poly: pack w into 3-byte/coeff storage.
        #[cfg(feature = "compressed-poly")]
        let w_comp = {
            let mut wc = compressed::CompressedVecK::new(k);
            for i in 0..k {
                // Reduce to [0, q-1] before packing.
                for c in w_tmp[i].iter_mut() {
                    *c = mod_q(*c);
                }
                wc.pack(i, &w_tmp[i]);
            }
            wc
        };

        // w1 = HighBits(w) — works on the full-poly tmp (before we drop it).
        #[cfg(not(feature = "low-stack"))]
        let mut w1 = [[0i32; N]; MAX_K];
        #[cfg(feature = "low-stack")]
        let mut w1 = poly_vec(k);
        decompose::high_bits_vec(&w_tmp, gamma2, &mut w1, k);

        // In non-compressed mode, keep w_tmp as "w" for later vec_sub.
        // In compressed mode, drop w_tmp (we'll read from w_comp).
        #[cfg(not(feature = "compressed-poly"))]
        let w = w_tmp;
        #[cfg(feature = "compressed-poly")]
        drop(w_tmp);

        let w1_encoded = encode::w1_encode::<P>(&w1);
        #[cfg(feature = "low-stack")]
        drop(w1);

        let mut c_tilde_buf = [0u8; 64];
        {
            let mut state = sha3::shake256();
            state.absorb(&mu);
            state.absorb(&w1_encoded);
            state.squeeze(&mut c_tilde_buf[..c_tilde_len]);
        }
        let c_tilde = &c_tilde_buf[..c_tilde_len];

        // --- challenge computation --------------------------------
        //
        // Default: c_hat = NTT(SampleInBall(c_tilde)), 2 KB stack.
        // compressed-challenge: compress c into 68 bytes and use
        // schoolbook multiplication in time domain, saving ~2 KB.
        let c = sample::sample_in_ball::<P>(c_tilde);
        #[cfg(not(feature = "compressed-challenge"))]
        let c_hat = {
            let mut ch = c;
            for coeff in ch.iter_mut() {
                *coeff = mod_q(*coeff);
            }
            ntt::ntt(&mut ch);
            ch
        };
        #[cfg(feature = "compressed-challenge")]
        let c_comp = {
            let mut cc = [0u8; compressed::COMPRESSED_CHALLENGE_BYTES];
            compressed::challenge_compress(&mut cc, &c, P::TAU);
            cc
        };
        // small-secret: also convert c to SmallPoly NTT for basemul.
        #[cfg(feature = "small-secret")]
        let c_small_ntt = {
            let mut cs = SmallPoly::from_i32(&c);
            smallpoly::small_ntt(&mut cs);
            cs
        };

        // ============================================================
        // union-buffer path: single 1 KB workspace reused per poly.
        // Processes L + K iterations sequentially, only z and h persist.
        // ============================================================
        #[cfg(feature = "union-buffer")]
        {
            let mut z = [[0i32; N]; MAX_L];
            let mut tmp = [0i32; N];
            let mut rejected = false;

            // Phase 1: z[i] = y[i] + c*s1[i]
            for l_idx in 0..l {
                #[cfg(all(not(feature = "compressed-challenge"), not(feature = "small-secret")))]
                {
                    tmp = ntt::pointwise_mul(&c_hat, &s1_hat[l_idx]);
                    ntt::ntt_inv(&mut tmp);
                }
                #[cfg(feature = "compressed-challenge")]
                {
                    tmp = [0i32; N];
                    compressed::schoolbook_mul_add(&mut tmp, &c_comp, &s1_hat[l_idx], P::TAU);
                }
                #[cfg(all(feature = "small-secret", not(feature = "compressed-challenge")))]
                {
                    tmp = smallpoly::small_basemul_invntt_widen(&c_small_ntt, &s1_small[l_idx]);
                }
                z[l_idx] = ntt::poly_add(&y[l_idx], &tmp);
            }
            if !check_norm_vec(&z, gamma1 - beta, l) {
                kappa += l as u16;
                continue;
            }

            // Phase 2: per k_idx — cs2, r0, ct0, hint
            let mut h = [[0i32; N]; MAX_K];
            let mut total_hints = 0usize;
            let mut wbuf = [0i32; N];

            for k_idx in 0..k {
                if rejected {
                    break;
                }
                // cs2 → tmp
                #[cfg(all(not(feature = "compressed-challenge"), not(feature = "small-secret")))]
                {
                    tmp = ntt::pointwise_mul(&c_hat, &s2_hat[k_idx]);
                    ntt::ntt_inv(&mut tmp);
                }
                #[cfg(feature = "compressed-challenge")]
                {
                    tmp = [0i32; N];
                    compressed::schoolbook_mul_add(&mut tmp, &c_comp, &s2_hat[k_idx], P::TAU);
                }
                #[cfg(all(feature = "small-secret", not(feature = "compressed-challenge")))]
                {
                    tmp = smallpoly::small_basemul_invntt_widen(&c_small_ntt, &s2_small[k_idx]);
                }

                // wbuf = w[k_idx] - cs2
                #[cfg(not(feature = "compressed-poly"))]
                for j in 0..N {
                    wbuf[j] = w[k_idx][j] - tmp[j];
                }
                #[cfg(feature = "compressed-poly")]
                w_comp.sub_into(k_idx, &tmp, &mut wbuf);

                // r0 check in tmp
                for j in 0..N {
                    tmp[j] = decompose::low_bits(wbuf[j], gamma2);
                }
                if !check_norm(&tmp, gamma2 - beta) {
                    rejected = true;
                    continue;
                }

                // ct0 → tmp
                #[cfg(not(feature = "compressed-challenge"))]
                {
                    tmp = ntt::pointwise_mul(&c_hat, &t0_hat[k_idx]);
                    ntt::ntt_inv(&mut tmp);
                }
                #[cfg(feature = "compressed-challenge")]
                {
                    tmp = [0i32; N];
                    compressed::schoolbook_mul_add(&mut tmp, &c_comp, &t0_hat[k_idx], P::TAU);
                }

                if !check_norm(&tmp, gamma2) {
                    rejected = true;
                    continue;
                }

                // hint for this k_idx
                for j in 0..N {
                    h[k_idx][j] = decompose::make_hint(mod_q(-tmp[j]), wbuf[j] + tmp[j], gamma2);
                    if h[k_idx][j] == 1 {
                        total_hints += 1;
                    }
                }
            }

            if rejected || total_hints > omega {
                kappa += l as u16;
                continue;
            }

            // Center z and encode
            for poly in z[..l].iter_mut() {
                for c in poly.iter_mut() {
                    *c = mod_q(*c);
                    if *c > Q / 2 {
                        *c -= Q;
                    }
                }
            }
            let sig = encode::sig_encode::<P>(c_tilde, &z, &h);
            return Ok(sig);
        }

        // ============================================================
        // Standard path (non-union-buffer)
        // ============================================================
        #[cfg(not(feature = "union-buffer"))]
        {
            // --- cs1 = ĉ · ŝ1, then z = y + cs1 --------------------
            #[cfg(not(feature = "low-stack"))]
            let mut cs1 = [[0i32; N]; MAX_L];
            #[cfg(feature = "low-stack")]
            let mut cs1 = poly_vec(l);

            #[cfg(feature = "sca-protected")]
            for i in 0..l {
                let prod = masked::masked_pointwise_mul_public(&s1_hat_m[i], &c_hat);
                cs1[i] = prod.unmask();
                ntt::ntt_inv(&mut cs1[i]);
                // refresh of s1_hat_m happens at the head of the
                // next rejection iteration (T1-A).
            }
            #[cfg(all(
                not(feature = "sca-protected"),
                not(feature = "compressed-challenge"),
                not(feature = "small-secret")
            ))]
            for i in 0..l {
                cs1[i] = ntt::pointwise_mul(&c_hat, &s1_hat[i]);
                ntt::ntt_inv(&mut cs1[i]);
            }
            #[cfg(all(not(feature = "sca-protected"), feature = "compressed-challenge"))]
            for i in 0..l {
                cs1[i] = [0i32; N];
                compressed::schoolbook_mul_add(&mut cs1[i], &c_comp, &s1_hat[i], P::TAU);
            }
            #[cfg(all(
                not(feature = "sca-protected"),
                feature = "small-secret",
                not(feature = "compressed-challenge")
            ))]
            for i in 0..l {
                cs1[i] = smallpoly::small_basemul_invntt_widen(&c_small_ntt, &s1_small[i]);
            }

            #[cfg(not(feature = "low-stack"))]
            let mut z = [[0i32; N]; MAX_L];
            #[cfg(feature = "low-stack")]
            let mut z = poly_vec(l);
            vec_add(&y, &cs1, &mut z, l);
            // cs1 no longer needed.
            #[cfg(feature = "low-stack")]
            drop(cs1);

            // --- cs2, w_minus_cs2, r0 --------------------------------
            #[cfg(not(feature = "low-stack"))]
            let mut cs2 = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut cs2 = poly_vec(k);

            #[cfg(feature = "sca-protected")]
            for i in 0..k {
                let prod = masked::masked_pointwise_mul_public(&s2_hat_m[i], &c_hat);
                cs2[i] = prod.unmask();
                ntt::ntt_inv(&mut cs2[i]);
                // refresh of s2_hat_m happens at the head of the
                // next rejection iteration (T1-A).
            }
            #[cfg(all(
                not(feature = "sca-protected"),
                not(feature = "compressed-challenge"),
                not(feature = "small-secret")
            ))]
            for i in 0..k {
                cs2[i] = ntt::pointwise_mul(&c_hat, &s2_hat[i]);
                ntt::ntt_inv(&mut cs2[i]);
            }
            #[cfg(all(not(feature = "sca-protected"), feature = "compressed-challenge"))]
            for i in 0..k {
                cs2[i] = [0i32; N];
                compressed::schoolbook_mul_add(&mut cs2[i], &c_comp, &s2_hat[i], P::TAU);
            }
            #[cfg(all(
                not(feature = "sca-protected"),
                feature = "small-secret",
                not(feature = "compressed-challenge")
            ))]
            for i in 0..k {
                cs2[i] = smallpoly::small_basemul_invntt_widen(&c_small_ntt, &s2_small[i]);
            }

            #[cfg(not(feature = "low-stack"))]
            let mut w_minus_cs2 = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut w_minus_cs2 = poly_vec(k);
            #[cfg(not(feature = "compressed-poly"))]
            vec_sub(&w, &cs2, &mut w_minus_cs2, k);
            #[cfg(feature = "compressed-poly")]
            for i in 0..k {
                w_comp.sub_into(i, &cs2[i], &mut w_minus_cs2[i]);
            }
            // cs2 and w/w_comp no longer needed for the norm checks.
            #[cfg(feature = "low-stack")]
            drop(cs2);
            #[cfg(all(feature = "low-stack", not(feature = "compressed-poly")))]
            drop(w);
            #[cfg(feature = "compressed-poly")]
            drop(w_comp);

            #[cfg(not(feature = "low-stack"))]
            let mut r0 = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut r0 = poly_vec(k);
            decompose::low_bits_vec(&w_minus_cs2, gamma2, &mut r0, k);

            // Norm checks. Standard build: early-abort for performance.
            // sca-ct-rejection build: collect all flags and decide at end.
            #[cfg(not(feature = "sca-ct-rejection"))]
            {
                if !check_norm_vec(&z, gamma1 - beta, l) {
                    kappa += l as u16;
                    continue;
                }
                if !check_norm_vec(&r0, gamma2 - beta, k) {
                    kappa += l as u16;
                    continue;
                }
            }
            #[cfg(feature = "sca-ct-rejection")]
            let mut _reject_flag = {
                let z_ok = check_norm_vec(&z, gamma1 - beta, l);
                let r0_ok = check_norm_vec(&r0, gamma2 - beta, k);
                !(z_ok & r0_ok)
            };
            // r0 no longer needed.
            #[cfg(feature = "low-stack")]
            drop(r0);

            // --- ct0, hint computation --------------------------------
            #[cfg(not(feature = "low-stack"))]
            let mut ct0 = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut ct0 = poly_vec(k);

            #[cfg(feature = "sca-protected")]
            for i in 0..k {
                let prod = masked::masked_pointwise_mul_public(&t0_hat_m[i], &c_hat);
                ct0[i] = prod.unmask();
                ntt::ntt_inv(&mut ct0[i]);
                // refresh of t0_hat_m happens at the head of the
                // next rejection iteration (T1-A).
            }
            #[cfg(all(not(feature = "sca-protected"), not(feature = "compressed-challenge")))]
            for i in 0..k {
                ct0[i] = ntt::pointwise_mul(&c_hat, &t0_hat[i]);
                ntt::ntt_inv(&mut ct0[i]);
            }
            #[cfg(all(not(feature = "sca-protected"), feature = "compressed-challenge"))]
            for i in 0..k {
                ct0[i] = [0i32; N];
                compressed::schoolbook_mul_add(&mut ct0[i], &c_comp, &t0_hat[i], P::TAU);
            }

            // Check ||ct0||_inf < gamma2
            #[cfg(not(feature = "sca-ct-rejection"))]
            {
                if !check_norm_vec(&ct0, gamma2, k) {
                    kappa += l as u16;
                    continue;
                }
            }
            #[cfg(feature = "sca-ct-rejection")]
            {
                _reject_flag |= !check_norm_vec(&ct0, gamma2, k);
            }

            // h = MakeHint(-ct0, w_minus_cs2 + ct0)
            #[cfg(not(feature = "low-stack"))]
            let mut w_cs2_ct0 = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut w_cs2_ct0 = poly_vec(k);
            vec_add(&w_minus_cs2, &ct0, &mut w_cs2_ct0, k);

            #[cfg(not(feature = "low-stack"))]
            let mut neg_ct0 = [[0i32; N]; MAX_K];
            #[cfg(feature = "low-stack")]
            let mut neg_ct0 = poly_vec(k);
            for i in 0..k {
                for j in 0..N {
                    neg_ct0[i][j] = mod_q(-ct0[i][j]);
                }
            }
            // ct0 and w_minus_cs2 no longer needed.
            #[cfg(feature = "low-stack")]
            {
                drop(ct0);
                drop(w_minus_cs2);
            }

            let (h, num_ones) = decompose::make_hint_vec(&neg_ct0, &w_cs2_ct0, gamma2, k);
            // neg_ct0 and w_cs2_ct0 no longer needed.
            #[cfg(feature = "low-stack")]
            {
                drop(neg_ct0);
                drop(w_cs2_ct0);
            }

            #[cfg(not(feature = "sca-ct-rejection"))]
            {
                if num_ones > omega {
                    kappa += l as u16;
                    continue;
                }
            }
            #[cfg(feature = "sca-ct-rejection")]
            {
                _reject_flag |= num_ones > omega;
                if _reject_flag {
                    kappa += l as u16;
                    continue;
                }
            }

            // Center z coefficients to [-gamma1+1, gamma1] before encoding
            for poly in z[..l].iter_mut() {
                for c in poly.iter_mut() {
                    *c = mod_q(*c);
                    if *c > Q / 2 {
                        *c -= Q;
                    }
                }
            }

            // Encode signature
            let sig = encode::sig_encode::<P>(c_tilde, &z, &h);
            return Ok(sig);
        } // end #[cfg(not(feature = "union-buffer"))]
    }
}

/// Verify a signature against a pre-formatted message.
///
/// Implements Algorithm 8 of FIPS 204 (ML-DSA.Verify_internal).
///
/// Recomputes the commitment w1' from the public key, signature components
/// (c_tilde, z, h), and the message hash mu. Verification succeeds when the
/// recomputed commitment hash matches the c_tilde embedded in the signature.
///
/// - `pk`: encoded public key (must be `P::PK_LEN` bytes).
/// - `m_prime`: pre-formatted message.
/// - `sig`: encoded signature (must be `P::SIG_LEN` bytes).
///
/// Returns `Ok(true)` if the signature is valid, `Ok(false)` otherwise.
///
/// # Errors
///
/// - [`MlDsaError::InvalidPublicKey`] if `pk` has the wrong length.
/// - [`MlDsaError::InvalidSignature`] if `sig` has the wrong length.
pub fn verify_internal<P: Params>(pk: &[u8], m_prime: &[u8], sig: &[u8]) -> Result<bool, MlDsaError> {
    let k = P::K;
    let l = P::L;
    let gamma1 = P::GAMMA1;
    let gamma2 = P::GAMMA2;
    let beta = P::BETA;
    let omega = P::OMEGA;
    let c_tilde_len = P::LAMBDA / 4;

    if pk.len() != P::PK_LEN {
        return Err(MlDsaError::InvalidPublicKey);
    }
    if sig.len() != P::SIG_LEN {
        return Err(MlDsaError::InvalidSignature);
    }

    // Decode public key
    let (rho, t1) = encode::pk_decode::<P>(pk);

    // tr = H(pk) (64 bytes)
    let mut tr = [0u8; 64];
    sha3::shake256_digest(pk, &mut tr);

    // Decode signature
    let (c_tilde, z, h) = match encode::sig_decode::<P>(sig) {
        Some(x) => x,
        None => return Ok(false),
    };

    // Check ||z||_inf < gamma1 - beta
    if !check_norm_vec(&z, gamma1 - beta, l) {
        return Ok(false);
    }

    // A-hat = ExpandA(rho)
    #[cfg(not(feature = "low-mem"))]
    let a_hat = sample::expand_a::<P>(&rho);

    // mu = H(tr || M')
    let mut mu = [0u8; 64];
    {
        let mut state = sha3::shake256();
        state.absorb(&tr);
        state.absorb(m_prime);
        state.squeeze(&mut mu);
    }

    // c = SampleInBall(c_tilde)
    let mut c = sample::sample_in_ball::<P>(&c_tilde);
    for coeff in c.iter_mut() {
        *coeff = mod_q(*coeff);
    }
    let mut c_hat = c;
    ntt::ntt(&mut c_hat);

    // z_hat = NTT(z)
    let mut z_hat = z;
    ntt_vec(&mut z_hat, l);

    // w'_approx = NTT^{-1}(A-hat * z_hat - c_hat * NTT(t1 * 2^d))
    // First compute NTT(t1 * 2^d)
    let mut t1_2d_hat = [[0i32; N]; MAX_K];
    for i in 0..k {
        for j in 0..N {
            t1_2d_hat[i][j] = mod_q(t1[i][j] * (1 << D));
        }
        ntt::ntt(&mut t1_2d_hat[i]);
    }

    // A-hat * z_hat
    let mut az = [[0i32; N]; MAX_K];
    #[cfg(not(feature = "low-mem"))]
    mat_vec_mul(&a_hat, &z_hat, k, l, &mut az);
    #[cfg(feature = "low-mem")]
    mat_vec_mul_lazy(&rho, &z_hat, k, l, &mut az);

    // c_hat * t1_2d_hat (component-wise)
    let mut ct1 = [[0i32; N]; MAX_K];
    for i in 0..k {
        ct1[i] = ntt::pointwise_mul(&c_hat, &t1_2d_hat[i]);
    }

    // w'_approx = NTT^{-1}(az - ct1)
    let mut w_approx = [[0i32; N]; MAX_K];
    vec_sub(&az, &ct1, &mut w_approx, k);
    ntt_inv_vec(&mut w_approx, k);

    // w1' = UseHint(h, w'_approx)
    let w1_prime = decompose::use_hint_vec(&h, &w_approx, gamma2, k);

    // Recompute c_tilde' = H(mu || w1Encode(w1'))
    let w1_encoded = encode::w1_encode::<P>(&w1_prime);
    let mut c_tilde_prime = vec![0u8; c_tilde_len];
    {
        let mut state = sha3::shake256();
        state.absorb(&mu);
        state.absorb(&w1_encoded);
        state.squeeze(&mut c_tilde_prime);
    }

    // Check c_tilde == c_tilde'
    // Also verify hint weight
    let mut hint_count = 0usize;
    for i in 0..k {
        for &c in h[i].iter() {
            hint_count += c as usize;
        }
    }
    if hint_count > omega {
        return Ok(false);
    }

    Ok(c_tilde == c_tilde_prime)
}

/// Generate an ML-DSA key pair.
///
/// Implements Algorithm 1 of FIPS 204 (ML-DSA.KeyGen). Draws 32 random
/// bytes from `rng` and delegates to [`keygen_internal`].
///
/// Returns `(pk, sk)` as byte vectors.
///
/// # Errors
///
/// Returns [`MlDsaError::RngFailure`] if the RNG cannot provide bytes.
pub fn keygen<P: Params>(rng: &mut dyn CryptoRng) -> Result<(Vec<u8>, Vec<u8>), MlDsaError> {
    let mut xi = [0u8; 32];
    rng.fill_bytes(&mut xi)?;
    let result = keygen_internal::<P>(&xi);
    Ok(result)
}

/// Sign a message with an optional context string (hedged mode).
///
/// Implements Algorithm 2 of FIPS 204 (ML-DSA.Sign). Constructs the
/// pre-formatted message `M' = 0x00 || len(ctx) || ctx || msg`, draws 32
/// random bytes for hedged signing, and calls `sign_internal`.
///
/// - `sk`: secret key (must be `P::SK_LEN` bytes).
/// - `msg`: message to sign.
/// - `ctx`: optional context string (at most 255 bytes).
/// - `rng`: source of randomness for the hedged nonce.
///
/// # Errors
///
/// - [`MlDsaError::ContextTooLong`] if `ctx` exceeds 255 bytes.
/// - [`MlDsaError::InvalidSecretKey`] if `sk` has the wrong length.
/// - [`MlDsaError::RngFailure`] if the RNG cannot provide bytes.
pub fn sign<P: Params>(sk: &[u8], msg: &[u8], ctx: &[u8], rng: &mut dyn CryptoRng) -> Result<Vec<u8>, MlDsaError> {
    if ctx.len() > 255 {
        return Err(MlDsaError::ContextTooLong);
    }
    if sk.len() != P::SK_LEN {
        return Err(MlDsaError::InvalidSecretKey);
    }

    // M' = 0x00 || len(ctx) || ctx || M
    let mut m_prime = Vec::with_capacity(1 + 1 + ctx.len() + msg.len());
    m_prime.push(0x00);
    m_prime.push(ctx.len() as u8);
    m_prime.extend_from_slice(ctx);
    m_prime.extend_from_slice(msg);

    // Random bytes for hedged signing
    let mut rnd = [0u8; 32];
    rng.fill_bytes(&mut rnd)?;

    sign_internal::<P>(sk, &m_prime, &rnd)
}

/// Verify a signature on a message with an optional context string.
///
/// Implements Algorithm 3 of FIPS 204 (ML-DSA.Verify). Constructs the
/// pre-formatted message `M' = 0x00 || len(ctx) || ctx || msg` and
/// delegates to [`verify_internal`].
///
/// - `pk`: public key (must be `P::PK_LEN` bytes).
/// - `msg`: the signed message.
/// - `ctx`: the context string used at signing time (at most 255 bytes).
/// - `sig`: the signature (must be `P::SIG_LEN` bytes).
///
/// Returns `Ok(true)` if the signature is valid, `Ok(false)` otherwise.
///
/// # Errors
///
/// - [`MlDsaError::ContextTooLong`] if `ctx` exceeds 255 bytes.
/// - [`MlDsaError::InvalidPublicKey`] if `pk` has the wrong length.
/// - [`MlDsaError::InvalidSignature`] if `sig` has the wrong length.
pub fn verify<P: Params>(pk: &[u8], msg: &[u8], ctx: &[u8], sig: &[u8]) -> Result<bool, MlDsaError> {
    if ctx.len() > 255 {
        return Err(MlDsaError::ContextTooLong);
    }
    if pk.len() != P::PK_LEN {
        return Err(MlDsaError::InvalidPublicKey);
    }
    if sig.len() != P::SIG_LEN {
        return Err(MlDsaError::InvalidSignature);
    }

    // M' = 0x00 || len(ctx) || ctx || M
    let mut m_prime = Vec::with_capacity(1 + 1 + ctx.len() + msg.len());
    m_prime.push(0x00);
    m_prime.push(ctx.len() as u8);
    m_prime.extend_from_slice(ctx);
    m_prime.extend_from_slice(msg);

    verify_internal::<P>(pk, &m_prime, sig)
}