rind 0.1.0

// Written for Graviola by Joe Birr-Pixton, 2026.
// SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT-0

//! ML-KEM-768 as standardised by FIPS-203.

use core::fmt;
use core::marker::PhantomData;

use crate::error::{Error, KeyFormatError};
use crate::low;
use crate::{
    low::{ct_copy, ct_equal, ct_select_i16},
    mid::{
        rng::{RandomSource, SystemRandom},
        sha3,
    },
};

mod sample_ntt;

/// An ML-KEM-768 decapsulation key.
pub(crate) struct DecapKey {
    ek_pke: EncapKey,
    dk_pke: [u8; K * 384],
    h_ek: [u8; 32],
    z: [u8; 32],
}

impl DecapKey {
    /// Generate a random [`DecapKey`] (which contains the corresponding [`EncapKey`]).
    ///
    /// This fails only if random material generation fails.
    pub(crate) fn generate() -> Result<Self, Error> {
        let _entry = low::Entry::new_secret();
        let mut seed = [0u8; 64];
        SystemRandom.fill(&mut seed)?;
        let key = Self::keygen_internal(&seed);
        low::zeroise(&mut seed);
        Ok(key)
    }

    /// Decapsulate ciphertext `c`, yielding a [`SharedSecret`].
    pub(crate) fn decaps(self, c: &Ciphertext) -> SharedSecret {
        let _entry = low::Entry::new_secret();
        self.decaps_internal(c)
    }

    /// Return the corresponding [`EncapKey`].
    pub(crate) fn encapsulation_key(&self) -> EncapKey {
        let _entry = low::Entry::new_public();
        self.ek_pke.clone()
    }

    /// Encode the key as bytes.
    ///
    /// This is used for testing.
    #[cfg(test)]
    pub(crate) fn as_bytes(&self) -> [u8; K * 768 + 96] {
        let _entry = low::Entry::new_secret();
        let mut out = [0u8; K * 768 + 96];
        out[..K * 384].copy_from_slice(&self.dk_pke);
        out[K * 384..K * 768].copy_from_slice(&self.ek_pke.t_hat);
        out[K * 768..K * 768 + 32].copy_from_slice(&self.ek_pke.rho);
        out[K * 768 + 32..K * 768 + 64].copy_from_slice(&self.h_ek);
        out[K * 768 + 64..].copy_from_slice(&self.z);
        out
    }

    pub(crate) fn keygen_internal(seed: &[u8; 64]) -> Self {
        // Preliminaries: split seed into d and z.
        let (d, z) = seed.split_at(32);
        let mut d = low::ct::into_secret(d.try_into().unwrap());

        // 1. (ekPKE, dkPKE) <- K-PKE.KeyGen(d)
        // 2. ek <- ekPKE
        let (ek_pke, dk_pke) = Self::kpke_keygen(&d);
        low::zeroise(&mut d);

        // 3. dk <- (dkPKE || ek || H(ek) || z)
        // 4. return (ek, dk)
        let mut h = sha3::Sha3_256Context::new();
        h.update(&ek_pke.t_hat);
        h.update(&ek_pke.rho);
        let h_ek = h.finish();

        Self {
            dk_pke,
            ek_pke,
            h_ek,
            z: low::ct::into_secret(z.try_into().unwrap()),
        }
    }

    fn decaps_internal(&self, c: &Ciphertext) -> SharedSecret {
        // Steps 1. - 4. are not relevant for us.
        // 5. m' <- K-PKE.Decrypt(dkPKE, c)
        let m_prime = self.kpke_decrypt(c);

        // 6. (K', r') <- G(m' || h)
        let mut g = sha3::Sha3_512Context::new();
        g.update(&m_prime.0);
        g.update(&self.h_ek);
        let mut kr_prime = g.finish();
        let (k_prime, r_prime) = kr_prime.split_at(32);
        let mut k_prime = SharedSecret(k_prime.try_into().unwrap());
        let r_prime = KpkeRandomness(r_prime.try_into().unwrap());
        low::zeroise(&mut kr_prime);

        // 7. K‾ <- J(z || c)
        let mut k_bar = SharedSecret([0u8; 32]);
        sha3::Shake256::new(&[&self.z, &c.0]).read(&mut k_bar.0);

        // 8. c' <- K-PKE.Encrypt(ekPKE, m', r')
        let mut c_prime = self.ek_pke.kpke_encrypt(m_prime, r_prime);

        // 9. - 11. if c != c' then K' <- K‾
        // 12. return K'
        ct_copy(ct_equal(&c.0, &c_prime.0), &mut k_prime.0, &k_bar.0);
        // A c' that differs from c re-encrypts m', the decryption of a
        // ciphertext the peer chose.
        low::zeroise(&mut c_prime.0);
        k_prime
    }

