krypteia-quantica 0.1.0

//! Encoding and decoding algorithms for ML-DSA (FIPS 204, Algorithms 9-28).
//!
//! Provides bit-level packing and unpacking of polynomials and polynomial
//! vectors into compact byte representations, as well as the Power2Round
//! decomposition used during key generation.
//!
//! All vector functions accept slices (`&[[i32; N]]`) instead of `&Vec<...>`
//! to avoid requiring heap-allocated containers.

use super::ntt::mod_q;
use super::params::{D, MAX_K, MAX_L, N, Params, Q};
use alloc::vec::Vec;

// ============================================================
// Bit/byte conversion utilities (Algorithms 9-13)
// ============================================================

/// Pack a polynomial whose coefficients lie in [0, b].
///
/// Implements Algorithm 16 of FIPS 204 (SimpleBitPack). Each coefficient
/// is stored using `bitlen(b)` bits in little-endian bit order.
///
/// - `w`: input polynomial with coefficients in [0, b].
/// - `b`: upper bound on coefficient values.
/// - `out`: output buffer (must have length >= N * bitlen(b) / 8).
pub fn simple_bit_pack(w: &[i32; N], b: u32, out: &mut [u8]) {
    let bits = 32 - b.leading_zeros() as usize; // bitlen(b)
    let mut bit_pos = 0usize;
    // Clear output
    for byte in out.iter_mut() {
        *byte = 0;
    }
    for i in 0..N {
        let val = w[i] as u32;
        for bit in 0..bits {
            if (val >> bit) & 1 == 1 {
                out[bit_pos / 8] |= 1 << (bit_pos % 8);
            }
            bit_pos += 1;
        }
    }
}

/// Unpack a polynomial whose coefficients lie in [0, b].
///
/// Implements Algorithm 18 of FIPS 204 (SimpleBitUnpack). Inverse of
/// [`simple_bit_pack`].
///
/// - `data`: packed byte data.
/// - `b`: upper bound on coefficient values.
/// - `w`: output polynomial (filled with decoded coefficients).
pub fn simple_bit_unpack(data: &[u8], b: u32, w: &mut [i32; N]) {
    let bits = 32 - b.leading_zeros() as usize;
    let mut bit_pos = 0usize;
    for i in 0..N {
        let mut val = 0u32;
        for bit in 0..bits {
            if (data[bit_pos / 8] >> (bit_pos % 8)) & 1 == 1 {
                val |= 1 << bit;
            }
            bit_pos += 1;
        }
        w[i] = val as i32;
    }
}

/// Pack a polynomial whose coefficients lie in [-a, b].
///
/// Implements Algorithm 17 of FIPS 204 (BitPack). Stores each coefficient
/// as `(b - coeff)`, mapping the range [-a, b] to [0, a+b], then packs
/// using `bitlen(a+b)` bits per coefficient.
///
/// - `w`: input polynomial with coefficients in [-a, b].
/// - `a`: magnitude of the negative bound.
/// - `b`: positive bound.
/// - `out`: output buffer (must have length >= N * bitlen(a+b) / 8).
pub fn bit_pack(w: &[i32; N], a: u32, b: u32, out: &mut [u8]) {
    let range = a + b;
    let bits = 32 - range.leading_zeros() as usize;
    let mut bit_pos = 0usize;
    for byte in out.iter_mut() {
        *byte = 0;
    }
    for i in 0..N {
        let val = (b as i32 - w[i]) as u32;
        for bit in 0..bits {
            if (val >> bit) & 1 == 1 {
                out[bit_pos / 8] |= 1 << (bit_pos % 8);
            }
            bit_pos += 1;
        }
    }
}

/// Unpack a polynomial whose coefficients lie in [-a, b].
///
/// Implements Algorithm 19 of FIPS 204 (BitUnpack). Inverse of [`bit_pack`].
///
/// - `data`: packed byte data.
/// - `a`: magnitude of the negative bound.
/// - `b`: positive bound.
/// - `w`: output polynomial (filled with decoded coefficients in [-a, b]).
pub fn bit_unpack(data: &[u8], a: u32, b: u32, w: &mut [i32; N]) {
    let range = a + b;
    let bits = 32 - range.leading_zeros() as usize;
    let mut bit_pos = 0usize;
    for i in 0..N {
        let mut val = 0u32;
        for bit in 0..bits {
            if (data[bit_pos / 8] >> (bit_pos % 8)) & 1 == 1 {
                val |= 1 << bit;
            }
            bit_pos += 1;
        }
        w[i] = b as i32 - val as i32;
    }
}

