From b044f3971bedd8c03a1718e59b878001e844f546 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Mon, 2 Nov 2026 10:14:23 +0100 Subject: [PATCH 01/10] ml-kem: deterministic benchmarks for all parameter sets --- ml-kem/benches/mlkem.rs | 96 +++++++++++++++++++++++++---------------- 1 file changed, 60 insertions(+), 36 deletions(-) diff --git a/ml-kem/benches/mlkem.rs b/ml-kem/benches/mlkem.rs index 740d857b..11e0e574 100644 --- a/ml-kem/benches/mlkem.rs +++ b/ml-kem/benches/mlkem.rs @@ -1,54 +1,78 @@ //! ML-KEM benchmarks. +//! +//! All inputs are fixed byte strings so that no RNG or system call runs inside a timed region, +//! and every output is passed through `black_box` so that no work can be optimized away. #![allow(missing_docs, clippy::unwrap_used)] -use ::kem::{Decapsulate, Encapsulate, Kem, KeyExport, KeyInit}; +use ::kem::{Decapsulate, KeyExport, KeyInit}; use core::hint::black_box; use criterion::{Criterion, criterion_group, criterion_main}; -use getrandom::SysRng; use ml_kem::*; -use rand_core::UnwrapErr; -fn criterion_benchmark(c: &mut Criterion) { - let mut rng = UnwrapErr(SysRng); - - // Key generation - c.bench_function("keygen", |b| { - b.iter(|| { - let (dk, ek) = MlKem768::generate_keypair_from_rng(&mut rng); - let _dk_bytes = black_box(dk.to_seed().unwrap()); - let _ek_bytes = black_box(ek.to_bytes()); +/// Benchmark every public operation of one parameter set. +macro_rules! bench_parameter_set { + ($c:expr, $name:literal, $module:ident) => {{ + type Dk = $module::DecapsulationKey; + type Ek = $module::EncapsulationKey; + + let seed = Seed::from([0x5a; 64]); + let m = B32::from([0xa5; 32]); + + let dk = Dk::new(&seed); + let ek_bytes = dk.encapsulation_key().to_bytes(); + let ek = Ek::new(&ek_bytes).unwrap(); + let (ct, _) = ek.encapsulate_deterministic(&m); + let mut bad_ct = ct.clone(); + bad_ct[0] ^= 1; + + let mut group = $c.benchmark_group($name); + + // Algorithm 19 (deterministic part): expand a seed into a decapsulation key. + group.bench_function("keygen", |b| { + b.iter(|| Dk::new(black_box(&seed))); }); - }); - let (dk, ek) = MlKem768::generate_keypair_from_rng(&mut rng); - let dk_bytes = dk.to_seed().unwrap(); - let ek_bytes = ek.to_bytes(); - let ek = ::EncapsulationKey::new(&ek_bytes).unwrap(); + // FIPS 203 §7.2 input checks plus H(ek). + group.bench_function("ek_parse", |b| { + b.iter(|| Ek::new(black_box(&ek_bytes)).unwrap()); + }); - // Encapsulation - c.bench_function("encapsulate", |b| { - b.iter(|| ek.encapsulate_with_rng(&mut rng)); - }); - let (ct, _sk) = ek.encapsulate_with_rng(&mut rng); + // Algorithm 17 with a fixed message. + group.bench_function("encapsulate", |b| { + b.iter(|| ek.encapsulate_deterministic(black_box(&m))); + }); - // Decapsulation - let dk = ::DecapsulationKey::new(&dk_bytes); + // Algorithm 18 on a valid ciphertext. + group.bench_function("decapsulate", |b| { + b.iter(|| dk.decapsulate(black_box(&ct))); + }); - c.bench_function("decapsulate", |b| { - b.iter(|| { - dk.decapsulate(&ct); + // Algorithm 18 on an invalid ciphertext (implicit rejection). Must cost the same as the + // valid case. + group.bench_function("decapsulate_reject", |b| { + b.iter(|| dk.decapsulate(black_box(&bad_ct))); }); - }); - - // Round trip - c.bench_function("round_trip", |b| { - b.iter(|| { - let (dk, ek) = MlKem768::generate_keypair_from_rng(&mut rng); - let (ct, _sk) = ek.encapsulate_with_rng(&mut rng); - dk.decapsulate(&ct); + + // One ephemeral exchange: key generation, encapsulation key parsing, encapsulation and + // decapsulation. + group.bench_function("round_trip", |b| { + b.iter(|| { + let dk = Dk::new(black_box(&seed)); + let ek = Ek::new(&dk.encapsulation_key().to_bytes()).unwrap(); + let (ct, k_send) = ek.encapsulate_deterministic(black_box(&m)); + (k_send, dk.decapsulate(&ct)) + }); }); - }); + + group.finish(); + }}; +} + +fn criterion_benchmark(c: &mut Criterion) { + bench_parameter_set!(c, "ML-KEM-512", ml_kem_512); + bench_parameter_set!(c, "ML-KEM-768", ml_kem_768); + bench_parameter_set!(c, "ML-KEM-1024", ml_kem_1024); } criterion_group!(benches, criterion_benchmark); From 4e1afa257faa4ce117d60ed74eb8684326c56e20 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Mon, 2 Nov 2026 17:32:51 +0100 Subject: [PATCH 02/10] ml-kem: add accumulated test vectors --- ml-kem/tests/accumulated.rs | 128 ++++++++++++++++++++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 ml-kem/tests/accumulated.rs diff --git a/ml-kem/tests/accumulated.rs b/ml-kem/tests/accumulated.rs new file mode 100644 index 00000000..68aecd55 --- /dev/null +++ b/ml-kem/tests/accumulated.rs @@ -0,0 +1,128 @@ +//! Accumulated test vectors, using the procedure of the C2SP CCTV "accumulated" vectors: +//! +//! +//! A SHAKE-128 instance with empty input provides `d`, `z`, `m` and a random ciphertext for each +//! test. The encapsulation key, the expanded decapsulation key, the ciphertext, the shared key from +//! encapsulation and the shared key from decapsulating the random ciphertext are absorbed into a +//! second SHAKE-128 instance, whose first 32 bytes of output are compared to the expected value. +//! +//! The hashes published by CCTV target the FIPS 203 initial public draft, whose K-PKE.KeyGen +//! derives its seeds as `G(d)` rather than `G(d || k)`, so they cannot match the final standard. +//! The expected values below were computed for FIPS 203 with libcrux-ml-kem 0.0.11. ml-kem 0.3.4 +//! matches the 100 and 10,000 test runs. +//! +//! The 100-test run is part of the normal test suite. The longer runs exercise the rare paths of +//! `SampleNTT` that need extra XOF blocks and are ignored by default: +//! +//! ```text +//! cargo test --release --test accumulated -- --ignored +//! ``` + +#![allow(clippy::unwrap_used, missing_docs)] + +use hex_literal::hex; +#[allow(deprecated)] +use ml_kem::ExpandedKeyEncoding; +use ml_kem::{ + B32, Seed, + kem::{Decapsulate, KeyExport, KeyInit}, +}; +use shake::{ExtendableOutput, Shake128, Update, XofReader}; + +macro_rules! accumulated_test { + ($name:ident, $module:ident, $count:expr, $expected:expr $(, $ignore:literal)?) => { + #[test] + $(#[ignore = $ignore])? + #[allow(deprecated)] + fn $name() { + let mut rng = Shake128::default().finalize_xof(); + let mut accumulator = Shake128::default(); + + for _ in 0..$count { + // `d || z` + let mut seed = Seed::default(); + rng.read(&mut seed); + let dk = ml_kem::$module::DecapsulationKey::new(&seed); + let ek = dk.encapsulation_key(); + accumulator.update(&ek.to_bytes()); + accumulator.update(&dk.to_expanded_bytes()); + + let mut m = B32::default(); + rng.read(&mut m); + let (ct, k) = ek.encapsulate_deterministic(&m); + assert_eq!(dk.decapsulate(&ct), k); + accumulator.update(&ct); + accumulator.update(&k); + + let mut random_ct = ml_kem::$module::Ciphertext::default(); + rng.read(&mut random_ct); + accumulator.update(&dk.decapsulate(&random_ct)); + } + + let mut digest = [0u8; 32]; + accumulator.finalize_xof().read(&mut digest); + assert_eq!(digest, $expected); + } + }; +} + +accumulated_test!( + accumulated_512_100, + ml_kem_512, + 100, + hex!("449120c6e320ef3e9fbfa2316e5f2d2e1e6dd37d8ff5d086d5d2db7d42aff0a1") +); +accumulated_test!( + accumulated_768_100, + ml_kem_768, + 100, + hex!("8d65b902f28edc683cebee2872962fd165a4d197c9e24ec74caa4470270df0b7") +); +accumulated_test!( + accumulated_1024_100, + ml_kem_1024, + 100, + hex!("c3ffe9ebecfa479c142656cbfbc6417efa05b77e994fe538eef4daed166363df") +); +accumulated_test!( + accumulated_512_10k, + ml_kem_512, + 10_000, + hex!("705dcffc87f4e67e35a09dcaa31772e86f3341bd3ccf1e78a5fef99ae6a35a13"), + "long-running" +); +accumulated_test!( + accumulated_768_10k, + ml_kem_768, + 10_000, + hex!("f959d18d3d1180121433bf0e05f11e7908cf9d03edc150b2b07cb90bef5bc1c1"), + "long-running" +); +accumulated_test!( + accumulated_1024_10k, + ml_kem_1024, + 10_000, + hex!("e3bf82b013307b2e9d47dde791ff6dfc82e694e6382404abdb948b908b75bad5"), + "long-running" +); +accumulated_test!( + accumulated_512_1m, + ml_kem_512, + 1_000_000, + hex!("21dd330d4355f2ae2876b9fa2b9de62ecaf76aca1d598de8db2b467d36e36a6a"), + "long-running" +); +accumulated_test!( + accumulated_768_1m, + ml_kem_768, + 1_000_000, + hex!("3b108396a277f2952ff3243a985c9709bcb95788c39b7b36a2c4e19d1a41e51e"), + "long-running" +); +accumulated_test!( + accumulated_1024_1m, + ml_kem_1024, + 1_000_000, + hex!("6377c4f0ecfdb32e63f7b58227960828784fe0b3e0e5e5e9f77be300f003512a"), + "long-running" +); From 5b7893cd6c1f714669454fcf12a388fb5c9b48ea Mon Sep 17 00:00:00 2001 From: ccipherx Date: Tue, 3 Nov 2026 14:05:12 +0100 Subject: [PATCH 03/10] ml-kem: montgomery ntt kernels with lazy reduction --- ml-kem/src/algebra.rs | 234 +++++++++--------- ml-kem/src/algebra/montgomery.rs | 396 +++++++++++++++++++++++++++++++ ml-kem/src/pke.rs | 13 +- 3 files changed, 531 insertions(+), 112 deletions(-) create mode 100644 ml-kem/src/algebra/montgomery.rs diff --git a/ml-kem/src/algebra.rs b/ml-kem/src/algebra.rs index 41285f3c..c2ea740b 100644 --- a/ml-kem/src/algebra.rs +++ b/ml-kem/src/algebra.rs @@ -3,10 +3,15 @@ use crate::{ crypto::{PRF, PrfOutput, XOF}, param::CbdSamplingSize, }; -use array::{Array, ArraySize, typenum::U256}; +use array::{Array, ArraySize, typenum::U1}; use module_lattice::{Encode, Field, MultiplyNtt, Truncate}; use sha3::digest::XofReader; +/// Section 4.3: arithmetic kernels for the NTT and for products in the NTT domain. +mod montgomery; + +pub(crate) use montgomery::{PreparedNtt, matrix_multiply, multiply_accumulate, prepare}; + module_lattice::define_field!(BaseField, u16, u32, u64, 3329); pub(crate) type Int = ::Int; @@ -146,44 +151,12 @@ pub(crate) trait Ntt { fn ntt(&self) -> Self::Output; } -/// One layer of the forward NTT butterfly. -/// -/// `LEN` is the butterfly half-length and `ITERATIONS = 128 / LEN` is the number of -/// butterfly groups in the layer. Making both compile-time constants lets the compiler -/// eliminate the iterator length calculation (`256 / (2 * LEN)`) that `step_by` would -/// otherwise compute with a `UDIV` instruction. -#[inline(always)] -fn ntt_layer(f: &mut Array, k: &mut usize) { - for i in 0..ITERATIONS { - let start = i * 2 * LEN; - let zeta = ZETA_POW_BITREV[*k]; - *k += 1; - - for j in start..(start + LEN) { - let t = zeta * f[j + LEN]; - f[j + LEN] = f[j] - t; - f[j] = f[j] + t; - } - } -} - /// Algorithm 9: `NTT` impl Ntt for Polynomial { type Output = NttPolynomial; fn ntt(&self) -> NttPolynomial { - let mut k = 1; - let mut f = self.0; - - ntt_layer::<128, 1>(&mut f, &mut k); - ntt_layer::<64, 2>(&mut f, &mut k); - ntt_layer::<32, 4>(&mut f, &mut k); - ntt_layer::<16, 8>(&mut f, &mut k); - ntt_layer::<8, 16>(&mut f, &mut k); - ntt_layer::<4, 32>(&mut f, &mut k); - ntt_layer::<2, 64>(&mut f, &mut k); - - f.into() + montgomery::ntt(&self.0) } } @@ -191,7 +164,7 @@ impl Ntt for Vector { type Output = NttVector; fn ntt(&self) -> NttVector { - NttVector::new(self.0.iter().map(Ntt::ntt).collect()) + NttVector::new(Array::from_fn(|i| self.0[i].ntt())) } } @@ -203,44 +176,12 @@ pub(crate) trait NttInverse { fn ntt_inverse(&self) -> Self::Output; } -/// One layer of the inverse NTT butterfly. -/// -/// See [`ntt_layer`] for the rationale behind the const generics. -#[inline(always)] -fn ntt_inverse_layer( - f: &mut Array, - k: &mut usize, -) { - for i in 0..ITERATIONS { - let start = i * 2 * LEN; - let zeta = ZETA_POW_BITREV[*k]; - *k -= 1; - - for j in start..(start + LEN) { - let t = f[j]; - f[j] = t + f[j + LEN]; - f[j + LEN] = zeta * (f[j + LEN] - t); - } - } -} - /// Algorithm 10: `NTT^{-1}` impl NttInverse for NttPolynomial { type Output = Polynomial; fn ntt_inverse(&self) -> Polynomial { - let mut f: Array = self.0.clone(); - let mut k = 127; - - ntt_inverse_layer::<2, 64>(&mut f, &mut k); - ntt_inverse_layer::<4, 32>(&mut f, &mut k); - ntt_inverse_layer::<8, 16>(&mut f, &mut k); - ntt_inverse_layer::<16, 8>(&mut f, &mut k); - ntt_inverse_layer::<32, 4>(&mut f, &mut k); - ntt_inverse_layer::<64, 2>(&mut f, &mut k); - ntt_inverse_layer::<128, 1>(&mut f, &mut k); - - Elem::new(3303) * &Polynomial::new(f) + Polynomial::new(montgomery::ntt_inverse(&self.0)) } } @@ -248,51 +189,18 @@ impl NttInverse for NttVector { type Output = Vector; fn ntt_inverse(&self) -> Vector { - Vector::new(self.0.iter().map(NttInverse::ntt_inverse).collect()) + Vector::new(Array::from_fn(|i| self.0[i].ntt_inverse())) } } /// Algorithm 11: `MultiplyNTTs` impl MultiplyNtt for BaseField { fn multiply_ntt(lhs: &NttPolynomial, rhs: &NttPolynomial) -> NttPolynomial { - let mut out = NttPolynomial::new(Array::default()); - - for i in 0..128 { - let (c0, c1) = base_case_multiply( - lhs.0[2 * i], - lhs.0[2 * i + 1], - rhs.0[2 * i], - rhs.0[2 * i + 1], - i, - ); - - out.0[2 * i] = c0; - out.0[2 * i + 1] = c1; - } - - out + let lhs = NttVector::::new(Array([lhs.clone()])); + multiply_accumulate(&lhs, &Array([PreparedNtt::new(rhs)])) } } -/// Algorithm 12: `BaseCaseMultiply` -/// -/// This is a hot loop. We promote to u64 so that we can do the absolute minimum number of -/// modular reductions, since these are the expensive operation. -#[inline] -fn base_case_multiply(a0: Elem, a1: Elem, b0: Elem, b1: Elem, i: usize) -> (Elem, Elem) { - let a0 = u32::from(a0.0); - let a1 = u32::from(a1.0); - let b0 = u32::from(b0.0); - let b1 = u32::from(b1.0); - let g = u32::from(GAMMA[i].0); - - let b1g = u32::from(BaseField::barrett_reduce(b1 * g)); - - let c0 = BaseField::barrett_reduce(a0 * b0 + a1 * b1g); - let c1 = BaseField::barrett_reduce(a0 * b1 + a1 * b0); - (Elem::new(c0), Elem::new(c1)) -} - /// Since the powers of zeta used in the `NTT` and `MultiplyNTTs` are fixed, we use pre-computed /// tables to avoid the need to compute the exponentiations at runtime. /// @@ -358,11 +266,11 @@ const GAMMA: [Elem; 128] = { mod test { use super::{ Array, B32, BaseField, Elem, Field, Int, Ntt, NttInverse, NttMatrix, NttPolynomial, - NttVector, PRF, Polynomial, U256, XOF, + NttVector, PRF, Polynomial, U1, XOF, }; use array::{ ArraySize, Flatten, - typenum::{U2, U3, U8}, + typenum::{U2, U3, U8, U256}, }; /// A polynomial with only a scalar component, to make simple test cases @@ -396,6 +304,120 @@ mod test { })) } + /// The previous `Elem`-based implementation of Algorithm 9, used as a test oracle. + pub(super) fn ntt_reference(p: &Polynomial) -> NttPolynomial { + let mut f = p.0; + let mut k = 1; + let mut len = 128; + while len >= 2 { + for start in (0..256).step_by(2 * len) { + let zeta = super::ZETA_POW_BITREV[k]; + k += 1; + for j in start..(start + len) { + let t = zeta * f[j + len]; + f[j + len] = f[j] - t; + f[j] = f[j] + t; + } + } + len /= 2; + } + NttPolynomial::new(f) + } + + /// The previous `Elem`-based implementation of Algorithm 10, used as a test oracle. + pub(super) fn ntt_inverse_reference(p: &NttPolynomial) -> Polynomial { + let mut f = p.0.clone(); + let mut k = 127; + let mut len = 2; + while len <= 128 { + for start in (0..256).step_by(2 * len) { + let zeta = super::ZETA_POW_BITREV[k]; + k -= 1; + for j in start..(start + len) { + let t = f[j]; + f[j] = t + f[j + len]; + f[j + len] = zeta * (f[j + len] - t); + } + } + len *= 2; + } + Elem::new(3303) * &Polynomial::new(f) + } + + /// Algorithms 11 and 12 written directly over `Elem`, summed over a vector. Test oracle. + pub(super) fn multiply_accumulate_reference( + a: &NttVector, + b: &NttVector, + ) -> NttPolynomial { + let mut out = NttPolynomial::default(); + for (a, b) in a.0.iter().zip(b.0.iter()) { + for i in 0..128 { + let (a0, a1) = (a.0[2 * i], a.0[2 * i + 1]); + let (b0, b1) = (b.0[2 * i], b.0[2 * i + 1]); + let gamma = super::GAMMA[i]; + out.0[2 * i] = out.0[2 * i] + a0 * b0 + a1 * b1 * gamma; + out.0[2 * i + 1] = out.0[2 * i + 1] + a0 * b1 + a1 * b0; + } + } + out + } + + /// Deterministic pseudo-random polynomial with coefficients in `[0, q)`. + #[allow( + clippy::cast_possible_truncation, + clippy::integer_division_remainder_used + )] + fn pseudo_random_poly(seed: u64) -> Array { + let mut x = seed.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1; + Array::from_fn(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + Elem::new(((x >> 16) % u64::from(BaseField::Q)) as Int) + }) + } + + /// The Montgomery kernels agree with the `Elem` reference on random and extreme inputs. + #[test] + #[allow(clippy::integer_division_remainder_used, reason = "test inputs")] + fn montgomery_kernels_match_reference() { + const FILLS: [Int; 4] = [0, 1, BaseField::Q - 1, (BaseField::Q - 1) / 2]; + let inputs: [Array; 210] = core::array::from_fn(|n| match n { + 0..200 => pseudo_random_poly(n as u64), + 200..204 => Array::from_fn(|_| Elem::new(FILLS[n - 200])), + _ => Array::from_fn(|i| { + Elem::new(if (i + n) % 2 == 0 { + BaseField::Q - 1 + } else { + 0 + }) + }), + }); + + for coeffs in &inputs { + let p = Polynomial::new(*coeffs); + assert_eq!(p.ntt(), ntt_reference(&p)); + + let p_hat = NttPolynomial::new(*coeffs); + assert_eq!(p_hat.ntt_inverse(), ntt_inverse_reference(&p_hat)); + } + + for pair in inputs.chunks_exact(6) { + let a: NttVector = NttVector::new(Array::from_fn(|i| NttPolynomial::new(pair[i]))); + let b: NttVector = + NttVector::new(Array::from_fn(|i| NttPolynomial::new(pair[i + 3]))); + assert_eq!( + super::multiply_accumulate(&a, &super::prepare(&b)), + multiply_accumulate_reference(&a, &b) + ); + assert_eq!(&a.0[0] * &b.0[0], { + let a1: NttVector = NttVector::new(Array([a.0[0].clone()])); + let b1: NttVector = NttVector::new(Array([b.0[0].clone()])); + multiply_accumulate_reference(&a1, &b1) + }); + } + } + #[test] #[allow(clippy::cast_possible_truncation)] fn polynomial_ops() { diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs new file mode 100644 index 00000000..e0966cd2 --- /dev/null +++ b/ml-kem/src/algebra/montgomery.rs @@ -0,0 +1,396 @@ +//! Signed 16-bit Montgomery arithmetic for the NTT and for products in the NTT domain. +//! +//! The generic `Elem` operators fully reduce after every addition, subtraction and +//! multiplication, using a hardened conditional subtraction. The kernels here instead keep +//! coefficients as partially reduced `i16` values and only use Montgomery and Barrett reductions, +//! which are pure arithmetic (multiply, shift, add). There are no selects, and so nothing that +//! depends on secret data for the compiler to turn into a branch. Each kernel converts its result +//! back to the canonical range `[0, q)` exactly once, at the end. +//! +//! Every bound below is stated for the inputs the kernels receive. Tests check the reductions +//! exhaustively over those domains, and debug builds check every addition for overflow. + +use super::{Elem, GAMMA, NttPolynomial, NttVector, ZETA_POW_BITREV}; +use array::{Array, ArraySize, typenum::U256}; + +/// The modulus `q`. +pub(super) const Q: i16 = 3329; + +/// `q^-1 mod 2^16`, as a signed value. +const QINV: i16 = -3327; + +/// The Montgomery radix `R = 2^16`, reduced mod `q`. +#[allow(clippy::integer_division_remainder_used, reason = "constant")] +const R_MOD_Q: i32 = (1 << 16) % Q as i32; + +/// `R^2 mod q`. Montgomery multiplication by this constant maps `x` to `x * R mod q`. +#[allow(clippy::integer_division_remainder_used, reason = "constant")] +const R2: i16 = centered((R_MOD_Q * R_MOD_Q) % Q as i32); + +/// `R / 128 mod q`. Montgomery multiplication by this constant divides by 128, which is the +/// final scaling step of `NTT^-1`. +#[allow(clippy::integer_division_remainder_used, reason = "constant")] +const INV_NTT_SCALE: i16 = centered((3303 * R_MOD_Q) % Q as i32); + +/// Map a value in `[0, q)` to its representative in `[-(q-1)/2, (q-1)/2]`. +#[allow(clippy::cast_possible_truncation, reason = "the value is below q")] +#[allow(clippy::integer_division_remainder_used, reason = "constant")] +const fn centered(x: i32) -> i16 { + if x > (Q as i32 - 1) / 2 { + (x - Q as i32) as i16 + } else { + x as i16 + } +} + +/// `ZETA_POW_BITREV` in the Montgomery domain: `zeta^BitRev7(i) * R mod q`, centered. +#[allow(clippy::integer_division_remainder_used, reason = "constant")] +const ZETAS: [i16; 128] = { + let mut out = [0; 128]; + let mut i = 0; + while i < 128 { + out[i] = centered((ZETA_POW_BITREV[i].0 as i32 * R_MOD_Q) % Q as i32); + i += 1; + } + out +}; + +/// `GAMMA` scaled by `R^2`: `zeta^(2 BitRev7(i) + 1) * R^2 mod q`, centered. Montgomery +/// multiplication of `x` by `GAMMAS_R2[i]` gives `x * GAMMA[i] * R mod q`. +#[allow(clippy::integer_division_remainder_used, reason = "constant")] +const GAMMAS_R2: [i16; 128] = { + let r2 = (R_MOD_Q * R_MOD_Q) % Q as i32; + let mut out = [0; 128]; + let mut i = 0; + while i < 128 { + out[i] = centered((GAMMA[i].0 as i32 * r2) % Q as i32); + i += 1; + } + out +}; + +/// Montgomery reduction: for `|a| < q * 2^15`, returns `r ≡ a * R^-1 (mod q)` with `|r| < q`. +#[inline(always)] +#[allow( + clippy::cast_possible_truncation, + reason = "truncation mod 2^16 is intended" +)] +fn montgomery_reduce(a: i32) -> i16 { + let t = (a as i16).wrapping_mul(QINV); + ((a - i32::from(t) * i32::from(Q)) >> 16) as i16 +} + +/// Montgomery multiplication: for `|a * b| < q * 2^15`, returns `r ≡ a * b * R^-1 (mod q)` with +/// `|r| < q`. +#[inline(always)] +fn fqmul(a: i16, b: i16) -> i16 { + montgomery_reduce(i32::from(a) * i32::from(b)) +} + +/// Barrett reduction: for any `a`, returns the representative of `a mod q` in +/// `[-(q-1)/2, (q-1)/2]`. +#[inline(always)] +#[allow( + clippy::cast_possible_truncation, + reason = "the result is below q in absolute value" +)] +fn barrett_reduce(a: i16) -> i16 { + #[allow(clippy::integer_division_remainder_used, reason = "constant")] + const V: i32 = ((1 << 26) + Q as i32 / 2) / Q as i32; + let t = (V * i32::from(a) + (1 << 25)) >> 26; + (i32::from(a) - t * i32::from(Q)) as i16 +} + +/// Map coefficients with `|a| < q` to their canonical representatives in `[0, q)`. +/// +/// Adding `q` to the negative coefficients is done with an arithmetic sign mask. LLVM recognizes +/// `(a >> 15) & q` as a select on the sign of `a`, and on targets without conditional moves (for +/// example `thumbv6m`) it lowers that select to a branch on the secret coefficient. The masks for +/// the whole polynomial are therefore passed through a single optimization barrier before they are +/// applied, which hides their origin while keeping both loops vectorizable. +#[inline(always)] +#[allow(clippy::cast_sign_loss, reason = "the result is in [0, q)")] +fn canonicalize(f: &[i16; 256]) -> Array { + let mut masks: [i16; 256] = core::array::from_fn(|i| f[i] >> 15); + core::hint::black_box(&mut masks); + Array::from_fn(|i| Elem::new((f[i] + (masks[i] & Q)) as u16)) +} + +/// Read a canonical coefficient in `[0, q)` as `i16`. +#[inline(always)] +#[allow(clippy::cast_possible_wrap, reason = "the value is below q")] +fn signed(x: Elem) -> i16 { + x.0 as i16 +} + +/// One layer of the forward NTT (Algorithm 9). +/// +/// Each layer adds at most `q` to the magnitude of a coefficient, since `fqmul` returns values +/// below `q` in absolute value. Inputs below `q` therefore stay below `8q < 2^15` after seven +/// layers, and no intermediate reduction is needed. +#[inline(always)] +fn ntt_layer(f: &mut [i16; 256], k: &mut usize) { + for block in f.chunks_exact_mut(2 * LEN).take(ITERATIONS) { + let zeta = ZETAS[*k]; + *k += 1; + + let (lo, hi) = block.split_at_mut(LEN); + for (a, b) in lo.iter_mut().zip(hi.iter_mut()) { + let t = fqmul(zeta, *b); + *b = *a - t; + *a += t; + } + } +} + +/// Algorithm 9: `NTT`, on coefficients in `[0, q)`. Returns coefficients in `[0, q)`. +pub(super) fn ntt(input: &Array) -> NttPolynomial { + let mut f: [i16; 256] = core::array::from_fn(|i| signed(input[i])); + let mut k = 1; + + ntt_layer::<128, 1>(&mut f, &mut k); + ntt_layer::<64, 2>(&mut f, &mut k); + ntt_layer::<32, 4>(&mut f, &mut k); + ntt_layer::<16, 8>(&mut f, &mut k); + ntt_layer::<8, 16>(&mut f, &mut k); + ntt_layer::<4, 32>(&mut f, &mut k); + ntt_layer::<2, 64>(&mut f, &mut k); + + for x in &mut f { + *x = barrett_reduce(*x); + } + NttPolynomial::new(canonicalize(&f)) +} + +/// One layer of the inverse NTT (Algorithm 10). +/// +/// Inputs are below `q` in absolute value. Sums and differences are below `2q`. The sum is +/// Barrett-reduced and the difference is Montgomery-multiplied, so outputs are again below `q`. +#[inline(always)] +fn ntt_inverse_layer(f: &mut [i16; 256], k: &mut usize) { + for block in f.chunks_exact_mut(2 * LEN).take(ITERATIONS) { + let zeta = ZETAS[*k]; + *k -= 1; + + let (lo, hi) = block.split_at_mut(LEN); + for (a, b) in lo.iter_mut().zip(hi.iter_mut()) { + let t = *a; + *a = barrett_reduce(t + *b); + *b = fqmul(zeta, *b - t); + } + } +} + +/// Algorithm 10: `NTT^-1`, on coefficients in `[0, q)`. Returns coefficients in `[0, q)`. +pub(super) fn ntt_inverse(input: &Array) -> Array { + let mut f: [i16; 256] = core::array::from_fn(|i| signed(input[i])); + let mut k = 127; + + ntt_inverse_layer::<2, 64>(&mut f, &mut k); + ntt_inverse_layer::<4, 32>(&mut f, &mut k); + ntt_inverse_layer::<8, 16>(&mut f, &mut k); + ntt_inverse_layer::<16, 8>(&mut f, &mut k); + ntt_inverse_layer::<32, 4>(&mut f, &mut k); + ntt_inverse_layer::<64, 2>(&mut f, &mut k); + ntt_inverse_layer::<128, 1>(&mut f, &mut k); + + // Multiply by 128^-1. The result is below q in absolute value. + for x in &mut f { + *x = fqmul(*x, INV_NTT_SCALE); + } + canonicalize(&f) +} + +/// The right-hand operand of `MultiplyNTTs` (Algorithm 11), prepared once so that it can be +/// multiplied by several left-hand operands. +/// +/// For each degree-one factor `b0 + b1 X` this stores `b0 * R`, `b1 * R` and `b1 * gamma * R` +/// (mod q, below q in absolute value). Products with canonical coefficients then need a single +/// Montgomery reduction, which removes the factor `R`, and `b1 * gamma` is computed once rather +/// than once per left-hand operand. +#[derive(Clone)] +pub(crate) struct PreparedNtt { + b0: [i16; 128], + b1: [i16; 128], + b1_gamma: [i16; 128], +} + +impl PreparedNtt { + pub(crate) fn new(b: &NttPolynomial) -> Self { + let mut out = Self { + b0: [0; 128], + b1: [0; 128], + b1_gamma: [0; 128], + }; + + for (i, pair) in b.0.chunks_exact(2).enumerate() { + out.b0[i] = fqmul(signed(pair[0]), R2); + out.b1[i] = fqmul(signed(pair[1]), R2); + out.b1_gamma[i] = fqmul(signed(pair[1]), GAMMAS_R2[i]); + } + + out + } +} + +/// Prepare every polynomial of a vector (see [`PreparedNtt`]). +pub(crate) fn prepare(v: &NttVector) -> Array { + Array::from_fn(|i| PreparedNtt::new(&v.0[i])) +} + +/// Algorithms 11 and 12 (`MultiplyNTTs` and `BaseCaseMultiply`), summed over a vector: +/// `sum_j a[j] * b[j]` in `T_q`. +/// +/// Products are accumulated in `i32` and reduced once per output coefficient. Coefficients of `a` +/// are in `[0, q)` and prepared values are below `q` in absolute value, so each output +/// coefficient gains less than `2 q^2` per term. With `K <= 4` terms the sum stays below +/// `8 q^2 < q * 2^15`, which is the input domain of `montgomery_reduce`. +pub(crate) fn multiply_accumulate( + a: &NttVector, + b: &Array, +) -> NttPolynomial { + debug_assert!(K::USIZE <= 4); + + let mut c0 = [0i32; 128]; + let mut c1 = [0i32; 128]; + + for (a, b) in a.0.iter().zip(b.iter()) { + for (i, pair) in a.0.chunks_exact(2).enumerate() { + let a0 = i32::from(pair[0].0); + let a1 = i32::from(pair[1].0); + c0[i] += a0 * i32::from(b.b0[i]) + a1 * i32::from(b.b1_gamma[i]); + c1[i] += a0 * i32::from(b.b1[i]) + a1 * i32::from(b.b0[i]); + } + } + + let mut out = [0i16; 256]; + for (i, pair) in out.chunks_exact_mut(2).enumerate() { + pair[0] = montgomery_reduce(c0[i]); + pair[1] = montgomery_reduce(c1[i]); + } + NttPolynomial::new(canonicalize(&out)) +} + +/// Multiply a `K x K` matrix by a prepared vector in `T_q`. +pub(crate) fn matrix_multiply( + rows: &Array, K>, + v: &Array, +) -> NttVector { + NttVector::new(Array::from_fn(|i| multiply_accumulate(&rows[i], v))) +} + +#[cfg(test)] +#[allow( + clippy::cast_possible_truncation, + clippy::cast_possible_wrap, + clippy::cast_sign_loss, + reason = "tests" +)] +#[allow( + clippy::integer_division_remainder_used, + reason = "test oracle on public values" +)] +mod tests { + use super::*; + use crate::algebra::BaseField; + use module_lattice::Field; + + fn modq(x: i64) -> i64 { + x.rem_euclid(i64::from(Q)) + } + + #[test] + fn constants() { + assert_eq!(i32::from(Q), i32::from(BaseField::Q)); + assert_eq!((i32::from(Q) * i32::from(QINV)) & 0xffff, 1); + assert_eq!(modq(i64::from(R2)), modq(1 << 32)); + assert_eq!(modq(i64::from(INV_NTT_SCALE) * 128), modq(1 << 16)); + for i in 0..128 { + let z = i64::from(ZETA_POW_BITREV[i].0); + let g = i64::from(GAMMA[i].0); + assert_eq!(modq(i64::from(ZETAS[i])), modq(z << 16)); + assert_eq!(modq(i64::from(GAMMAS_R2[i])), modq(modq(g << 16) << 16)); + assert!(ZETAS[i].abs() <= (Q - 1) / 2); + assert!(GAMMAS_R2[i].abs() <= (Q - 1) / 2); + } + } + + /// `montgomery_reduce` over its whole documented domain, with a stride that keeps the test + /// fast in debug builds, plus both ends of the domain. + #[test] + fn montgomery_reduce_domain() { + let bound = i32::from(Q) << 15; + let check = |a: i32| { + let r = montgomery_reduce(a); + assert!(r.abs() < Q, "a = {a}"); + assert_eq!(modq(i64::from(r) << 16), modq(i64::from(a)), "a = {a}"); + }; + let mut a = -bound + 1; + while a < bound { + check(a); + a += 97; + } + for a in [-bound + 1, -bound + 2, -1, 0, 1, bound - 2, bound - 1] { + check(a); + } + } + + /// `fqmul` for every multiplier used by the kernels and every input magnitude they can see. + #[test] + fn fqmul_all_constants() { + let constants = ZETAS + .iter() + .chain(GAMMAS_R2.iter()) + .chain([R2, INV_NTT_SCALE].iter()); + for &c in constants { + for a in (-8 * i32::from(Q) + 1)..(8 * i32::from(Q)) { + let a = a as i16; + let r = fqmul(a, c); + assert!(r.abs() < Q); + assert_eq!(modq(i64::from(r) << 16), modq(i64::from(a) * i64::from(c))); + } + } + } + + #[test] + fn barrett_reduce_exhaustive() { + for a in i16::MIN..=i16::MAX { + let r = barrett_reduce(a); + assert!(r.abs() <= (Q - 1) / 2, "a = {a}"); + assert_eq!(modq(i64::from(r)), modq(i64::from(a)), "a = {a}"); + } + } + + #[test] + fn canonicalize_exhaustive() { + let inputs: [i16; 2 * Q as usize - 1] = core::array::from_fn(|i| i as i16 - Q + 1); + for chunk in inputs.chunks(256) { + let mut f = [0i16; 256]; + f[..chunk.len()].copy_from_slice(chunk); + let out = canonicalize(&f); + for (a, r) in f.iter().zip(out.iter()) { + assert!(r.0 < BaseField::Q); + assert_eq!(i64::from(r.0), modq(i64::from(*a))); + } + } + } + + /// The largest possible accumulated product stays inside the Montgomery domain. + #[test] + fn multiply_accumulate_extremes() { + use array::typenum::U4; + + let max = Elem::new(BaseField::Q - 1); + let a: NttVector = NttVector::new(Array::from_fn(|_| { + NttPolynomial::new(Array::from_fn(|_| max)) + })); + for b in [max, Elem::new(1), Elem::new(0), Elem::new(1665)] { + let v: NttVector = NttVector::new(Array::from_fn(|_| { + NttPolynomial::new(Array::from_fn(|_| b)) + })); + let fast = multiply_accumulate(&a, &prepare(&v)); + let slow = super::super::test::multiply_accumulate_reference(&a, &v); + assert_eq!(fast, slow); + } + } +} diff --git a/ml-kem/src/pke.rs b/ml-kem/src/pke.rs index 29cc3cf1..2eb449af 100644 --- a/ml-kem/src/pke.rs +++ b/ml-kem/src/pke.rs @@ -1,7 +1,7 @@ use crate::B32; use crate::algebra::{ - Ntt, NttInverse, NttMatrix, NttVector, Polynomial, Vector, matrix_sample_ntt, sample_poly_cbd, - sample_poly_vec_cbd, + Ntt, NttInverse, NttMatrix, NttVector, Polynomial, Vector, matrix_multiply, matrix_sample_ntt, + multiply_accumulate, prepare, sample_poly_cbd, sample_poly_vec_cbd, }; use crate::compress::Compress; use crate::crypto::{G, PRF}; @@ -88,7 +88,7 @@ where let e_hat = e.ntt(); // Compute the public value - let t_hat = &(&A_hat * &s_hat) + &e_hat; + let t_hat = &matrix_multiply(&A_hat.0, &prepare(&s_hat)) + &e_hat; // Assemble the keys let dk = DecryptionKey { s_hat }; @@ -108,7 +108,7 @@ where v.decompress::(); let u_hat = u.ntt(); - let sTu = (&self.s_hat * &u_hat).ntt_inverse(); + let sTu = multiply_accumulate(&self.s_hat, &prepare(&u_hat)).ntt_inverse(); let mut w = &v - &sTu; Encode::::encode(w.compress::()) } @@ -151,13 +151,14 @@ where let A_hat_t: NttMatrix = matrix_sample_ntt(&self.rho, true); let r_hat: NttVector = r.ntt(); - let ATr: Vector = (&A_hat_t * &r_hat).ntt_inverse(); + let r_hat = prepare(&r_hat); + let ATr: Vector = matrix_multiply(&A_hat_t.0, &r_hat).ntt_inverse(); let mut u = ATr + e1; let mut mu: Polynomial = Encode::::decode(message); mu.decompress::(); - let tTr: Polynomial = (&self.t_hat * &r_hat).ntt_inverse(); + let tTr: Polynomial = multiply_accumulate(&self.t_hat, &r_hat).ntt_inverse(); let mut v = &(&tTr + &e2) + μ let c1 = Encode::::encode(u.compress::()); From b318f4cd79ece2c866f03d8b5b4163d187ff1d36 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Wed, 4 Nov 2026 11:48:37 +0100 Subject: [PATCH 04/10] ml-kem: bit-sliced cbd sampling --- ml-kem/src/algebra.rs | 103 +++++++++++++++++++++---------- ml-kem/src/algebra/montgomery.rs | 2 +- 2 files changed, 70 insertions(+), 35 deletions(-) diff --git a/ml-kem/src/algebra.rs b/ml-kem/src/algebra.rs index c2ea740b..151f1b9b 100644 --- a/ml-kem/src/algebra.rs +++ b/ml-kem/src/algebra.rs @@ -4,7 +4,7 @@ use crate::{ param::CbdSamplingSize, }; use array::{Array, ArraySize, typenum::U1}; -use module_lattice::{Encode, Field, MultiplyNtt, Truncate}; +use module_lattice::{Field, MultiplyNtt, Truncate}; use sha3::digest::XofReader; /// Section 4.3: arithmetic kernels for the NTT and for products in the NTT domain. @@ -96,38 +96,39 @@ pub(crate) fn matrix_sample_ntt(rho: &B32, transpose: bool) -> Ntt /// Algorithm 8: `SamplePolyCBD_eta(B)` /// -/// To avoid all the bitwise manipulation in the algorithm as written, we reuse the logic in -/// `ByteDecode`. We decode the PRF output into integers with `2 * eta` bits, then use -/// `count_ones` to perform the summation described in the algorithm. +/// The bits are summed in parallel within a machine word (SWAR), as in the CRYSTALS-Kyber +/// reference implementation. Every group of `eta` bytes encodes four coefficients of `2 * eta` +/// bits each. Adding the word to itself shifted right by `1..eta` places, masked to every +/// `eta`-th bit, leaves the sum of each run of `eta` bits in an `eta`-bit field. Coefficient `j` +/// is then the field at bit `2 * eta * j` minus the field at bit `2 * eta * j + eta`. /// /// The PRF output is secret, so each coefficient is computed arithmetically: no memory access /// or branch depends on the sampled bits (in particular, there is no lookup table). +#[allow(clippy::cast_possible_truncation, reason = "x and y are at most eta")] pub(crate) fn sample_poly_cbd(B: &PrfOutput) -> Polynomial where Eta: CbdSamplingSize, { - let vals: Polynomial = Encode::::decode(B); - Polynomial::new( - vals.0 - .iter() - .map(|val| cbd_coefficient::(val.0)) - .collect(), - ) -} + let eta = Eta::USIZE; + let ones: u32 = (0..8).fold(0, |m, k| m | (1 << (k * eta))); + let field: u32 = (1 << eta) - 1; -/// Map `2 * eta` sampled bits to a CBD coefficient: the number of ones in the low `eta` bits -/// minus the number of ones in the high `eta` bits, reduced mod `q`. -fn cbd_coefficient(bits: Int) -> Elem -where - Eta: CbdSamplingSize, -{ - let low_mask: Int = (1 << Eta::USIZE) - 1; - let x = >::truncate((bits & low_mask).count_ones()); - let y = >::truncate((bits >> Eta::USIZE).count_ones()); + let mut f = [0i16; 256]; + for (bytes, coeffs) in B.chunks_exact(eta).zip(f.chunks_exact_mut(4)) { + let t = bytes + .iter() + .rev() + .fold(0u32, |t, &b| (t << 8) | u32::from(b)); + let d = (0..eta).fold(0, |d, k| d + ((t >> k) & ones)); + + for (j, c) in coeffs.iter_mut().enumerate() { + let x = (d >> (2 * eta * j)) & field; + let y = (d >> (2 * eta * j + eta)) & field; + *c = x as i16 - y as i16; + } + } - // x, y <= eta, so `x + Q - y` is in `[Q - eta, Q + eta]` and one conditional subtraction - // fully reduces it. - Elem::new(BaseField::small_reduce(x + BaseField::Q - y)) + Polynomial::new(montgomery::canonicalize(&f)) } pub(crate) fn sample_poly_vec_cbd(sigma: &B32, start_n: u8) -> Vector @@ -655,21 +656,55 @@ mod test { test_sample(&sample, &CBD3); } - /// Check `cbd_coefficient` against Algorithm 8 (`x - y mod q`, where `x` and `y` are the sums - /// of the first and last `eta` bits) for every possible input. - #[test] + /// Algorithm 8 exactly as written in FIPS 203, used as a test oracle. #[allow( clippy::integer_division_remainder_used, reason = "test oracle on public values" )] - fn cbd_coefficient_matches_definition() { + fn sample_poly_cbd_reference( + B: &crate::crypto::PrfOutput, + ) -> Polynomial { + let eta = Eta::USIZE; + let bit = |i: usize| Int::from((B[i / 8] >> (i % 8)) & 1); + Polynomial::new(Array::from_fn(|i| { + let x: Int = (0..eta).map(|j| bit(2 * i * eta + j)).sum(); + let y: Int = (0..eta).map(|j| bit(2 * i * eta + eta + j)).sum(); + Elem::new((x + BaseField::Q - y) % BaseField::Q) + })) + } + + /// Check `sample_poly_cbd` against Algorithm 8 with every `2 * eta`-bit pattern in every + /// coefficient position, and on pseudo-random inputs. + #[test] + #[allow( + clippy::cast_possible_truncation, + clippy::integer_division_remainder_used, + reason = "test inputs" + )] + fn sample_poly_cbd_matches_definition() { fn check() { - let eta = Eta::U32; - for bits in 0..(1u16 << (2 * eta)) { - let x: u16 = (0..eta).map(|i| (bits >> i) & 1).sum(); - let y: u16 = (eta..2 * eta).map(|i| (bits >> i) & 1).sum(); - let expected = (x + BaseField::Q - y) % BaseField::Q; - assert_eq!(super::cbd_coefficient::(bits), Elem::new(expected)); + let width = 2 * Eta::USIZE; + for offset in 0..(1 << width) { + let mut B = crate::crypto::PrfOutput::::default(); + for i in 0..256 { + let pattern = (i + offset) % (1 << width); + for j in 0..width { + let pos = i * width + j; + B[pos / 8] |= (((pattern >> j) & 1) as u8) << (pos % 8); + } + } + assert_eq!( + super::sample_poly_cbd::(&B), + sample_poly_cbd_reference::(&B) + ); + } + + for seed in 0..64u8 { + let B = PRF::(&B32::from([seed; 32]), seed); + assert_eq!( + super::sample_poly_cbd::(&B), + sample_poly_cbd_reference::(&B) + ); } } diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs index e0966cd2..652cdac9 100644 --- a/ml-kem/src/algebra/montgomery.rs +++ b/ml-kem/src/algebra/montgomery.rs @@ -110,7 +110,7 @@ fn barrett_reduce(a: i16) -> i16 { /// applied, which hides their origin while keeping both loops vectorizable. #[inline(always)] #[allow(clippy::cast_sign_loss, reason = "the result is in [0, q)")] -fn canonicalize(f: &[i16; 256]) -> Array { +pub(super) fn canonicalize(f: &[i16; 256]) -> Array { let mut masks: [i16; 256] = core::array::from_fn(|i| f[i] >> 15); core::hint::black_box(&mut masks); Array::from_fn(|i| Elem::new((f[i] + (masks[i] & Q)) as u16)) From 352366faa76126ab3bae34463cdf949837380597 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Wed, 4 Nov 2026 19:20:05 +0100 Subject: [PATCH 05/10] ml-kem: select-free polynomial add and sub --- ml-kem/src/algebra.rs | 22 +++++++++++++++++++++ ml-kem/src/algebra/montgomery.rs | 34 +++++++++++++++++++++++++++++++- ml-kem/src/pke.rs | 11 ++++++----- 3 files changed, 61 insertions(+), 6 deletions(-) diff --git a/ml-kem/src/algebra.rs b/ml-kem/src/algebra.rs index 151f1b9b..8b9919ad 100644 --- a/ml-kem/src/algebra.rs +++ b/ml-kem/src/algebra.rs @@ -169,6 +169,28 @@ impl Ntt for Vector { } } +/// `a + b` in `R_q`. +pub(crate) fn poly_add(a: &Polynomial, b: &Polynomial) -> Polynomial { + Polynomial::new(montgomery::add(&a.0, &b.0)) +} + +/// `a - b` in `R_q`. +pub(crate) fn poly_sub(a: &Polynomial, b: &Polynomial) -> Polynomial { + Polynomial::new(montgomery::sub(&a.0, &b.0)) +} + +/// `a + b` for vectors over `R_q`. +pub(crate) fn vector_add(a: &Vector, b: &Vector) -> Vector { + Vector::new(Array::from_fn(|i| poly_add(&a.0[i], &b.0[i]))) +} + +/// `a + b` for vectors over `T_q`. +pub(crate) fn ntt_vector_add(a: &NttVector, b: &NttVector) -> NttVector { + NttVector::new(Array::from_fn(|i| { + NttPolynomial::new(montgomery::add(&a.0[i].0, &b.0[i].0)) + })) +} + /// The inverse NTT is the reverse of the Number Theoretic Transform, converting coefficient-wise /// products back into standard polynomial form while preserving correctness modulo the same prime. #[allow(clippy::module_name_repetitions)] diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs index 652cdac9..c26935c1 100644 --- a/ml-kem/src/algebra/montgomery.rs +++ b/ml-kem/src/algebra/montgomery.rs @@ -101,7 +101,7 @@ fn barrett_reduce(a: i16) -> i16 { (i32::from(a) - t * i32::from(Q)) as i16 } -/// Map coefficients with `|a| < q` to their canonical representatives in `[0, q)`. +/// Map coefficients in `[-q, q)` to their canonical representatives in `[0, q)`. /// /// Adding `q` to the negative coefficients is done with an arithmetic sign mask. LLVM recognizes /// `(a >> 15) & q` as a select on the sign of `a`, and on targets without conditional moves (for @@ -116,6 +116,20 @@ pub(super) fn canonicalize(f: &[i16; 256]) -> Array { Array::from_fn(|i| Elem::new((f[i] + (masks[i] & Q)) as u16)) } +/// Coefficient-wise `a + b`, for coefficients in `[0, q)`. +pub(super) fn add(a: &Array, b: &Array) -> Array { + // `a + b - q` is in `[-q, q - 2]`. + let f: [i16; 256] = core::array::from_fn(|i| signed(a[i]) + signed(b[i]) - Q); + canonicalize(&f) +} + +/// Coefficient-wise `a - b`, for coefficients in `[0, q)`. +pub(super) fn sub(a: &Array, b: &Array) -> Array { + // `a - b` is in `(-q, q)`. + let f: [i16; 256] = core::array::from_fn(|i| signed(a[i]) - signed(b[i])); + canonicalize(&f) +} + /// Read a canonical coefficient in `[0, q)` as `i16`. #[inline(always)] #[allow(clippy::cast_possible_wrap, reason = "the value is below q")] @@ -375,6 +389,24 @@ mod tests { } } + /// `add` and `sub` agree with the `Elem` operators for every pair of field elements. + #[test] + fn add_sub_exhaustive() { + for a in 0..BaseField::Q { + let lhs = Array::from_fn(|_| Elem::new(a)); + for chunk in 0..BaseField::Q.div_ceil(256) { + let rhs: Array = + Array::from_fn(|i| Elem::new((chunk * 256 + i as u16) % BaseField::Q)); + let sum = add(&lhs, &rhs); + let difference = sub(&lhs, &rhs); + for i in 0..256 { + assert_eq!(sum[i], lhs[i] + rhs[i]); + assert_eq!(difference[i], lhs[i] - rhs[i]); + } + } + } + } + /// The largest possible accumulated product stays inside the Montgomery domain. #[test] fn multiply_accumulate_extremes() { diff --git a/ml-kem/src/pke.rs b/ml-kem/src/pke.rs index 2eb449af..b2c24fe7 100644 --- a/ml-kem/src/pke.rs +++ b/ml-kem/src/pke.rs @@ -1,7 +1,8 @@ use crate::B32; use crate::algebra::{ Ntt, NttInverse, NttMatrix, NttVector, Polynomial, Vector, matrix_multiply, matrix_sample_ntt, - multiply_accumulate, prepare, sample_poly_cbd, sample_poly_vec_cbd, + multiply_accumulate, ntt_vector_add, poly_add, poly_sub, prepare, sample_poly_cbd, + sample_poly_vec_cbd, vector_add, }; use crate::compress::Compress; use crate::crypto::{G, PRF}; @@ -88,7 +89,7 @@ where let e_hat = e.ntt(); // Compute the public value - let t_hat = &matrix_multiply(&A_hat.0, &prepare(&s_hat)) + &e_hat; + let t_hat = ntt_vector_add(&matrix_multiply(&A_hat.0, &prepare(&s_hat)), &e_hat); // Assemble the keys let dk = DecryptionKey { s_hat }; @@ -109,7 +110,7 @@ where let u_hat = u.ntt(); let sTu = multiply_accumulate(&self.s_hat, &prepare(&u_hat)).ntt_inverse(); - let mut w = &v - &sTu; + let mut w = poly_sub(&v, &sTu); Encode::::encode(w.compress::()) } @@ -153,13 +154,13 @@ where let r_hat: NttVector = r.ntt(); let r_hat = prepare(&r_hat); let ATr: Vector = matrix_multiply(&A_hat_t.0, &r_hat).ntt_inverse(); - let mut u = ATr + e1; + let mut u = vector_add(&ATr, &e1); let mut mu: Polynomial = Encode::::decode(message); mu.decompress::(); let tTr: Polynomial = multiply_accumulate(&self.t_hat, &r_hat).ntt_inverse(); - let mut v = &(&tTr + &e2) + μ + let mut v = poly_add(&poly_add(&tTr, &e2), &mu); let c1 = Encode::::encode(u.compress::()); let c2 = Encode::::encode(v.compress::()); From 5b9965ec8446427125190b056c55e26a63da3a36 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Thu, 5 Nov 2026 15:37:44 +0100 Subject: [PATCH 06/10] ml-kem: cache the expanded matrix in heap keys --- ml-kem/src/algebra.rs | 22 ++++++++- ml-kem/src/algebra/montgomery.rs | 23 +++++++++- ml-kem/src/pke.rs | 79 +++++++++++++++++++++++++------- 3 files changed, 104 insertions(+), 20 deletions(-) diff --git a/ml-kem/src/algebra.rs b/ml-kem/src/algebra.rs index 8b9919ad..dcf1a85f 100644 --- a/ml-kem/src/algebra.rs +++ b/ml-kem/src/algebra.rs @@ -10,7 +10,9 @@ use sha3::digest::XofReader; /// Section 4.3: arithmetic kernels for the NTT and for products in the NTT domain. mod montgomery; -pub(crate) use montgomery::{PreparedNtt, matrix_multiply, multiply_accumulate, prepare}; +pub(crate) use montgomery::{ + PreparedNtt, matrix_multiply, matrix_transpose_multiply, multiply_accumulate, prepare, +}; module_lattice::define_field!(BaseField, u16, u32, u64, 3329); @@ -327,6 +329,24 @@ mod test { })) } + /// Multiplying by the transpose of a matrix agrees with transposing it first. + #[test] + fn matrix_transpose_multiply() { + let rho = B32::from([7; 32]); + let a: NttMatrix = super::matrix_sample_ntt(&rho, false); + let v: NttVector = + super::matrix_sample_ntt::(&B32::from([8; 32]), false).0[0].clone(); + let v = super::prepare(&v); + assert_eq!( + super::matrix_transpose_multiply(&a.0, &v), + super::matrix_multiply(&matrix_transpose(&a).0, &v) + ); + assert_eq!( + super::matrix_transpose_multiply(&a.0, &v), + super::matrix_multiply(&super::matrix_sample_ntt::(&rho, true).0, &v) + ); + } + /// The previous `Elem`-based implementation of Algorithm 9, used as a test oracle. pub(super) fn ntt_reference(p: &Polynomial) -> NttPolynomial { let mut f = p.0; diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs index c26935c1..477eb4a4 100644 --- a/ml-kem/src/algebra/montgomery.rs +++ b/ml-kem/src/algebra/montgomery.rs @@ -264,11 +264,19 @@ pub(crate) fn multiply_accumulate( b: &Array, ) -> NttPolynomial { debug_assert!(K::USIZE <= 4); + accumulate(a.0.iter().zip(b.iter())) +} +/// The sum of the products of `(a, b)` pairs in `T_q`, for at most four pairs. See +/// [`multiply_accumulate`] for the bounds. +#[inline(always)] +fn accumulate<'a>( + terms: impl Iterator, +) -> NttPolynomial { let mut c0 = [0i32; 128]; let mut c1 = [0i32; 128]; - for (a, b) in a.0.iter().zip(b.iter()) { + for (a, b) in terms { for (i, pair) in a.0.chunks_exact(2).enumerate() { let a0 = i32::from(pair[0].0); let a1 = i32::from(pair[1].0); @@ -285,7 +293,7 @@ pub(crate) fn multiply_accumulate( NttPolynomial::new(canonicalize(&out)) } -/// Multiply a `K x K` matrix by a prepared vector in `T_q`. +/// Multiply a `K x K` matrix, given by its rows, by a prepared vector in `T_q`. pub(crate) fn matrix_multiply( rows: &Array, K>, v: &Array, @@ -293,6 +301,17 @@ pub(crate) fn matrix_multiply( NttVector::new(Array::from_fn(|i| multiply_accumulate(&rows[i], v))) } +/// Multiply the transpose of a `K x K` matrix, given by its rows, by a prepared vector in `T_q`. +pub(crate) fn matrix_transpose_multiply( + rows: &Array, K>, + v: &Array, +) -> NttVector { + debug_assert!(K::USIZE <= 4); + NttVector::new(Array::from_fn(|i| { + accumulate(rows.iter().map(|row| &row.0[i]).zip(v.iter())) + })) +} + #[cfg(test)] #[allow( clippy::cast_possible_truncation, diff --git a/ml-kem/src/pke.rs b/ml-kem/src/pke.rs index b2c24fe7..30530484 100644 --- a/ml-kem/src/pke.rs +++ b/ml-kem/src/pke.rs @@ -1,8 +1,8 @@ use crate::B32; use crate::algebra::{ Ntt, NttInverse, NttMatrix, NttVector, Polynomial, Vector, matrix_multiply, matrix_sample_ntt, - multiply_accumulate, ntt_vector_add, poly_add, poly_sub, prepare, sample_poly_cbd, - sample_poly_vec_cbd, vector_add, + matrix_transpose_multiply, multiply_accumulate, ntt_vector_add, poly_add, poly_sub, prepare, + sample_poly_cbd, sample_poly_vec_cbd, vector_add, }; use crate::compress::Compress; use crate::crypto::{G, PRF}; @@ -93,7 +93,12 @@ where // Assemble the keys let dk = DecryptionKey { s_hat }; - let ek = EncryptionKey { t_hat, rho }; + let ek = EncryptionKey { + t_hat, + rho, + #[cfg(feature = "alloc")] + A_hat, + }; (dk, ek) } @@ -128,15 +133,49 @@ where /// An `EncryptionKey` provides the ability to encrypt a value so that it can only be /// decrypted by the holder of the corresponding decapsulation key. -#[derive(Clone, Default, Debug, Eq, PartialEq)] +/// +/// There is deliberately no `Default` impl: with the `alloc` feature a default key would pair +/// `rho = 0` with an all-zero matrix instead of the matrix expanded from `rho`. +#[derive(Clone)] pub(crate) struct EncryptionKey