    /// This is Algorithm 13: K-PKE.KeyGen(d) -> (ek_PKE, dk_PKE)
    fn kpke_keygen(d: &[u8; 32]) -> (EncapKey, [u8; K * 384]) {
        // 1. (𝜌, 𝜎) <- G(d || k)
        let mut g = sha3::Sha3_512Context::new();
        g.update(d);
        g.update(&[K_BYTE]);
        let mut g = g.finish();
        let (rho, sigma) = g.split_at(32);
        // rho seeds the public matrix, so it is public.
        let rho: [u8; 32] = low::ct::into_public(rho.try_into().unwrap());
        let sigma = KpkeRandomness(sigma.try_into().unwrap());
        low::zeroise(&mut g);

        // 3. - 7.
        let a_hat = Coeffs::sample_poly(&rho);

        // 8. - 15.
        let (s, e) = Coeffs::sample_poly_cbd_dual(&sigma.0);

        // 16. ŝ <- NTT(s)
        let mut s_hat = s.ntt();

        // 17. ê <- NTT(e)
        let e_hat = e.ntt();

        // 18. t^ = A^ ∘ ŝ + ê
        let t_hat = a_hat.mul_add(&s_hat, &e_hat);

        let encap = EncapKey {
            t_hat: low::ct::into_public(t_hat.to_bytes()),
            rho,
            transpose_a_hat: a_hat.transpose(),
        };

        s_hat.reduce_in_place();
        let decap = s_hat.to_bytes();

        (encap, decap)
    }

    /// This is Algorithm 15: K-PKE.Decrypt(dk_PKE, c)
    fn kpke_decrypt(&self, c: &Ciphertext) -> Message {
        // 1. c_1 <- c[0 ∶ 32 * du * k]
        // 2. c_2 <- c[32 * du * k ∶ 32 * (du * k + dv)]
        let (c_1, c_2) = c.0.split_at(32 * DU * K);

        // 3. u' <- Decompress_du(ByteDecode_du(c1))
        let u_prime = Coeffs::decompress_from_bytes_du_10(c_1.try_into().unwrap());
        // 4. v' <- Decompress_dv(ByteDecode_dv(c2))
        let v_prime = Coeffs::decompress_from_bytes_dv_4(c_2.try_into().unwrap());

        // 5. s^ <- ByteDecode_12(dk_PKE)
        let s_hat = Coeffs::from_bytes(&self.dk_pke);

        // 6. w <- v' - NTT-1(transpose(s^) * NTT(u'))
        let w = v_prime.sub(&s_hat.mul_and_inverse_ntt(&u_prime.ntt()));

        // 7. m <- ByteEncode_1(Compress_1(w))
        w.to_message_bits()
    }
}

// Wiped whole, so no byte of the key survives, padding included.
impl Drop for DecapKey {
    fn drop(&mut self) {
        low::zeroise_value(self);
    }
}

impl low::Zeroable for DecapKey {}

/// An ML-KEM-768 encapsulation key.
#[derive(Clone)]
pub(crate) struct EncapKey {
    t_hat: [u8; K * 384],
    rho: [u8; 32],
    transpose_a_hat: Coeffs<{ K * K * N }, Ntt>,
}

impl EncapKey {
    /// Create a new [`EncapKey`] from bytes.
    pub(crate) fn from_bytes(input: &[u8; K * 384 + 32]) -> Result<Self, Error> {
        let _entry = low::Entry::new_public();
        let (t_hat, rho) = input.split_at(K * 384);
        let t_hat = t_hat.try_into().unwrap();
        let rho = rho.try_into().unwrap();

        // 2. (Modulus check) Perform the computation
        //
        //   test <- ByteEncode_12(ByteDecode_12(ek[0 ∶ 384k]))
        //
        // If test != ek[0 ∶ 384k], then input checking failed.
        if t_hat != Coeffs::from_bytes(&t_hat).to_bytes() {
            return Err(KeyFormatError::MlKemModulus.into());
        }

        let transpose_a_hat = Coeffs::sample_poly_transposed(&rho);

        Ok(Self {
            t_hat,
            rho,
            transpose_a_hat,
        })
    }