/// Pack a hint vector into bytes.
///
/// Implements Algorithm 20 of FIPS 204 (HintBitPack). The hint vector `h`
/// consists of k binary polynomials with at most `omega` total non-zero
/// entries. The output uses `omega + k` bytes: the first `omega` bytes
/// store the indices of non-zero coefficients, and the last k bytes store
/// cumulative index counts.
///
/// - `h`: hint vector (k polynomials with entries in {0, 1}).
/// - `out`: output buffer (must have length omega + k).
pub fn hint_bit_pack<P: Params>(h: &[[i32; N]], out: &mut [u8]) {
    let omega = P::OMEGA;
    let k = P::K;
    // out has length omega + k
    for byte in out.iter_mut() {
        *byte = 0;
    }
    let mut idx = 0usize;
    for i in 0..k {
        for j in 0..N {
            if h[i][j] != 0 {
                out[idx] = j as u8;
                idx += 1;
            }
        }
        out[omega + i] = idx as u8;
    }
}

/// Unpack a hint vector from bytes.
///
/// Implements Algorithm 21 of FIPS 204 (HintBitUnpack). Inverse of
/// [`hint_bit_pack`]. Performs validity checks on the encoding: indices
/// must be strictly increasing within each polynomial, and cumulative
/// counts must be non-decreasing and within bounds.
///
/// Returns `None` if the encoding is malformed.
pub fn hint_bit_unpack<P: Params>(data: &[u8]) -> Option<[[i32; N]; MAX_K]> {
    let omega = P::OMEGA;
    let k = P::K;
    let mut h = [[0i32; N]; MAX_K];
    let mut idx = 0usize;
    for i in 0..k {
        let upper = data[omega + i] as usize;
        if upper < idx || upper > omega {
            return None;
        }
        let first = idx;
        while idx < upper {
            // Check ordering (indices must be strictly increasing within each polynomial)
            if idx > first && data[idx] <= data[idx - 1] {
                return None;
            }
            let j = data[idx] as usize;
            if j >= N {
                return None;
            }
            h[i][j] = 1;
            idx += 1;
        }
    }
    // Check remaining bytes are zero
    while idx < omega {
        if data[idx] != 0 {
            return None;
        }
        idx += 1;
    }
    Some(h)
}

/// Encode a public key as bytes.
///
/// Implements Algorithm 22 of FIPS 204 (pkEncode). The public key consists
/// of the 32-byte seed `rho` followed by the k polynomials of `t1`, each
/// packed with 10 bits per coefficient (since t1 values lie in [0, 1023]).
///
/// - `rho`: 32-byte public seed for matrix A generation.
/// - `t1`: vector of k polynomials (the high bits of t).
///
/// Returns a byte vector of length `P::PK_LEN`.
pub fn pk_encode<P: Params>(rho: &[u8; 32], t1: &[[i32; N]]) -> Vec<u8> {
    let k = P::K;
    // t1 coefficients are in [0, 2^{bitlen(q-1)-d} - 1] = [0, 2^10 - 1] = [0, 1023]
    let coeff_bits = 10; // bitlen(q-1) - d = 23 - 13
    let poly_bytes = N * coeff_bits / 8; // 256 * 10 / 8 = 320
    let mut pk = vec![0u8; P::PK_LEN];
    pk[..32].copy_from_slice(rho);
    for i in 0..k {
        let offset = 32 + i * poly_bytes;
        simple_bit_pack(&t1[i], 1023, &mut pk[offset..offset + poly_bytes]);
    }
    pk
}

/// Decode a public key from bytes.
///
/// Implements Algorithm 23 of FIPS 204 (pkDecode). Inverse of [`pk_encode`].
///
/// Returns `(rho, t1)` where `rho` is the 32-byte seed and `t1` is the
/// fixed array of k polynomials with coefficients in [0, 1023].
pub fn pk_decode<P: Params>(pk: &[u8]) -> ([u8; 32], [[i32; N]; MAX_K]) {
    let k = P::K;
    let poly_bytes = 320; // 256 * 10 / 8
    let mut rho = [0u8; 32];
    rho.copy_from_slice(&pk[..32]);
    let mut t1 = [[0i32; N]; MAX_K];
    for i in 0..k {
        let offset = 32 + i * poly_bytes;
        simple_bit_unpack(&pk[offset..offset + poly_bytes], 1023, &mut t1[i]);
    }
    (rho, t1)
}