where P: PkeParams, { t_hat: NttVector, rho: B32, + + /// The matrix `A_hat` expanded from `rho`. Expanding it runs `SampleNTT` `k^2` times, which + /// is the most expensive step of `K-PKE.Encrypt`, so keys that live on the heap keep a copy. + /// It is public and fully determined by `rho`. + #[cfg(feature = "alloc")] + A_hat: NttMatrix, } +// Handwritten to leave out `A_hat`, which is determined by `rho` +impl

fmt::Debug for EncryptionKey

+where + P: PkeParams, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("EncryptionKey") + .field("t_hat", &self.t_hat) + .field("rho", &self.rho) + .finish() + } +} + +// Handwritten to leave out `A_hat`, which is determined by `rho` +impl

PartialEq for EncryptionKey

+where + P: PkeParams, +{ + fn eq(&self, other: &Self) -> bool { + self.t_hat == other.t_hat && self.rho == other.rho + } +} + +impl

Eq for EncryptionKey

where P: PkeParams {} + impl

EncryptionKey

where P: PkeParams, @@ -150,10 +189,14 @@ where let prf_output = PRF::(randomness, 2 * P::K::U8); let e2: Polynomial = sample_poly_cbd::(&prf_output); - let A_hat_t: NttMatrix = matrix_sample_ntt(&self.rho, true); + #[cfg(feature = "alloc")] + let A_hat = &self.A_hat; + #[cfg(not(feature = "alloc"))] + let A_hat = &matrix_sample_ntt::(&self.rho, false); + let r_hat: NttVector = r.ntt(); let r_hat = prepare(&r_hat); - let ATr: Vector = matrix_multiply(&A_hat_t.0, &r_hat).ntt_inverse(); + let ATr: Vector = matrix_transpose_multiply(&A_hat.0, &r_hat).ntt_inverse(); let mut u = vector_add(&ATr, &e1); let mut mu: Polynomial = Encode::::decode(message); @@ -179,12 +222,8 @@ where /// Returns [`InvalidKey`] in the event that the key fails the encapsulation key checks /// specified in FIPS 203 §7.2. pub(crate) fn from_bytes(enc: &EncodedEncryptionKey

) -> Result { - let (t_hat, rho) = P::split_ek(enc); - let t_hat = P::decode_u12(t_hat); - let ret = Self { - t_hat, - rho: rho.clone(), - }; + let (t_hat_enc, rho) = P::split_ek(enc); + let t_hat = P::decode_u12(t_hat_enc); // Check the candidate encapsulation key is valid using the method specified in FIPS 203 // §7.2 ML-KEM Encapsulation: @@ -211,12 +250,18 @@ where // // #1 is performed by the `EncodedEncryptionKey` type, and the following check vicariously // performs #2 by encoding the integer-mod-q array using our implementation of ByteEncode₁₂ - // and comparing the resulting serialization to see if it round-trips. - if &ret.to_bytes() == enc { - Ok(ret) - } else { - Err(InvalidKey) + // and comparing the resulting serialization to see if it round-trips. It runs before + // `A_hat` is expanded, so that invalid keys are rejected cheaply. + if P::encode_u12(&t_hat) != *t_hat_enc { + return Err(InvalidKey); } + + Ok(Self { + t_hat, + rho: rho.clone(), + #[cfg(feature = "alloc")] + A_hat: matrix_sample_ntt(rho, false), + }) } } From fcd56d1df76f45c421b166f66b8c667d6869c059 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Fri, 6 Nov 2026 12:09:18 +0100 Subject: [PATCH 07/10] ml-kem: 16-bit montgomery and barrett reductions --- ml-kem/src/algebra/montgomery.rs | 87 ++++++++++++++++++++++++++------ 1 file changed, 71 insertions(+), 16 deletions(-) diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs index 477eb4a4..e7f7214c 100644 --- a/ml-kem/src/algebra/montgomery.rs +++ b/ml-kem/src/algebra/montgomery.rs @@ -55,6 +55,17 @@ const ZETAS: [i16; 128] = { out }; +/// `qinv_multiple` of each of `ZETAS`. +const ZETAS_QINV: [i16; 128] = { + let mut out = [0; 128]; + let mut i = 0; + while i < 128 { + out[i] = qinv_multiple(ZETAS[i]); + i += 1; + } + out +}; + /// `GAMMA` scaled by `R^2`: `zeta^(2 BitRev7(i) + 1) * R^2 mod q`, centered. Montgomery /// multiplication of `x` by `GAMMAS_R2[i]` gives `x * GAMMA[i] * R mod q`. #[allow(clippy::integer_division_remainder_used, reason = "constant")] @@ -69,6 +80,17 @@ const GAMMAS_R2: [i16; 128] = { out }; +/// `qinv_multiple` of each of `GAMMAS_R2`. +const GAMMAS_R2_QINV: [i16; 128] = { + let mut out = [0; 128]; + let mut i = 0; + while i < 128 { + out[i] = qinv_multiple(GAMMAS_R2[i]); + i += 1; + } + out +}; + /// Montgomery reduction: for `|a| < q * 2^15`, returns `r ≡ a * R^-1 (mod q)` with `|r| < q`. #[inline(always)] #[allow( @@ -81,24 +103,53 @@ fn montgomery_reduce(a: i32) -> i16 { } /// Montgomery multiplication: for `|a * b| < q * 2^15`, returns `r ≡ a * b * R^-1 (mod q)` with -/// `|r| < q`. -#[inline(always)] +/// `|r| < q`. Reference for [`fqmul_const`]. +#[cfg(test)] fn fqmul(a: i16, b: i16) -> i16 { montgomery_reduce(i32::from(a) * i32::from(b)) } -/// Barrett reduction: for any `a`, returns the representative of `a mod q` in -/// `[-(q-1)/2, (q-1)/2]`. +/// The high 16 bits of the product of `a` and `b`. #[inline(always)] #[allow( clippy::cast_possible_truncation, - reason = "the result is below q in absolute value" + reason = "the high half of a product of two i16 values fits in i16" )] +fn mulhi(a: i16, b: i16) -> i16 { + ((i32::from(a) * i32::from(b)) >> 16) as i16 +} + +/// `b * q^-1 mod 2^16`, the second constant that [`fqmul_const`] needs for a multiplier `b`. +const fn qinv_multiple(b: i16) -> i16 { + b.wrapping_mul(QINV) +} + +/// Montgomery multiplication by a constant `b`, given `b_qinv = qinv_multiple(b)`. The result is +/// exactly `fqmul(a, b)`. +/// +/// Montgomery reduction subtracts `t * q`, where `t = a * b * q^-1 mod 2^16`, from `a * b`. The +/// low 16 bits of the two products are equal, so the result is the difference of their high +/// halves (Seiler, ePrint 2018/039). Computed this way, every operation works on 16-bit lanes, +/// which compilers map onto SIMD instructions such as `pmulhw` and `pmullw`. +#[inline(always)] +fn fqmul_const(a: i16, b: i16, b_qinv: i16) -> i16 { + mulhi(a, b) - mulhi(a.wrapping_mul(b_qinv), Q) +} + +/// Barrett reduction: for any `a`, returns the representative of `a mod q` in +/// `[-(q-1)/2, (q-1)/2]`. +/// +/// The quotient `round(a * V / 2^26)`, with `V = round(2^26 / q)`, is computed from the high half +/// of `a * V` as `(mulhi(a, V) + 2^9) >> 10`, which is equal because nested floors of divisions +/// by integers collapse. The quotient has absolute value at most 10, and the result fits in +/// `i16`, so the wrapping operations below are exact. +#[inline(always)] +#[allow(clippy::cast_possible_truncation, reason = "V is below 2^15")] fn barrett_reduce(a: i16) -> i16 { #[allow(clippy::integer_division_remainder_used, reason = "constant")] - const V: i32 = ((1 << 26) + Q as i32 / 2) / Q as i32; - let t = (V * i32::from(a) + (1 << 25)) >> 26; - (i32::from(a) - t * i32::from(Q)) as i16 + const V: i16 = (((1 << 26) + Q as i32 / 2) / Q as i32) as i16; + let t = (mulhi(a, V) + (1 << 9)) >> 10; + a.wrapping_sub(t.wrapping_mul(Q)) } /// Map coefficients in `[-q, q)` to their canonical representatives in `[0, q)`. @@ -145,12 +196,12 @@ fn signed(x: Elem) -> i16 { #[inline(always)] fn ntt_layer(f: &mut [i16; 256], k: &mut usize) { for block in f.chunks_exact_mut(2 * LEN).take(ITERATIONS) { - let zeta = ZETAS[*k]; + let (zeta, zeta_qinv) = (ZETAS[*k], ZETAS_QINV[*k]); *k += 1; let (lo, hi) = block.split_at_mut(LEN); for (a, b) in lo.iter_mut().zip(hi.iter_mut()) { - let t = fqmul(zeta, *b); + let t = fqmul_const(*b, zeta, zeta_qinv); *b = *a - t; *a += t; } @@ -183,14 +234,14 @@ pub(super) fn ntt(input: &Array) -> NttPolynomial { #[inline(always)] fn ntt_inverse_layer(f: &mut [i16; 256], k: &mut usize) { for block in f.chunks_exact_mut(2 * LEN).take(ITERATIONS) { - let zeta = ZETAS[*k]; + let (zeta, zeta_qinv) = (ZETAS[*k], ZETAS_QINV[*k]); *k -= 1; let (lo, hi) = block.split_at_mut(LEN); for (a, b) in lo.iter_mut().zip(hi.iter_mut()) { let t = *a; *a = barrett_reduce(t + *b); - *b = fqmul(zeta, *b - t); + *b = fqmul_const(*b - t, zeta, zeta_qinv); } } } @@ -210,7 +261,7 @@ pub(super) fn ntt_inverse(input: &Array) -> Array { // Multiply by 128^-1. The result is below q in absolute value. for x in &mut f { - *x = fqmul(*x, INV_NTT_SCALE); + *x = fqmul_const(*x, INV_NTT_SCALE, qinv_multiple(INV_NTT_SCALE)); } canonicalize(&f) } @@ -238,9 +289,10 @@ impl PreparedNtt { }; for (i, pair) in b.0.chunks_exact(2).enumerate() { - out.b0[i] = fqmul(signed(pair[0]), R2); - out.b1[i] = fqmul(signed(pair[1]), R2); - out.b1_gamma[i] = fqmul(signed(pair[1]), GAMMAS_R2[i]); + let r2_qinv = qinv_multiple(R2); + out.b0[i] = fqmul_const(signed(pair[0]), R2, r2_qinv); + out.b1[i] = fqmul_const(signed(pair[1]), R2, r2_qinv); + out.b1_gamma[i] = fqmul_const(signed(pair[1]), GAMMAS_R2[i], GAMMAS_R2_QINV[i]); } out @@ -379,6 +431,7 @@ mod tests { for a in (-8 * i32::from(Q) + 1)..(8 * i32::from(Q)) { let a = a as i16; let r = fqmul(a, c); + assert_eq!(fqmul_const(a, c, qinv_multiple(c)), r); assert!(r.abs() < Q); assert_eq!(modq(i64::from(r) << 16), modq(i64::from(a) * i64::from(c))); } @@ -389,6 +442,8 @@ mod tests { fn barrett_reduce_exhaustive() { for a in i16::MIN..=i16::MAX { let r = barrett_reduce(a); + let t = (20159 * i32::from(a) + (1 << 25)) >> 26; + assert_eq!(i32::from(r), i32::from(a) - t * i32::from(Q), "a = {a}"); assert!(r.abs() <= (Q - 1) / 2, "a = {a}"); assert_eq!(modq(i64::from(r)), modq(i64::from(a)), "a = {a}"); } From f330642bacb655614210f09dd5a9e02efe86fb07 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Fri, 6 Nov 2026 18:44:29 +0100 Subject: [PATCH 08/10] ml-kem: expand the matrix first in encrypt --- ml-kem/src/pke.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/ml-kem/src/pke.rs b/ml-kem/src/pke.rs index 30530484..cb66803b 100644 --- a/ml-kem/src/pke.rs +++ b/ml-kem/src/pke.rs @@ -183,17 +183,19 @@ where /// Encrypt the specified message for the holder of the corresponding decryption key, using the /// provided randomness, according the `K-PKE.Encrypt` procedure. pub(crate) fn encrypt(&self, message: &B32, randomness: &B32) -> Ciphertext

{ + // Expand the public matrix before sampling anything secret, in the order of Algorithm 14 + // (steps 4-8, then 9-17). + #[cfg(feature = "alloc")] + let A_hat = &self.A_hat; + #[cfg(not(feature = "alloc"))] + let A_hat = &matrix_sample_ntt::(&self.rho, false); + let r = sample_poly_vec_cbd::(randomness, 0); let e1 = sample_poly_vec_cbd::(randomness, P::K::U8); let prf_output = PRF::(randomness, 2 * P::K::U8); let e2: Polynomial = sample_poly_cbd::(&prf_output); - #[cfg(feature = "alloc")] - let A_hat = &self.A_hat; - #[cfg(not(feature = "alloc"))] - let A_hat = &matrix_sample_ntt::(&self.rho, false); - let r_hat: NttVector = r.ntt(); let r_hat = prepare(&r_hat); let ATr: Vector = matrix_transpose_multiply(&A_hat.0, &r_hat).ntt_inverse(); From 9b4e0b7b92578e5ce7bd3d981e53447656abc8e2 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Sat, 7 Nov 2026 16:02:56 +0100 Subject: [PATCH 09/10] ml-kem: keep canonicalization out of line --- ml-kem/src/algebra/montgomery.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs index e7f7214c..07d0dd37 100644 --- a/ml-kem/src/algebra/montgomery.rs +++ b/ml-kem/src/algebra/montgomery.rs @@ -159,7 +159,10 @@ fn barrett_reduce(a: i16) -> i16 { /// example `thumbv6m`) it lowers that select to a branch on the secret coefficient. The masks for /// the whole polynomial are therefore passed through a single optimization barrier before they are /// applied, which hides their origin while keeping both loops vectorizable. -#[inline(always)] +/// +/// Every kernel ends with this pass, so it is kept out of line: inlining it everywhere costs about +/// 10 KiB of code on thumbv7em and is not measurably faster. +#[inline(never)] #[allow(clippy::cast_sign_loss, reason = "the result is in [0, q)")] pub(super) fn canonicalize(f: &[i16; 256]) -> Array { let mut masks: [i16; 256] = core::array::from_fn(|i| f[i] >> 15); @@ -321,7 +324,6 @@ pub(crate) fn multiply_accumulate( /// The sum of the products of `(a, b)` pairs in `T_q`, for at most four pairs. See /// [`multiply_accumulate`] for the bounds. -#[inline(always)] fn accumulate<'a>( terms: impl Iterator, ) -> NttPolynomial { From a471d154011a8b073509b0663614b319b424af41 Mon Sep 17 00:00:00 2001 From: ccipherx Date: Sun, 8 Nov 2026 13:26:40 +0100 Subject: [PATCH 10/10] ml-kem: shorten comments and add key and boundary tests --- ml-kem/benches/mlkem.rs | 17 ++-- ml-kem/src/algebra.rs | 23 ++---- ml-kem/src/algebra/montgomery.rs | 138 +++++++++++-------------------- ml-kem/src/pke.rs | 45 +++++++--- ml-kem/tests/accumulated.rs | 21 ++--- 5 files changed, 102 insertions(+), 142 deletions(-) diff --git a/ml-kem/benches/mlkem.rs b/ml-kem/benches/mlkem.rs index 11e0e574..bd4cf8f2 100644 --- a/ml-kem/benches/mlkem.rs +++ b/ml-kem/benches/mlkem.rs @@ -1,7 +1,6 @@ //! ML-KEM benchmarks. //! -//! All inputs are fixed byte strings so that no RNG or system call runs inside a timed region, -//! and every output is passed through `black_box` so that no work can be optimized away. +//! Inputs are fixed, so no RNG runs in timed loops. Outputs go through `black_box`. #![allow(missing_docs, clippy::unwrap_used)] @@ -10,7 +9,7 @@ use core::hint::black_box; use criterion::{Criterion, criterion_group, criterion_main}; use ml_kem::*; -/// Benchmark every public operation of one parameter set. +/// Benchmarks one parameter set. macro_rules! bench_parameter_set { ($c:expr, $name:literal, $module:ident) => {{ type Dk = $module::DecapsulationKey; @@ -28,34 +27,30 @@ macro_rules! bench_parameter_set { let mut group = $c.benchmark_group($name); - // Algorithm 19 (deterministic part): expand a seed into a decapsulation key. + // KeyGen from a seed. group.bench_function("keygen", |b| { b.iter(|| Dk::new(black_box(&seed))); }); - // FIPS 203 §7.2 input checks plus H(ek). + // §7.2 checks plus H(ek). group.bench_function("ek_parse", |b| { b.iter(|| Ek::new(black_box(&ek_bytes)).unwrap()); }); - // Algorithm 17 with a fixed message. group.bench_function("encapsulate", |b| { b.iter(|| ek.encapsulate_deterministic(black_box(&m))); }); - // Algorithm 18 on a valid ciphertext. group.bench_function("decapsulate", |b| { b.iter(|| dk.decapsulate(black_box(&ct))); }); - // Algorithm 18 on an invalid ciphertext (implicit rejection). Must cost the same as the - // valid case. + // Implicit rejection. Must cost the same as `decapsulate`. group.bench_function("decapsulate_reject", |b| { b.iter(|| dk.decapsulate(black_box(&bad_ct))); }); - // One ephemeral exchange: key generation, encapsulation key parsing, encapsulation and - // decapsulation. + // KeyGen, ek parsing, Encaps and Decaps. group.bench_function("round_trip", |b| { b.iter(|| { let dk = Dk::new(black_box(&seed)); diff --git a/ml-kem/src/algebra.rs b/ml-kem/src/algebra.rs index dcf1a85f..8adf5bdd 100644 --- a/ml-kem/src/algebra.rs +++ b/ml-kem/src/algebra.rs @@ -7,7 +7,6 @@ use array::{Array, ArraySize, typenum::U1}; use module_lattice::{Field, MultiplyNtt, Truncate}; use sha3::digest::XofReader; -/// Section 4.3: arithmetic kernels for the NTT and for products in the NTT domain. mod montgomery; pub(crate) use montgomery::{ @@ -98,11 +97,8 @@ pub(crate) fn matrix_sample_ntt(rho: &B32, transpose: bool) -> Ntt /// Algorithm 8: `SamplePolyCBD_eta(B)` /// -/// The bits are summed in parallel within a machine word (SWAR), as in the CRYSTALS-Kyber -/// reference implementation. Every group of `eta` bytes encodes four coefficients of `2 * eta` -/// bits each. Adding the word to itself shifted right by `1..eta` places, masked to every -/// `eta`-th bit, leaves the sum of each run of `eta` bits in an `eta`-bit field. Coefficient `j` -/// is then the field at bit `2 * eta * j` minus the field at bit `2 * eta * j + eta`. +/// Bits are summed in parallel within a word (SWAR), as in the CRYSTALS-Kyber reference code. +/// Each `eta` bytes hold four coefficients. /// /// The PRF output is secret, so each coefficient is computed arithmetically: no memory access /// or branch depends on the sampled bits (in particular, there is no lookup table). @@ -329,7 +325,6 @@ mod test { })) } - /// Multiplying by the transpose of a matrix agrees with transposing it first. #[test] fn matrix_transpose_multiply() { let rho = B32::from([7; 32]); @@ -347,7 +342,7 @@ mod test { ); } - /// The previous `Elem`-based implementation of Algorithm 9, used as a test oracle. + /// Algorithm 9 over `Elem`. Test oracle. pub(super) fn ntt_reference(p: &Polynomial) -> NttPolynomial { let mut f = p.0; let mut k = 1; @@ -367,7 +362,7 @@ mod test { NttPolynomial::new(f) } - /// The previous `Elem`-based implementation of Algorithm 10, used as a test oracle. + /// Algorithm 10 over `Elem`. Test oracle. pub(super) fn ntt_inverse_reference(p: &NttPolynomial) -> Polynomial { let mut f = p.0.clone(); let mut k = 127; @@ -387,7 +382,7 @@ mod test { Elem::new(3303) * &Polynomial::new(f) } - /// Algorithms 11 and 12 written directly over `Elem`, summed over a vector. Test oracle. + /// Algorithms 11 and 12 over `Elem`, summed over a vector. Test oracle. pub(super) fn multiply_accumulate_reference( a: &NttVector, b: &NttVector, @@ -405,7 +400,7 @@ mod test { out } - /// Deterministic pseudo-random polynomial with coefficients in `[0, q)`. + /// Deterministic pseudo-random polynomial. #[allow( clippy::cast_possible_truncation, clippy::integer_division_remainder_used @@ -420,7 +415,6 @@ mod test { }) } - /// The Montgomery kernels agree with the `Elem` reference on random and extreme inputs. #[test] #[allow(clippy::integer_division_remainder_used, reason = "test inputs")] fn montgomery_kernels_match_reference() { @@ -698,7 +692,7 @@ mod test { test_sample(&sample, &CBD3); } - /// Algorithm 8 exactly as written in FIPS 203, used as a test oracle. + /// Algorithm 8 as written in FIPS 203. Test oracle. #[allow( clippy::integer_division_remainder_used, reason = "test oracle on public values" @@ -715,8 +709,7 @@ mod test { })) } - /// Check `sample_poly_cbd` against Algorithm 8 with every `2 * eta`-bit pattern in every - /// coefficient position, and on pseudo-random inputs. + /// Every `2 * eta`-bit pattern in every position, plus PRF outputs. #[test] #[allow( clippy::cast_possible_truncation, diff --git a/ml-kem/src/algebra/montgomery.rs b/ml-kem/src/algebra/montgomery.rs index 07d0dd37..601843e8 100644 --- a/ml-kem/src/algebra/montgomery.rs +++ b/ml-kem/src/algebra/montgomery.rs @@ -1,38 +1,29 @@ -//! Signed 16-bit Montgomery arithmetic for the NTT and for products in the NTT domain. +//! Signed 16-bit Montgomery arithmetic for the NTT and NTT-domain products. //! -//! The generic `Elem` operators fully reduce after every addition, subtraction and -//! multiplication, using a hardened conditional subtraction. The kernels here instead keep -//! coefficients as partially reduced `i16` values and only use Montgomery and Barrett reductions, -//! which are pure arithmetic (multiply, shift, add). There are no selects, and so nothing that -//! depends on secret data for the compiler to turn into a branch. Each kernel converts its result -//! back to the canonical range `[0, q)` exactly once, at the end. -//! -//! Every bound below is stated for the inputs the kernels receive. Tests check the reductions -//! exhaustively over those domains, and debug builds check every addition for overflow. +//! Coefficients stay partially reduced as `i16` and are reduced only with Montgomery and Barrett +//! (multiply, shift, add), so there are no secret-dependent selects. Kernels return `[0, q)`. use super::{Elem, GAMMA, NttPolynomial, NttVector, ZETA_POW_BITREV}; use array::{Array, ArraySize, typenum::U256}; -/// The modulus `q`. pub(super) const Q: i16 = 3329; -/// `q^-1 mod 2^16`, as a signed value. +/// `q^-1 mod 2^16`. const QINV: i16 = -3327; -/// The Montgomery radix `R = 2^16`, reduced mod `q`. +/// `R = 2^16 mod q`. #[allow(clippy::integer_division_remainder_used, reason = "constant")] const R_MOD_Q: i32 = (1 << 16) % Q as i32; -/// `R^2 mod q`. Montgomery multiplication by this constant maps `x` to `x * R mod q`. +/// `R^2 mod q`. #[allow(clippy::integer_division_remainder_used, reason = "constant")] const R2: i16 = centered((R_MOD_Q * R_MOD_Q) % Q as i32); -/// `R / 128 mod q`. Montgomery multiplication by this constant divides by 128, which is the -/// final scaling step of `NTT^-1`. +/// `R / 128 mod q`, the final scaling of `NTT^-1` in Montgomery form. #[allow(clippy::integer_division_remainder_used, reason = "constant")] const INV_NTT_SCALE: i16 = centered((3303 * R_MOD_Q) % Q as i32); -/// Map a value in `[0, q)` to its representative in `[-(q-1)/2, (q-1)/2]`. +/// Maps `[0, q)` to `[-(q-1)/2, (q-1)/2]`. #[allow(clippy::cast_possible_truncation, reason = "the value is below q")] #[allow(clippy::integer_division_remainder_used, reason = "constant")] const fn centered(x: i32) -> i16 { @@ -43,7 +34,7 @@ const fn centered(x: i32) -> i16 { } } -/// `ZETA_POW_BITREV` in the Montgomery domain: `zeta^BitRev7(i) * R mod q`, centered. +/// `ZETA_POW_BITREV[i] * R mod q`, centered. #[allow(clippy::integer_division_remainder_used, reason = "constant")] const ZETAS: [i16; 128] = { let mut out = [0; 128]; @@ -55,7 +46,6 @@ const ZETAS: [i16; 128] = { out }; -/// `qinv_multiple` of each of `ZETAS`. const ZETAS_QINV: [i16; 128] = { let mut out = [0; 128]; let mut i = 0; @@ -66,8 +56,7 @@ const ZETAS_QINV: [i16; 128] = { out }; -/// `GAMMA` scaled by `R^2`: `zeta^(2 BitRev7(i) + 1) * R^2 mod q`, centered. Montgomery -/// multiplication of `x` by `GAMMAS_R2[i]` gives `x * GAMMA[i] * R mod q`. +/// `GAMMA[i] * R^2 mod q`, centered. #[allow(clippy::integer_division_remainder_used, reason = "constant")] const GAMMAS_R2: [i16; 128] = { let r2 = (R_MOD_Q * R_MOD_Q) % Q as i32; @@ -80,7 +69,6 @@ const GAMMAS_R2: [i16; 128] = { out }; -/// `qinv_multiple` of each of `GAMMAS_R2`. const GAMMAS_R2_QINV: [i16; 128] = { let mut out = [0; 128]; let mut i = 0; @@ -91,7 +79,7 @@ const GAMMAS_R2_QINV: [i16; 128] = { out }; -/// Montgomery reduction: for `|a| < q * 2^15`, returns `r ≡ a * R^-1 (mod q)` with `|r| < q`. +/// `a * R^-1 mod q` in `(-q, q)`, for `|a| < q * 2^15`. #[inline(always)] #[allow( clippy::cast_possible_truncation, @@ -102,14 +90,13 @@ fn montgomery_reduce(a: i32) -> i16 { ((a - i32::from(t) * i32::from(Q)) >> 16) as i16 } -/// Montgomery multiplication: for `|a * b| < q * 2^15`, returns `r ≡ a * b * R^-1 (mod q)` with -/// `|r| < q`. Reference for [`fqmul_const`]. +/// `a * b * R^-1 mod q`. Test reference for [`fqmul_const`]. #[cfg(test)] fn fqmul(a: i16, b: i16) -> i16 { montgomery_reduce(i32::from(a) * i32::from(b)) } -/// The high 16 bits of the product of `a` and `b`. +/// High 16 bits of `a * b`. #[inline(always)] #[allow( clippy::cast_possible_truncation, @@ -119,30 +106,20 @@ fn mulhi(a: i16, b: i16) -> i16 { ((i32::from(a) * i32::from(b)) >> 16) as i16 } -/// `b * q^-1 mod 2^16`, the second constant that [`fqmul_const`] needs for a multiplier `b`. +/// `b * q^-1 mod 2^16`. const fn qinv_multiple(b: i16) -> i16 { b.wrapping_mul(QINV) } -/// Montgomery multiplication by a constant `b`, given `b_qinv = qinv_multiple(b)`. The result is -/// exactly `fqmul(a, b)`. -/// -/// Montgomery reduction subtracts `t * q`, where `t = a * b * q^-1 mod 2^16`, from `a * b`. The -/// low 16 bits of the two products are equal, so the result is the difference of their high -/// halves (Seiler, ePrint 2018/039). Computed this way, every operation works on 16-bit lanes, -/// which compilers map onto SIMD instructions such as `pmulhw` and `pmullw`. +/// `fqmul(a, b)` using only 16-bit operations, given `b_qinv = qinv_multiple(b)` +/// (Seiler, ePrint 2018/039). Vectorizes to `pmulhw`/`pmullw`. #[inline(always)] fn fqmul_const(a: i16, b: i16, b_qinv: i16) -> i16 { mulhi(a, b) - mulhi(a.wrapping_mul(b_qinv), Q) } -/// Barrett reduction: for any `a`, returns the representative of `a mod q` in -/// `[-(q-1)/2, (q-1)/2]`. -/// -/// The quotient `round(a * V / 2^26)`, with `V = round(2^26 / q)`, is computed from the high half -/// of `a * V` as `(mulhi(a, V) + 2^9) >> 10`, which is equal because nested floors of divisions -/// by integers collapse. The quotient has absolute value at most 10, and the result fits in -/// `i16`, so the wrapping operations below are exact. +/// `a mod q` in `[-(q-1)/2, (q-1)/2]`. The wrapping operations are exact because the result fits +/// in `i16`. #[inline(always)] #[allow(clippy::cast_possible_truncation, reason = "V is below 2^15")] fn barrett_reduce(a: i16) -> i16 { @@ -152,16 +129,10 @@ fn barrett_reduce(a: i16) -> i16 { a.wrapping_sub(t.wrapping_mul(Q)) } -/// Map coefficients in `[-q, q)` to their canonical representatives in `[0, q)`. +/// Maps `[-q, q)` to `[0, q)`. /// -/// Adding `q` to the negative coefficients is done with an arithmetic sign mask. LLVM recognizes -/// `(a >> 15) & q` as a select on the sign of `a`, and on targets without conditional moves (for -/// example `thumbv6m`) it lowers that select to a branch on the secret coefficient. The masks for -/// the whole polynomial are therefore passed through a single optimization barrier before they are -/// applied, which hides their origin while keeping both loops vectorizable. -/// -/// Every kernel ends with this pass, so it is kept out of line: inlining it everywhere costs about -/// 10 KiB of code on thumbv7em and is not measurably faster. +/// The barrier keeps LLVM from turning the sign mask into a branch on secret data, which it does +/// on targets without conditional moves (thumbv6m). Out of line to save code size. #[inline(never)] #[allow(clippy::cast_sign_loss, reason = "the result is in [0, q)")] pub(super) fn canonicalize(f: &[i16; 256]) -> Array { @@ -170,32 +141,27 @@ pub(super) fn canonicalize(f: &[i16; 256]) -> Array { Array::from_fn(|i| Elem::new((f[i] + (masks[i] & Q)) as u16)) } -/// Coefficient-wise `a + b`, for coefficients in `[0, q)`. +/// `a + b` for canonical inputs. pub(super) fn add(a: &Array, b: &Array) -> Array { - // `a + b - q` is in `[-q, q - 2]`. + // In `[-q, q - 2]`. let f: [i16; 256] = core::array::from_fn(|i| signed(a[i]) + signed(b[i]) - Q); canonicalize(&f) } -/// Coefficient-wise `a - b`, for coefficients in `[0, q)`. +/// `a - b` for canonical inputs. pub(super) fn sub(a: &Array, b: &Array) -> Array { - // `a - b` is in `(-q, q)`. + // In `(-q, q)`. let f: [i16; 256] = core::array::from_fn(|i| signed(a[i]) - signed(b[i])); canonicalize(&f) } -/// Read a canonical coefficient in `[0, q)` as `i16`. #[inline(always)] #[allow(clippy::cast_possible_wrap, reason = "the value is below q")] fn signed(x: Elem) -> i16 { x.0 as i16 } -/// One layer of the forward NTT (Algorithm 9). -/// -/// Each layer adds at most `q` to the magnitude of a coefficient, since `fqmul` returns values -/// below `q` in absolute value. Inputs below `q` therefore stay below `8q < 2^15` after seven -/// layers, and no intermediate reduction is needed. +/// One NTT layer. Each layer grows values by less than `q`, so they stay below `8q < 2^15`. #[inline(always)] fn ntt_layer(f: &mut [i16; 256], k: &mut usize) { for block in f.chunks_exact_mut(2 * LEN).take(ITERATIONS) { @@ -211,7 +177,7 @@ fn ntt_layer(f: &mut [i16; 256], k: & } } -/// Algorithm 9: `NTT`, on coefficients in `[0, q)`. Returns coefficients in `[0, q)`. +/// Algorithm 9: `NTT`. pub(super) fn ntt(input: &Array) -> NttPolynomial { let mut f: [i16; 256] = core::array::from_fn(|i| signed(input[i])); let mut k = 1; @@ -230,10 +196,7 @@ pub(super) fn ntt(input: &Array) -> NttPolynomial { NttPolynomial::new(canonicalize(&f)) } -/// One layer of the inverse NTT (Algorithm 10). -/// -/// Inputs are below `q` in absolute value. Sums and differences are below `2q`. The sum is -/// Barrett-reduced and the difference is Montgomery-multiplied, so outputs are again below `q`. +/// One inverse NTT layer. Values stay below `q` in absolute value. #[inline(always)] fn ntt_inverse_layer(f: &mut [i16; 256], k: &mut usize) { for block in f.chunks_exact_mut(2 * LEN).take(ITERATIONS) { @@ -249,7 +212,7 @@ fn ntt_inverse_layer(f: &mut [i16; 25 } } -/// Algorithm 10: `NTT^-1`, on coefficients in `[0, q)`. Returns coefficients in `[0, q)`. +/// Algorithm 10: `NTT^-1`. pub(super) fn ntt_inverse(input: &Array) -> Array { let mut f: [i16; 256] = core::array::from_fn(|i| signed(input[i])); let mut k = 127; @@ -262,20 +225,15 @@ pub(super) fn ntt_inverse(input: &Array) -> Array { ntt_inverse_layer::<64, 2>(&mut f, &mut k); ntt_inverse_layer::<128, 1>(&mut f, &mut k); - // Multiply by 128^-1. The result is below q in absolute value. + // Scale by 128^-1. for x in &mut f { *x = fqmul_const(*x, INV_NTT_SCALE, qinv_multiple(INV_NTT_SCALE)); } canonicalize(&f) } -/// The right-hand operand of `MultiplyNTTs` (Algorithm 11), prepared once so that it can be -/// multiplied by several left-hand operands. -/// -/// For each degree-one factor `b0 + b1 X` this stores `b0 * R`, `b1 * R` and `b1 * gamma * R` -/// (mod q, below q in absolute value). Products with canonical coefficients then need a single -/// Montgomery reduction, which removes the factor `R`, and `b1 * gamma` is computed once rather -/// than once per left-hand operand. +/// Right operand of Algorithm 11 as `b0 * R`, `b1 * R` and `b1 * gamma * R`, computed once and +/// reused across matrix rows. #[derive(Clone)] pub(crate) struct PreparedNtt { b0: [i16; 128], @@ -302,18 +260,13 @@ impl PreparedNtt { } } -/// Prepare every polynomial of a vector (see [`PreparedNtt`]). pub(crate) fn prepare(v: &NttVector) -> Array { Array::from_fn(|i| PreparedNtt::new(&v.0[i])) } -/// Algorithms 11 and 12 (`MultiplyNTTs` and `BaseCaseMultiply`), summed over a vector: -/// `sum_j a[j] * b[j]` in `T_q`. +/// Algorithms 11 and 12 summed over a vector: `sum_j a[j] * b[j]`. /// -/// Products are accumulated in `i32` and reduced once per output coefficient. Coefficients of `a` -/// are in `[0, q)` and prepared values are below `q` in absolute value, so each output -/// coefficient gains less than `2 q^2` per term. With `K <= 4` terms the sum stays below -/// `8 q^2 < q * 2^15`, which is the input domain of `montgomery_reduce`. +/// One reduction per coefficient: with `K <= 4` the `i32` sum stays below `8 q^2 < q * 2^15`. pub(crate) fn multiply_accumulate( a: &NttVector, b: &Array, @@ -322,8 +275,7 @@ pub(crate) fn multiply_accumulate( accumulate(a.0.iter().zip(b.iter())) } -/// The sum of the products of `(a, b)` pairs in `T_q`, for at most four pairs. See -/// [`multiply_accumulate`] for the bounds. +/// Sum of products of at most four `(a, b)` pairs. fn accumulate<'a>( terms: impl Iterator, ) -> NttPolynomial { @@ -347,7 +299,7 @@ fn accumulate<'a>( NttPolynomial::new(canonicalize(&out)) } -/// Multiply a `K x K` matrix, given by its rows, by a prepared vector in `T_q`. +/// `A * v`, for `A` given by rows. pub(crate) fn matrix_multiply( rows: &Array, K>, v: &Array, @@ -355,7 +307,7 @@ pub(crate) fn matrix_multiply( NttVector::new(Array::from_fn(|i| multiply_accumulate(&rows[i], v))) } -/// Multiply the transpose of a `K x K` matrix, given by its rows, by a prepared vector in `T_q`. +/// `A^T * v`, for `A` given by rows. pub(crate) fn matrix_transpose_multiply( rows: &Array, K>, v: &Array, @@ -402,8 +354,15 @@ mod tests { } } - /// `montgomery_reduce` over its whole documented domain, with a stride that keeps the test - /// fast in debug builds, plus both ends of the domain. + #[test] + fn centered_boundaries() { + assert_eq!(centered(0), 0); + assert_eq!(centered(1664), 1664); + assert_eq!(centered(1665), -1664); + assert_eq!(centered(3328), -1); + } + + /// Strided over the whole domain, plus both ends. #[test] fn montgomery_reduce_domain() { let bound = i32::from(Q) << 15; @@ -422,7 +381,7 @@ mod tests { } } - /// `fqmul` for every multiplier used by the kernels and every input magnitude they can see. + /// Every kernel constant, every input below `8q`. #[test] fn fqmul_all_constants() { let constants = ZETAS @@ -465,7 +424,6 @@ mod tests { } } - /// `add` and `sub` agree with the `Elem` operators for every pair of field elements. #[test] fn add_sub_exhaustive() { for a in 0..BaseField::Q { @@ -483,7 +441,7 @@ mod tests { } } - /// The largest possible accumulated product stays inside the Montgomery domain. + /// Largest possible sums stay inside the Montgomery domain. #[test] fn multiply_accumulate_extremes() { use array::typenum::U4; diff --git a/ml-kem/src/pke.rs b/ml-kem/src/pke.rs index cb66803b..c3a48618 100644 --- a/ml-kem/src/pke.rs +++ b/ml-kem/src/pke.rs @@ -134,8 +134,7 @@ where /// An `EncryptionKey` provides the ability to encrypt a value so that it can only be /// decrypted by the holder of the corresponding decapsulation key. /// -/// There is deliberately no `Default` impl: with the `alloc` feature a default key would pair -/// `rho = 0` with an all-zero matrix instead of the matrix expanded from `rho`. +/// No `Default`: the cached matrix must match `rho`. #[derive(Clone)] pub(crate) struct EncryptionKey

where @@ -144,14 +143,12 @@ where t_hat: NttVector, rho: B32, - /// The matrix `A_hat` expanded from `rho`. Expanding it runs `SampleNTT` `k^2` times, which - /// is the most expensive step of `K-PKE.Encrypt`, so keys that live on the heap keep a copy. - /// It is public and fully determined by `rho`. + /// Public `A_hat` expanded from `rho`, cached because expanding it dominates `encrypt`. #[cfg(feature = "alloc")] A_hat: NttMatrix, } -// Handwritten to leave out `A_hat`, which is determined by `rho` +// Leaves out `A_hat`, which follows from `rho` impl

fmt::Debug for EncryptionKey

where P: PkeParams, @@ -164,7 +161,7 @@ where } } -// Handwritten to leave out `A_hat`, which is determined by `rho` +// Leaves out `A_hat`, which follows from `rho` impl