    /// Encapsulate a random shared secret, returning the shared secret and a ciphertext.
    ///
    /// The ciphertext can be returned to the holder of the [`DecapKey`] who can then
    /// derive the same [`SharedSecret`].
    pub(crate) fn encaps(&self) -> Result<(SharedSecret, Ciphertext), Error> {
        let _entry = low::Entry::new_secret();
        let mut m = Message([0; 32]);
        SystemRandom.fill(&mut m.0)?;
        Ok(self.encaps_internal(m))
    }

    /// Encode the [`EncapKey`] as bytes.
    pub(crate) fn as_bytes(&self) -> [u8; K * 384 + 32] {
        let _entry = low::Entry::new_public();
        let mut out = [0u8; K * 384 + 32];
        out[..K * 384].copy_from_slice(&self.t_hat);
        out[K * 384..].copy_from_slice(&self.rho);
        out
    }

    pub(crate) fn encaps_internal(&self, m: Message) -> (SharedSecret, Ciphertext) {
        low::ct::secret_slice(&m.0);

        // 1. (K, r) <- G(m || H(ek))
        let mut h = sha3::Sha3_256Context::new();
        h.update(&self.t_hat);
        h.update(&self.rho);
        let h_ek = h.finish();

        let mut g = sha3::Sha3_512Context::new();
        g.update(&m.0);
        g.update(&h_ek);
        let mut kr = g.finish();
        let (k, r) = kr.split_at(32);
        let k = SharedSecret(k.try_into().unwrap());
        let r = KpkeRandomness(r.try_into().unwrap());
        low::zeroise(&mut kr);

        // 2. c <- K-PKE.Encrypt(ek, m, r)
        let c = self.kpke_encrypt(m, r);
        (k, Ciphertext(low::ct::into_public(c.0)))
    }

    fn kpke_encrypt(&self, m: Message, r: KpkeRandomness) -> Ciphertext {
        // 2. t^ <- ByteDecode12(ekPKE[0 ∶ 384k])
        let t_hat = Coeffs::from_bytes(&self.t_hat);

        // 3. 𝜌 <- ekPKE[384k ∶ 384k + 32]
        //  - only required for a_hat expansion
        // 4. - 8. rolled into self.transpose_a_hat.

        // 9. - 12.
        // 13. - 16.
        let (y, e_1) = Coeffs::sample_poly_cbd_dual(&r.0);

        // 17. e_2 <- SamplePolyCBD(PRF (r, N))
        let mut buf = [0u8; 128];
        sha3::Shake256::new(&[&r.0, &[K_BYTE * 2]]).read(&mut buf);
        let mut e_2 = Coeffs::zero();
        sample_cbd2(&buf, &mut e_2.0);
        low::zeroise(&mut buf);

        // 18. y^ <- NTT(y)
        let y_hat = y.ntt();

        // 19. u <- NTT^-1(transpose(A^) ∘ y^) + e_1
        let u = self.transpose_a_hat.mul_and_inverse_ntt(&y_hat).add(&e_1);

        // 20. µ <- Decompress_1(ByteDecode_1(m))
        let mu = Coeffs::from_message_bits(m);

        // 21. v <- NTT^-1(transpose(t^) ∘ y^) + e_2 + µ
        let v = t_hat.mul_and_inverse_ntt(&y_hat).add(&e_2).add(&mu);

        let mut c = Ciphertext([0; _]);
        let (c_1, c_2) = c.0.split_at_mut(32 * DU * K);

        // 22. c_1 <- ByteEncode_du(Compress_du(u))
        u.compress_into_bytes_du_10(c_1.try_into().unwrap());
        // 23. c_2 <- ByteEncode_dv(Compress_dv(v))
        v.compress_into_bytes_dv_4(c_2.try_into().unwrap());

        // 24. return c <- (c_1 || c_2)
        c
    }
}

/// An ML-KEM-768 ciphertext.
#[derive(Debug, Clone)]
pub(crate) struct Ciphertext([u8; 32 * (DU * K + DV)]);

impl From<[u8; 1088]> for Ciphertext {
    fn from(value: [u8; 1088]) -> Self {
        Self(value)
    }
}

impl AsRef<[u8; 1088]> for Ciphertext {
    fn as_ref(&self) -> &[u8; 1088] {
        &self.0
    }
}