/// Encode a secret key as bytes.
///
/// Implements Algorithm 24 of FIPS 204 (skEncode). The secret key is the
/// concatenation of `rho` (32 bytes), `K` (32 bytes), `tr` (64 bytes),
/// the l polynomials of `s1` and k polynomials of `s2` (each packed with
/// `bitlen(2*eta)` bits per coefficient), and the k polynomials of `t0`
/// (each packed with 13 bits per coefficient).
///
/// - `rho`: 32-byte public seed.
/// - `k_seed`: 32-byte secret seed K.
/// - `tr`: 64-byte hash of the public key.
/// - `s1`: secret vector of l polynomials with coefficients in [-eta, eta].
/// - `s2`: secret vector of k polynomials with coefficients in [-eta, eta].
/// - `t0`: low-order bits vector (k polynomials from Power2Round).
///
/// Returns a byte vector of length `P::SK_LEN`.
pub fn sk_encode<P: Params>(
    rho: &[u8; 32],
    k_seed: &[u8; 32],
    tr: &[u8; 64],
    s1: &[[i32; N]],
    s2: &[[i32; N]],
    t0: &[[i32; N]],
) -> Vec<u8> {
    let eta = P::ETA as u32;
    let l = P::L;
    let k = P::K;
    let eta_bits = P::BITLEN_2ETA;
    let poly_eta_bytes = N * eta_bits / 8;
    let d = D;
    let poly_t0_bytes = N * d / 8; // 256 * 13 / 8 = 416

    let mut sk = vec![0u8; P::SK_LEN];
    let mut offset = 0;

    sk[offset..offset + 32].copy_from_slice(rho);
    offset += 32;
    sk[offset..offset + 32].copy_from_slice(k_seed);
    offset += 32;
    sk[offset..offset + 64].copy_from_slice(tr);
    offset += 64;

    for i in 0..l {
        bit_pack(&s1[i], eta, eta, &mut sk[offset..offset + poly_eta_bytes]);
        offset += poly_eta_bytes;
    }
    for i in 0..k {
        bit_pack(&s2[i], eta, eta, &mut sk[offset..offset + poly_eta_bytes]);
        offset += poly_eta_bytes;
    }
    for i in 0..k {
        // t0 coefficients are in [-(2^{d-1}-1), 2^{d-1}] = [-4095, 4096]
        // We use bit_pack with a=2^{d-1}-1=4095 and b=2^{d-1}=4096
        bit_pack(&t0[i], 4095, 4096, &mut sk[offset..offset + poly_t0_bytes]);
        offset += poly_t0_bytes;
    }

    sk
}

/// Decode a secret key from bytes.
///
/// Implements Algorithm 25 of FIPS 204 (skDecode). Inverse of [`sk_encode`].
///
/// Returns `(rho, K, tr, s1, s2, t0)` using fixed arrays.
pub fn sk_decode<P: Params>(
    sk: &[u8],
) -> (
    [u8; 32],
    [u8; 32],
    [u8; 64],
    [[i32; N]; MAX_L],
    [[i32; N]; MAX_K],
    [[i32; N]; MAX_K],
) {
    let eta = P::ETA as u32;
    let l = P::L;
    let k = P::K;
    let eta_bits = P::BITLEN_2ETA;
    let poly_eta_bytes = N * eta_bits / 8;
    let d = D;
    let poly_t0_bytes = N * d / 8;

    let mut offset = 0;
    let mut rho = [0u8; 32];
    rho.copy_from_slice(&sk[offset..offset + 32]);
    offset += 32;
    let mut k_seed = [0u8; 32];
    k_seed.copy_from_slice(&sk[offset..offset + 32]);
    offset += 32;
    let mut tr = [0u8; 64];
    tr.copy_from_slice(&sk[offset..offset + 64]);
    offset += 64;

    let mut s1 = [[0i32; N]; MAX_L];
    for i in 0..l {
        bit_unpack(&sk[offset..offset + poly_eta_bytes], eta, eta, &mut s1[i]);
        offset += poly_eta_bytes;
    }
    let mut s2 = [[0i32; N]; MAX_K];
    for i in 0..k {
        bit_unpack(&sk[offset..offset + poly_eta_bytes], eta, eta, &mut s2[i]);
        offset += poly_eta_bytes;
    }
    let mut t0 = [[0i32; N]; MAX_K];
    for i in 0..k {
        bit_unpack(&sk[offset..offset + poly_t0_bytes], 4095, 4096, &mut t0[i]);
        offset += poly_t0_bytes;
    }

    (rho, k_seed, tr, s1, s2, t0)
}

/// Decode only the seeds (rho, K, tr) from a secret key, without
/// unpacking any polynomial. Returns 128 bytes of stack.
pub fn sk_decode_seeds<P: Params>(sk: &[u8]) -> ([u8; 32], [u8; 32], [u8; 64]) {
    let mut rho = [0u8; 32];
    rho.copy_from_slice(&sk[..32]);
    let mut k_seed = [0u8; 32];
    k_seed.copy_from_slice(&sk[32..64]);
    let mut tr = [0u8; 64];
    tr.copy_from_slice(&sk[64..128]);
    (rho, k_seed, tr)
}