PartialEq for EncryptionKey

where P: PkeParams, @@ -183,8 +180,7 @@ where /// Encrypt the specified message for the holder of the corresponding decryption key, using the /// provided randomness, according the `K-PKE.Encrypt` procedure. pub(crate) fn encrypt(&self, message: &B32, randomness: &B32) -> Ciphertext

{ - // Expand the public matrix before sampling anything secret, in the order of Algorithm 14 - // (steps 4-8, then 9-17). + // Algorithm 14 order: the matrix (steps 4-8) before the secrets (steps 9-17). #[cfg(feature = "alloc")] let A_hat = &self.A_hat; #[cfg(not(feature = "alloc"))] @@ -253,7 +249,7 @@ where // #1 is performed by the `EncodedEncryptionKey` type, and the following check vicariously // performs #2 by encoding the integer-mod-q array using our implementation of ByteEncode₁₂ // and comparing the resulting serialization to see if it round-trips. It runs before - // `A_hat` is expanded, so that invalid keys are rejected cheaply. + // expanding `A_hat`, so invalid keys are rejected cheaply. if P::encode_u12(&t_hat) != *t_hat_enc { return Err(InvalidKey); } @@ -349,4 +345,33 @@ mod test { key_inequality_test::(); key_inequality_test::(); } + + /// `PartialEq` compares `t_hat` and `rho`. + #[test] + fn encryption_key_equality() { + let (_, ek) = DecryptionKey::::generate(&B32::from([1; 32])); + assert_eq!(ek, ek.clone()); + + let mut other_rho = ek.clone(); + other_rho.rho[0] ^= 1; + assert_ne!(ek, other_rho); + + let mut other_t_hat = ek.clone(); + other_t_hat.t_hat.0[0].0[0].0 ^= 1; + assert_ne!(ek, other_t_hat); + } + + /// `Debug` matches the former derived output and omits the cache. + #[test] + fn encryption_key_debug() { + extern crate alloc; + use alloc::format; + + let (_, ek) = DecryptionKey::::generate(&B32::from([1; 32])); + let expected = format!( + "EncryptionKey {{ t_hat: {:?}, rho: {:?} }}", + ek.t_hat, ek.rho + ); + assert_eq!(format!("{ek:?}"), expected); + } } diff --git a/ml-kem/tests/accumulated.rs b/ml-kem/tests/accumulated.rs index 68aecd55..25f7832d 100644 --- a/ml-kem/tests/accumulated.rs +++ b/ml-kem/tests/accumulated.rs @@ -1,22 +1,11 @@ -//! Accumulated test vectors, using the procedure of the C2SP CCTV "accumulated" vectors: +//! Accumulated test vectors, following the C2SP CCTV procedure: //! //! -//! A SHAKE-128 instance with empty input provides `d`, `z`, `m` and a random ciphertext for each -//! test. The encapsulation key, the expanded decapsulation key, the ciphertext, the shared key from -//! encapsulation and the shared key from decapsulating the random ciphertext are absorbed into a -//! second SHAKE-128 instance, whose first 32 bytes of output are compared to the expected value. +//! One SHAKE-128 stream supplies the inputs and a second one hashes all outputs. CCTV's published +//! hashes target the FIPS 203 draft (`G(d)` instead of `G(d || k)`), so these values are for final +//! FIPS 203, computed with libcrux-ml-kem 0.0.11. //! -//! The hashes published by CCTV target the FIPS 203 initial public draft, whose K-PKE.KeyGen -//! derives its seeds as `G(d)` rather than `G(d || k)`, so they cannot match the final standard. -//! The expected values below were computed for FIPS 203 with libcrux-ml-kem 0.0.11. ml-kem 0.3.4 -//! matches the 100 and 10,000 test runs. -//! -//! The 100-test run is part of the normal test suite. The longer runs exercise the rare paths of -//! `SampleNTT` that need extra XOF blocks and are ignored by default: -//! -//! ```text -//! cargo test --release --test accumulated -- --ignored -//! ``` +//! Longer runs are ignored by default: `cargo test --release --test accumulated -- --ignored` #![allow(clippy::unwrap_used, missing_docs)]