/// A K-PKE message.
pub(crate) struct Message(pub(crate) [u8; 32]);

impl Drop for Message {
    fn drop(&mut self) {
        low::zeroise(&mut self.0);
    }
}

/// K-PKE randomness input: `sigma` for key generation, `r` for encryption.
struct KpkeRandomness([u8; 32]);

impl Drop for KpkeRandomness {
    fn drop(&mut self) {
        low::zeroise(&mut self.0);
    }
}

/// An ML-KEM shared secret.
pub(crate) struct SharedSecret([u8; 32]);

impl AsRef<[u8; 32]> for SharedSecret {
    fn as_ref(&self) -> &[u8; 32] {
        &self.0
    }
}

impl fmt::Debug for SharedSecret {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("SharedSecret").finish_non_exhaustive()
    }
}

impl Drop for SharedSecret {
    fn drop(&mut self) {
        low::zeroise(&mut self.0);
    }
}

/// Tracks what the values in a `Coeffs` mean.
trait Domain: Clone {}

/// Coefficients in normal domain.
#[derive(Clone)]
struct Normal;
impl Domain for Normal {}

/// Coefficients in NTT domain.
#[derive(Clone)]
struct Ntt;
impl Domain for Ntt {}

/// Multiplication cache values.
#[derive(Clone)]
struct MulCache;
impl Domain for MulCache {}

/// Polynomial coefficients.
///
/// Arithmetic on coefficients, here and in the compression and sampling
/// helpers, wraps: their bounds rule out overflow, and a debug build's
/// overflow check would branch on secret values.
#[derive(Clone, Debug)]
#[repr(align(32))]
struct Coeffs<const C: usize, D: Domain>([i16; C], PhantomData<D>);

// Most coefficients are secret: noise, keys, and the message and its encryption.
impl<const C: usize, D: Domain> Drop for Coeffs<C, D> {
    fn drop(&mut self) {
        low::zeroise(&mut self.0);
    }
}

impl<const C: usize, D: Domain> Coeffs<C, D> {
    const fn zero() -> Self {
        Self([0; C], PhantomData)
    }

    /// Reduce coefficients into [0, Q)
    fn reduce_in_place(&mut self) {
        for k in self.0.as_chunks_mut().0 {
            low::mlkem_reduce(k);
        }
    }
}

impl Coeffs<{ K * K * N }, Ntt> {
    fn sample_poly(rho: &[u8; 32]) -> Self {
        let mut r = Self(sample_ntt::sample_ntt(rho), PhantomData);
        r.reorder();
        r
    }

    fn sample_poly_transposed(rho: &[u8; 32]) -> Self {
        let mut r = Self(sample_ntt::sample_ntt_transposed(rho), PhantomData);
        r.reorder();
        r
    }

    // a * s + e
    fn mul_add(
        &self,
        s: &Coeffs<{ K * N }, Ntt>,
        e: &Coeffs<{ K * N }, Ntt>,
    ) -> Coeffs<{ K * N }, Ntt> {
        let mut r = e.clone();

        let s_precomp = s.mul_cache();

        let mut term = Coeffs::<_, Ntt>::zero();

        for (aa, rr) in self
            .0
            .as_chunks()
            .0
            .iter()
            .zip(r.0.as_chunks_mut::<N>().0.iter_mut())
        {
            low::mlkem_basemul_k3(&mut term.0, aa, &s.0, &s_precomp.0);
            low::mlkem_tomont(&mut term.0);

            for (r, t) in rr.iter_mut().zip(term.0.iter()) {
                *r = r.wrapping_add(*t);
            }
        }

        r.reduce_in_place();
        r
    }

    fn reorder(&mut self) {
        for ch in self.0.as_chunks_mut().0 {
            low::mlkem_unpack(ch);
        }
    }

    fn transpose(mut self) -> Self {
        let (chunks, _) = self.0.as_chunks_mut::<N>();

        for i in 0..K {
            for j in (i + 1)..K {
                chunks.swap(i * K + j, j * K + i);
            }
        }

        self
    }