/// Decode a single polynomial of s1 from the packed secret key.
///
/// `idx` must be in `0..P::L`.
pub fn sk_decode_s1<P: Params>(sk: &[u8], idx: usize, out: &mut [i32; N]) {
    let eta = P::ETA as u32;
    let eta_bits = P::BITLEN_2ETA;
    let poly_eta_bytes = N * eta_bits / 8;
    let base = 128; // rho(32) + K(32) + tr(64)
    let offset = base + idx * poly_eta_bytes;
    bit_unpack(&sk[offset..offset + poly_eta_bytes], eta, eta, out);
}

/// Decode a single polynomial of s2 from the packed secret key.
///
/// `idx` must be in `0..P::K`.
pub fn sk_decode_s2<P: Params>(sk: &[u8], idx: usize, out: &mut [i32; N]) {
    let eta = P::ETA as u32;
    let l = P::L;
    let eta_bits = P::BITLEN_2ETA;
    let poly_eta_bytes = N * eta_bits / 8;
    let base = 128 + l * poly_eta_bytes;
    let offset = base + idx * poly_eta_bytes;
    bit_unpack(&sk[offset..offset + poly_eta_bytes], eta, eta, out);
}

/// Decode a single polynomial of t0 from the packed secret key.
///
/// `idx` must be in `0..P::K`.
pub fn sk_decode_t0<P: Params>(sk: &[u8], idx: usize, out: &mut [i32; N]) {
    let l = P::L;
    let k = P::K;
    let eta_bits = P::BITLEN_2ETA;
    let poly_eta_bytes = N * eta_bits / 8;
    let d = D;
    let poly_t0_bytes = N * d / 8;
    let base = 128 + (l + k) * poly_eta_bytes;
    let offset = base + idx * poly_t0_bytes;
    bit_unpack(&sk[offset..offset + poly_t0_bytes], 4095, 4096, out);
}

/// Encode a signature as bytes.
///
/// Implements Algorithm 26 of FIPS 204 (sigEncode). The signature is the
/// concatenation of `c_tilde` (lambda/4 bytes), the l polynomials of `z`
/// (each packed with `1 + bitlen(gamma1-1)` bits per coefficient), and the
/// hint vector `h` (packed via [`hint_bit_pack`]).
///
/// - `c_tilde`: commitment hash (lambda/4 bytes).
/// - `z`: response vector (l polynomials with coefficients in [-(gamma1-1), gamma1]).
/// - `h`: hint vector (k binary polynomials).
///
/// Returns a byte vector of length `P::SIG_LEN`.
pub fn sig_encode<P: Params>(c_tilde: &[u8], z: &[[i32; N]], h: &[[i32; N]]) -> Vec<u8> {
    let l = P::L;
    let gamma1 = P::GAMMA1 as u32;
    let gamma1_bits = P::BITLEN_GAMMA1_MINUS1 + 1; // bitlen(gamma1-1)+1 = 1+bitlen(gamma1-1)
    let poly_z_bytes = N * gamma1_bits / 8;
    let c_tilde_len = P::LAMBDA / 4;

    let mut sig = vec![0u8; P::SIG_LEN];
    let mut offset = 0;

    sig[offset..offset + c_tilde_len].copy_from_slice(&c_tilde[..c_tilde_len]);
    offset += c_tilde_len;

    for i in 0..l {
        // z coefficients are in [-(gamma1-1), gamma1], so a=gamma1-1, b=gamma1
        bit_pack(&z[i], gamma1 - 1, gamma1, &mut sig[offset..offset + poly_z_bytes]);
        offset += poly_z_bytes;
    }

    hint_bit_pack::<P>(h, &mut sig[offset..]);

    sig
}

/// Decode a signature from bytes.
///
/// Implements Algorithm 27 of FIPS 204 (sigDecode). Inverse of [`sig_encode`].
///
/// Returns `Some((c_tilde, z, h))` on success, or `None` if the hint
/// encoding is malformed. Uses fixed arrays instead of Vec.
pub fn sig_decode<P: Params>(sig: &[u8]) -> Option<(Vec<u8>, [[i32; N]; MAX_L], [[i32; N]; MAX_K])> {
    let l = P::L;
    let gamma1 = P::GAMMA1 as u32;
    let gamma1_bits = P::BITLEN_GAMMA1_MINUS1 + 1;
    let poly_z_bytes = N * gamma1_bits / 8;
    let c_tilde_len = P::LAMBDA / 4;

    let mut offset = 0;
    let c_tilde = sig[offset..offset + c_tilde_len].to_vec();
    offset += c_tilde_len;

    let mut z = [[0i32; N]; MAX_L];
    for i in 0..l {
        bit_unpack(&sig[offset..offset + poly_z_bytes], gamma1 - 1, gamma1, &mut z[i]);
        offset += poly_z_bytes;
    }

    let h = hint_bit_unpack::<P>(&sig[offset..])?;

    Some((c_tilde, z, h))
}