    fn mul_and_inverse_ntt(&self, y: &Coeffs<{ K * N }, Ntt>) -> Coeffs<{ K * N }, Normal> {
        let mut r = Coeffs::<_, Ntt>::zero();

        let y_precomp = y.mul_cache();

        let mut term = Coeffs::<_, Ntt>::zero();

        for (aa, rr) in self
            .0
            .as_chunks()
            .0
            .iter()
            .zip(r.0.as_chunks_mut::<N>().0.iter_mut())
        {
            low::mlkem_basemul_k3(&mut term.0, aa, &y.0, &y_precomp.0);

            for (r, t) in rr.iter_mut().zip(term.0.iter()) {
                *r = r.wrapping_add(*t);
            }
        }

        r.reduce_in_place();
        r.inverse_ntt()
    }
}

impl Coeffs<{ K * N }, Ntt> {
    /// This is `ByteDecode_12(bytes)`.
    fn from_bytes(bytes: &[u8; 1152]) -> Self {
        let mut r = Coeffs::zero();

        let (r_chunks, _) = r.0.as_chunks_mut();
        let (b_chunks, _) = bytes.as_chunks();

        for (rr, bb) in r_chunks.iter_mut().zip(b_chunks.iter()) {
            low::mlkem_frombytes(rr, bb);
            low::mlkem_reduce(rr);
        }

        r
    }

    /// This is `ByteEncode_12(self)`
    fn to_bytes(&self) -> [u8; 1152] {
        let mut r = [0u8; 1152];

        let (c_chunks, _) = self.0.as_chunks();
        let (r_chunks, _) = r.as_chunks_mut();

        for (rr, cc) in r_chunks.iter_mut().zip(c_chunks.iter()) {
            low::mlkem_tobytes(rr, cc);
        }

        r
    }

    /// This is `NTT-1(self * s)`
    fn mul_and_inverse_ntt(&self, s: &Coeffs<{ K * N }, Ntt>) -> Coeffs<{ N }, Normal> {
        let mut r = Coeffs::zero();

        let s_precomp = s.mul_cache();

        low::mlkem_basemul_k3(&mut r.0, &self.0, &s.0, &s_precomp.0);
        // nb. no need for a tomont here, as `inverse_ntt` incorporates that

        r.inverse_ntt()
    }

    fn mul_cache(&self) -> Coeffs<{ K * 128 }, MulCache> {
        let mut r = Coeffs::zero();

        let (m, _) = self.0.as_chunks();
        let (rr, _) = r.0.as_chunks_mut();

        low::mlkem_mulcache_compute(&mut rr[0], &m[0]);
        low::mlkem_mulcache_compute(&mut rr[1], &m[1]);
        low::mlkem_mulcache_compute(&mut rr[2], &m[2]);

        r
    }
}

impl<const C: usize> Coeffs<C, Ntt> {
    fn inverse_ntt(mut self) -> Coeffs<C, Normal> {
        for ch in self.0.as_chunks_mut().0 {
            low::mlkem_intt(ch);
        }
        Coeffs(self.0, PhantomData)
    }
}

/// `Compress_4(x) = round((2^4 / q) * x) mod 2^4`, for x in [0, Q)
///
/// Division-free so no variable-time division instruction is emitted
/// (cf. KyberSlash); the constant is exact over the whole input range.
fn compress_4(x: i16) -> u8 {
    debug_assert!((0..Q).contains(&low::ct::into_public(x)));
    let c = x as u64;
    let c = c << 4;
    let c = c.wrapping_add(Q_HALF as u64);
    let c = c.wrapping_mul(80635);
    let c = c >> 28;
    let c = c & 0xf;
    c as u8
}

/// `Decompress_4(y) = round((q / 2^4) * y)`, for y in [0, 2^4)
fn decompress_4(y: u16) -> i16 {
    debug_assert!(y < 16);
    let d = y as u32;
    let d = d * (Q as u32);
    let d = d + 8;
    let d = d >> 4;
    d as i16
}

/// `Compress_10(x) = round((2^10 / q) * x) mod 2^10`, for x in [0, Q)
fn compress_10(x: i16) -> u16 {
    debug_assert!((0..Q).contains(&low::ct::into_public(x)));
    let c = x as u64;
    let c = c << 10;
    let c = c.wrapping_add(Q_HALF as u64);
    let c = c.wrapping_mul(1290167);
    let c = c >> 32;
    let c = c & 0x3ff;
    c as u16
}

/// `Decompress_10(y) = round((q / 2^10) * y)`, for y in [0, 2^10)
fn decompress_10(y: u16) -> i16 {
    debug_assert!(y < 1024);
    let d = y as u32;
    let d = d * (Q as u32);
    let d = d + 512;
    let d = d >> 10;
    d as i16
}

/// `Compress_1(x)` for eight coefficients in `[0, Q)`
fn compress_1_x8(coeffs: &[i16; 8]) -> u8 {
    let mut r = 0;

    for (i, x) in coeffs.iter().enumerate() {
        debug_assert!((0..Q).contains(&low::ct::into_public(*x)));
        let x = (*x as u64) << 1;
        let x = x.wrapping_add(Q_HALF as u64);
        let x = x.wrapping_mul(80635);
        let x = x >> 28;
        r |= ((x & 1) as u8) << i;
    }

    r
}

impl Coeffs<{ K * N }, Normal> {
    /// This is dual, K-wise `SamplePolyCBD`.
    ///
    /// In other words, it produces two noise vectors, diversifying the seed `sigma`
    /// with indices from 0..2K.
    fn sample_poly_cbd_dual(sigma: &[u8; 32]) -> (Self, Self) {
        // We need 2 * K samples.
        let mut samples = [[0; 128]; 6];

        // Prepare 6 SHAKE256 inputs.
        let mut bufs = [[0; 40]; 6];
        for (n, buf) in bufs.iter_mut().enumerate() {
            buf[..32].copy_from_slice(sigma);
            buf[32] = n as u8;
            buf[33] = sha3::SHAKE_PAD_BYTE;
        }

        sha3::Shake256::one_shot_sextet(&bufs, &mut samples);
        low::zeroise(bufs.as_flattened_mut());

        let mut a = Coeffs::zero();
        let mut b = Coeffs::zero();

        for (sample, poly) in samples.iter().zip(
            a.0.as_chunks_mut()
                .0
                .iter_mut()
                .chain(b.0.as_chunks_mut().0.iter_mut()),
        ) {
            sample_cbd2(sample, poly);
        }
        low::zeroise(samples.as_flattened_mut());

        (a, b)
    }

    /// K-wise NTT
    fn ntt(mut self) -> Coeffs<{ K * N }, Ntt> {
        for ch in self.0.as_chunks_mut().0 {
            low::mlkem_ntt(ch);
        }
        Coeffs(self.0, PhantomData)
    }

    /// This is `ByteEncode_10(Compress_10(u))`
    ///
    /// Each group of four coefficients becomes five bytes,
    /// little-endian bit order per `BitsToBytes`.
    fn compress_into_bytes_du_10(&self, c_1: &mut [u8; 960]) {
        let (quads, _) = self.0.as_chunks::<4>();
        let (out_chunks, _) = c_1.as_chunks_mut::<5>();

        for (quad, out) in quads.iter().zip(out_chunks.iter_mut()) {
            let t0 = compress_10(quad[0]);
            let t1 = compress_10(quad[1]);
            let t2 = compress_10(quad[2]);
            let t3 = compress_10(quad[3]);

            out[0] = t0 as u8;
            out[1] = ((t0 >> 8) | (t1 << 2)) as u8;
            out[2] = ((t1 >> 6) | (t2 << 4)) as u8;
            out[3] = ((t2 >> 4) | (t3 << 6)) as u8;
            out[4] = (t3 >> 2) as u8;
        }
    }

    /// This is `Decompress_10(ByteDecode_10(c_1))`
    ///
    /// The inverse of [`Self::compress_into_bytes_du_10`]: each group of
    /// five bytes yields four ten-bit values, little-endian bit order.
    fn decompress_from_bytes_du_10(c_1: &[u8; 960]) -> Self {
        let mut r = Self::zero();

        let (in_chunks, _) = c_1.as_chunks::<5>();
        let (quads, _) = r.0.as_chunks_mut::<4>();

        for (inp, quad) in in_chunks.iter().zip(quads.iter_mut()) {
            let t0 = (inp[0] as u16) | ((inp[1] as u16 & 0x3) << 8);
            let t1 = ((inp[1] as u16) >> 2) | ((inp[2] as u16 & 0xf) << 6);
            let t2 = ((inp[2] as u16) >> 4) | ((inp[3] as u16 & 0x3f) << 4);
            let t3 = ((inp[3] as u16) >> 6) | ((inp[4] as u16) << 2);

            quad[0] = decompress_10(t0);
            quad[1] = decompress_10(t1);
            quad[2] = decompress_10(t2);
            quad[3] = decompress_10(t3);
        }

        r
    }
}