/// Encode the high-order bits vector w1 as bytes.
///
/// Implements Algorithm 28 of FIPS 204 (w1Encode). The w1 coefficients lie
/// in [0, (q-1)/(2*gamma2) - 1] and are packed using [`simple_bit_pack`].
/// The encoded output is hashed together with the message digest to form
/// the commitment hash c_tilde during signing and verification.
///
/// - `w1`: vector of k polynomials (high bits from Decompose).
///
/// Returns the packed byte representation of w1.
pub fn w1_encode<P: Params>(w1: &[[i32; N]]) -> Vec<u8> {
    let k = P::K;
    let gamma2 = P::GAMMA2;
    // w1 coefficients are in [0, (q-1)/(2*gamma2)]
    let max_w1 = ((Q - 1) / (2 * gamma2) - 1) as u32;
    let bits = 32 - max_w1.leading_zeros() as usize;
    let poly_bytes = N * bits / 8;
    let mut out = vec![0u8; k * poly_bytes];
    for i in 0..k {
        let offset = i * poly_bytes;
        simple_bit_pack(&w1[i], max_w1, &mut out[offset..offset + poly_bytes]);
    }
    out
}

/// Decompose a coefficient into high and low parts using a power-of-2 divisor.
///
/// Implements Algorithm 35 of FIPS 204 (Power2Round). Splits `r` into
/// `(r1, r0)` such that `r = r1 * 2^d + r0` with `r0` in the centered
/// range `[-(2^{d-1} - 1), 2^{d-1}]`.
///
/// Used during key generation to compress the public vector t into t1
/// (stored in the public key) and t0 (stored in the secret key).
pub fn power2round(r: i32) -> (i32, i32) {
    let rp = mod_q(r);
    // r0 = r mod+ 2^d (centered modular reduction)
    let two_d = 1i32 << D;
    let half = two_d >> 1;
    let mut r0v = rp % two_d; // [0, 2^d - 1]
    if r0v > half {
        r0v -= two_d;
    }
    let r1 = (rp - r0v) / two_d;
    (r1, r0v)
}

/// Apply [`power2round`] to every coefficient of a polynomial vector.
///
/// Returns `(t1, t0)` as fixed arrays where each element satisfies
/// `t[i][j] = t1[i][j] * 2^d + t0[i][j]`.
pub fn power2round_vec(t: &[[i32; N]], len: usize) -> ([[i32; N]; MAX_K], [[i32; N]; MAX_K]) {
    let mut t1 = [[0i32; N]; MAX_K];
    let mut t0 = [[0i32; N]; MAX_K];
    for i in 0..len {
        for j in 0..N {
            let (r1, r0) = power2round(t[i][j]);
            t1[i][j] = r1;
            t0[i][j] = r0;
        }
    }
    (t1, t0)
}

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

    #[test]
    fn test_simple_bit_pack_unpack() {
        let mut w = [0i32; N];
        for i in 0..N {
            w[i] = (i as i32 * 3) % 1024;
        }
        let mut buf = [0u8; 320]; // 256 * 10 / 8
        simple_bit_pack(&w, 1023, &mut buf);
        let mut w2 = [0i32; N];
        simple_bit_unpack(&buf, 1023, &mut w2);
        assert_eq!(w, w2);
    }

    #[test]
    fn test_bit_pack_unpack() {
        let mut w = [0i32; N];
        for i in 0..N {
            w[i] = (i as i32 % 5) - 2; // values in [-2, 2]
        }
        // a=2, b=2, a+b=4, bitlen(4)=3, 256*3/8=96
        let mut buf = [0u8; 96];
        bit_pack(&w, 2, 2, &mut buf);
        let mut w2 = [0i32; N];
        bit_unpack(&buf, 2, 2, &mut w2);
        assert_eq!(w, w2);
    }

    #[test]
    fn test_power2round() {
        let (r1, r0) = power2round(1234567);
        assert_eq!(r1 * (1 << D) + r0, 1234567);
        assert!(r0.abs() <= (1 << (D - 1)));
    }
}