impl Coeffs<N, Normal> {
    /// This is `Decompress_1(ByteDecode_1(m))`
    ///
    /// `ByteDecode_1` is just `BytesToBits`.
    /// `Decompress_1` maps set bits to Q/2.
    fn from_message_bits(m: Message) -> Self {
        let mut r = Self::zero();

        for (i, byte) in m.0.iter().enumerate() {
            for bit in 0..8 {
                r.0[i * 8 + bit] = ct_select_i16(byte >> bit & 1, Q_HALF, 0);
            }
        }

        r
    }

    /// This is `ByteEncode_1(Compress_1(m))`
    fn to_message_bits(&self) -> Message {
        let mut r = Message([0; 32]);

        for (i, r) in r.0.iter_mut().enumerate() {
            *r = compress_1_x8(&self.0[i * 8..(i + 1) * 8].try_into().unwrap());
        }

        r
    }

    /// This is `ByteEncode_4(Compress_4(v))`
    ///
    /// Each pair of coefficients becomes one byte, low nibble first.
    fn compress_into_bytes_dv_4(&self, c_2: &mut [u8; 128]) {
        let (pairs, _) = self.0.as_chunks::<2>();
        for (pair, byte) in pairs.iter().zip(c_2.iter_mut()) {
            *byte = compress_4(pair[0]) | (compress_4(pair[1]) << 4);
        }
    }

    /// This is `Decompress_4(ByteDecode_4(c_2))`
    ///
    /// The inverse of [`Self::compress_into_bytes_dv_4`]: each byte yields
    /// two four-bit values, low nibble first.
    fn decompress_from_bytes_dv_4(c_2: &[u8; 128]) -> Self {
        let mut r = Self::zero();

        let (pairs, _) = r.0.as_chunks_mut::<2>();
        for (byte, pair) in c_2.iter().zip(pairs.iter_mut()) {
            pair[0] = decompress_4((byte & 0xf) as u16);
            pair[1] = decompress_4((byte >> 4) as u16);
        }

        r
    }
}

impl<const C: usize> Coeffs<C, Normal> {
    /// Simple addition
    fn add(mut self, term: &Self) -> Self {
        for (ss, tt) in self.0.iter_mut().zip(term.0.iter()) {
            *ss = ss.wrapping_add(*tt);
        }
        self.reduce_in_place();
        self
    }

    /// Simple subtraction
    fn sub(mut self, term: &Self) -> Self {
        for (ss, tt) in self.0.iter_mut().zip(term.0.iter()) {
            *ss = ss.wrapping_sub(*tt);
        }
        self.reduce_in_place();
        self
    }
}

fn sample_cbd2(buf: &[u8; 128], out: &mut [i16; 256]) {
    let (buf_chunks, _) = buf.as_chunks::<4>();
    let (out_chunks, _) = out.as_chunks_mut::<8>();
    for (in_bytes, out_coeffs) in buf_chunks.iter().zip(out_chunks.iter_mut()) {
        let t = u32::from_le_bytes(*in_bytes);
        let d = (t & 0x5555_5555).wrapping_add((t >> 1) & 0x5555_5555);

        for (j, coeff) in out_coeffs.iter_mut().enumerate() {
            let a = ((d >> (4 * j)) & 0x3) as i16;
            let b = ((d >> (4 * j + 2)) & 0x3) as i16;
            *coeff = a.wrapping_sub(b);
        }
    }
}

const K: usize = 3;
const K_BYTE: u8 = K as u8;
const DU: usize = 10;
const DV: usize = 4;

const Q: i16 = 3329;
const Q_HALF: i16 = (Q + 1) / 2;
const N: usize = 256;

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

    #[test]
    fn shared_secret_debug() {
        assert_eq!(format!("{:?}", SharedSecret([0u8; 32])), "SharedSecret(..)");
    }

    #[test]
    fn pairwise() {
        let d = DecapKey::generate().unwrap();
        let (ess, ct) = d.encapsulation_key().encaps().unwrap();
        let dss = d.decaps(&ct);
        assert_eq!(low::ct::into_public(ess.0), low::ct::into_public(dss.0));
    }

    #[test]
    fn encaps_modulus_test() {
        assert_eq!(
            EncapKey::from_bytes(&[0xffu8; 1184]).err(),
            Some(KeyFormatError::MlKemModulus.into())
        );
    }
}