Mikeo/lowmem improvements

CLAUDE.md

@@ -28,8 +28,9 @@ fixed.
 
 ## Toolchain
 
-- Uses Rust **nightly** (pinned in `rust-toolchain.toml`) — `core/src/lib.rs` uses `#![feature(adt_const_params)]`.
-- 2024 edition (set workspace-wide in the root `Cargo.toml`).
+- Builds on Rust **stable**: there is no toolchain pin, and no crate enables a `#![feature(...)]`
+  gate, so nightly-only tooling (`-Z` flags and the like) is not available.
+- 2024 edition (set workspace-wide in the root `Cargo.toml`), which needs Rust 1.85 or later.
 
 ## Common commands
 
@@ -108,6 +109,13 @@ Repo mechanics behind those rules, which the documents don't spell out:
 - **CLI commands stream.** The `cli/` binary is stdin→stdout with ~1 KB buffers so commands compose in shell
   pipelines; preserve that when adding subcommands.
 - Trait → factory → CLI is the wiring path for a new primitive; see [the workspace architecture](#the-core--core-test-framework--factory-spine) above for the crates involved.
+- **Comments describe the code that is there, not the road that led to it.** Do not add a comment
+  explaining a transient design decision — an approach that was tried and abandoned, what an earlier
+  version did, why one formulation was chosen over another that is no longer present — or describing
+  a design the code does not use. Such comments are noise: they age badly, and a reader has to work
+  out that they describe nothing in front of them. A comment that explains why the present code is
+  the way it is, and that a naive edit would break it (a spec step, an invariant, a constraint), is
+  wanted; a comment narrating how it got that way is not.
 
 ## Working from specifications
 

alpha_0.1.3_release_notes.md

@@ -7,3 +7,4 @@
 * bug fixes to the way SHA3/SHAKE handled absorbing and squeezing a partial final byte.
 * Design discussions about whether core::traits::XOF (in the abstract) should allow interleaving absorb -> squeeze ->
   absorb (ie "absorb-after-squeeze). Outcome: absorb-after-squeeze forbidden. Could be changed in the future.
+* Further reductions to the memory usage of the mldsa-lowmemory and mlkem-lowmemory crates. 

crypto/mldsa-lowmemory/src/aux_functions.rs

@@ -6,7 +6,7 @@ use crate::params::{
     GAMMA1_2_POW_17, GAMMA1_2_POW_19, GAMMA2_Q_MINUS_1_OVER_32, GAMMA2_Q_MINUS_1_OVER_88,
     MLDSAParams,
 };
-use crate::polynomial::Polynomial;
+use crate::polynomial::{HintRow, Polynomial, ZEROED_HINT_ROW, hint_set};
 use bouncycastle_core::traits::XOF;
 use bouncycastle_utils::secret::ZeroizablePrimitive;
 
@@ -358,13 +358,15 @@ pub(crate) fn unpack_z_row<P: MLDSAParams, const SIG_LEN: usize>(
     if z.check_norm(P::gamma1_minus_beta) { Err(()) } else { Ok(z) }
 }
 /// Part of unpacking the sig value
+///
+/// Returns the decoded row, or `None` if the encoded hint is malformed.
 pub(crate) fn unpack_h_row<P: MLDSAParams, const SIG_LEN: usize>(
     row: usize,
     sig: &[u8; SIG_LEN],
-) -> Option<Polynomial> {
+) -> Option<HintRow> {
     debug_assert!(row < P::k);
 
-    let mut h = Polynomial::new();
+    let mut out = ZEROED_HINT_ROW;
 
     // skip over the other stuff in the encoded sig value
     let pos = P::C_TILDE_LEN + P::l * P::POLY_Z_PACKED_LEN;
@@ -401,7 +403,7 @@ pub(crate) fn unpack_h_row<P: MLDSAParams, const SIG_LEN: usize>(
             return None;
         }
         // 12: 𝐡[𝑖]_𝑦[Index] ← 1
-        h[sig[pos + j] as usize] = 1;
+        hint_set(&mut out, sig[pos + j] as usize);
 
         // 13: Index ← Index + 1
         //  > done by for loop
@@ -418,7 +420,7 @@ pub(crate) fn unpack_h_row<P: MLDSAParams, const SIG_LEN: usize>(
         }
     }
 
-    Some(h)
+    Some(out)
 }
 
 /// Algorithm 29 SampleInBall(𝜌)
@@ -544,12 +546,22 @@ pub(crate) fn rej_bounded_poly<P: MLDSAParams>(rho: &[u8; 64], nonce: &[u8; 2])
     h.absorb(rho).expect("absorb before squeeze is infallible");
     h.absorb(nonce).expect("absorb before squeeze is infallible");
 
-    // SHAKE is fairly inefficient if only 3 bytes are squeezed at a time, so the implementation does a block instead.
-    // size is not a limitation as long as it is a multiple of 3.
-    // 312 seems to be the sweet spot after some experimentation
-    // which is possibly also related with the average rejection rate.
-    // Also, 312 is a multiple of 8 (efficient for SHAKE)
-    let mut z_arr = [0u8; 312];
+    // Deviation from FIPS 204, Algorithm 31 step 5, which squeezes one byte per loop iteration:
+    // H is SHAKE256, which produces a whole 136-byte block per Keccak permutation, so squeezing a
+    // byte at a time wastes most of each block. The squeeze is buffered instead, and the refill
+    // below makes the byte stream — and therefore the output — identical to the spec's.
+    //
+    // 272 is exactly two SHAKE256 blocks (2 × 136), so filling the buffer costs two permutations
+    // with nothing stranded in the sponge's output queue, and it covers the whole polynomial in a
+    // single squeeze almost always. Per FIPS 204 §C, each iteration consumes one byte and yields
+    // Binomial(2, θ) coefficients, θ = 15/16 for η = 2 and 9/16 for η = 4. The worst case is
+    // η = 4 (ML-DSA-65): 228 bytes needed on average, and over 300k simulated seeds the largest
+    // requirement was 276 bytes, so the refill runs for roughly 1 seed in 100,000. For η = 2
+    // (ML-DSA-44/87) it is 137 bytes on average and never exceeded 150.
+    //
+    // This is a buffer, not the iteration cap of FIPS 204 Table 3 (481 bytes for RejBoundedPoly):
+    // the loop refills rather than giving up, so no cap is imposed.
+    let mut z_arr = [0u8; 272];
     h.squeeze_out(&mut z_arr);
     let mut idx: usize = 0;
 

crypto/mldsa-lowmemory/src/lib.rs

@@ -82,38 +82,37 @@
 //!
 //!
 //! ## Algorithm Peak Memory Usage
-//! The table below shows peak memory usage of the ML-DSA algorithms and the rough performanc (throughput) impact.
-//!
-//! Measuring peak application memory usage can be a bit tricky, and the numbers obtained depend heavily on how the
-//! measurement harness is designed. Here, we aim to provide a conservative measurement, meaning that we are aiming for an
-//! over-estimate so that any deployment within an existing application will use incrementally less additional memory
-//! than the amount stated here.
 //!
-//! Our measurement methodology is to compile a simple standalone HelloWorld application that only calls the function under test
-//! with as minimal as possible hard-coded data (such as keys or ciphertexts) and measure the peak memory usage of running
-//! the compiled binary using `valgrind --tool=massif --heap=no --stack=yes`. The flags for heap and stack
-//! reflect the fact that this is a `no_std` rust application and therefore the cryptographic functions use no heap memory.
-//! The measurements may over-estimate by as much as 3 kb since that that's the measured peak memory usage of a do-nothing
-//! HelloWorld rust application.
+//! The table below shows peak memory usage of the ML-DSA algorithms and the rough performanc (throughput) impact.
 //!
-//! | Algorithm                 | Peak swap memory usage (kB) | Throughput (ops/s)  |
+//! | Algorithm                 | Peak stack memory usage (kB) | Throughput (ops/s)  |
 //! |---------------------------|-----------------------------|---------------------|
-//! | MLDSA44_lowmemory/KeyGen  | 12.6 (113.8)                | 11,800     (11,300) |
-//! | MLDSA65_lowmemory/KeyGen  | 15.0 (124.1)                | 5,500      (7.000)  |
-//! | MLDSA87_lowmemory/KeyGen  | 15.2 (197.8)                | 3,300      (4,200)  |
-//! | MLDSA44_lowmemory/Sign    | 24.8 (117.7)                | 850        (4,000)  |
-//! | MLDSA65_lowmemory/Sign    | 28.2 (159.6)                | 580        (2,900)  |
-//! | MLDSA87_lowmemory/Sign    | 31.1 (236.7)                | 315        (2,000)  |
-//! | MLDSA44_lowmemory/Verify  | 17.1 (73.0)                 | 10,100     (14,000) |
-//! | MLDSA65_lowmemory/Verify  | 18.0 (134.4)                | 6,300      (8,400)  |
-//! | MLDSA87_lowmemory/Verify  | 20.6 (211.6)                | 3,500      (5,000)  |
+//! | MLDSA44_lowmemory/KeyGen  | 12.3 (62.3)                 | 11,700     (11,400) |
+//! | MLDSA65_lowmemory/KeyGen  | 14.8 (93.9)                 | 6,800      (7,000)  |
+//! | MLDSA87_lowmemory/KeyGen  | 16.8 (141.4)                | 3,900      (4,200)  |
+//! | MLDSA44_lowmemory/Sign    | 22.3 (84.0)                 | 990        (3,400)  |
+//! | MLDSA65_lowmemory/Sign    | 27.9 (124.7)                | 490        (2,400)  |
+//! | MLDSA87_lowmemory/Sign    | 31.6 (182.6)                | 330        (1,700)  |
+//! | MLDSA44_lowmemory/Verify  | 15.2 (52.7)                 | 11,200     (13,200) |
+//! | MLDSA65_lowmemory/Verify  | 19.1 (81.3)                 | 6,400      (8,000)  |
+//! | MLDSA87_lowmemory/Verify  | 22.3 (123.7)                | 3,700      (4,800)  |
 //!
 //! Values in parentheses are the comparison values from the un-optimized implementation in the \[bouncycastle_mldsa] crate.
-//! Size numbers were collected with valgrind using a simple main program that calls only the measured function.
-//! Performance throughput numbers were collected on my laptop using the library's provided benchmarks, so
-//! performance they should be taken with an extreme grain of salt.
 //!
+//! **Caveates**
+//!
+//! Throughput numbers are meant to show relative difference between the two implementations and not be
+//! absolute performance measurements.
+//!
+//! Size numbers were collected with valgrind using a simple main program that calls only the measured function.
+//! Measurements include the public key and signature buffer, which is a realistic setting, but
+//! numbers would be lower if those were on the heap and excluded from the measurement.
 //! Actual values may vary based on build configuration and target architecture.
+//! The peak also depends on how the compiler lays out the measured function's stack frame relative
+//! to its caller and whether it elides copies at the call boundary, which is outside
+//! the library's control. The memory benchmarking framework included in the library attempts to
+//! This effect is small for the low-memory crates (well under 1 kB) and larger for the full implementations
+//! in parentheses (a few kB), so treat the latter as approximate.
 //!
 //! # Usage
 //!

crypto/mldsa-lowmemory/src/low_memory_helpers.rs

@@ -5,7 +5,6 @@
 use crate::aux_functions::{bit_unpack_eta_out, expand_mask_poly, rej_ntt_poly, unpack_z_row};
 use crate::params::MLDSAParams;
 use crate::polynomial::Polynomial;
-use bouncycastle_core::errors::SignatureError;
 use bouncycastle_utils::secret::{Secret, ZeroizablePrimitive};
 
 #[inline(always)]
@@ -43,8 +42,8 @@ pub(crate) fn compute_w_row<P: MLDSAParams>(
 pub(crate) fn compute_wp_approx_row<P: MLDSAParams, const SIG_LEN: usize>(
     rho: &[u8; 32],
     sig: &[u8; SIG_LEN],
-    t1: &Polynomial,
-    c: &Polynomial,
+    t1: Polynomial,
+    c_hat: &Polynomial,
     idx: usize,
 ) -> Result<Polynomial, ()> {
     // Algorithm 8: line 9: 𝐰′_approx ← NTT−1(𝐀_hat ∘ NTT(𝐳) − NTT(𝑐) ∘ NTT(𝐭1 ⋅ 2^𝑑))
@@ -72,15 +71,15 @@ pub(crate) fn compute_wp_approx_row<P: MLDSAParams, const SIG_LEN: usize>(
         Az_acc.add_ntt(&tmp);
     }
 
-    let ct1 = compute_ct1(t1.clone(), c.clone());
-    fn compute_ct1(mut t1_i: Polynomial, mut c: Polynomial) -> Polynomial {
-        t1_i.shift_left_d();
-        t1_i.ntt();
-        c.ntt();
-        t1_i.multiply_ntt(&c);
+    // NTT(𝑐) ∘ NTT(𝐭1 ⋅ 2^𝑑), computed in place in the buffer `t1` arrived in.
+    let ct1 = {
+        let mut ct1 = t1;
+        ct1.shift_left_d();
+        ct1.ntt();
+        ct1.multiply_ntt(c_hat);
 
-        t1_i
-    }
+        ct1
+    };
 
     Az_acc.sub(&ct1);
     Az_acc.inv_ntt();
@@ -89,34 +88,39 @@ pub(crate) fn compute_wp_approx_row<P: MLDSAParams, const SIG_LEN: usize>(
     Ok(Az_acc)
 }
 
+/// 𝐳 is written into `z_out`, and `true` returned. `false` means the norm check rejected this
+/// component, and `z_out` then holds a partial value that the caller must discard.
 pub(crate) fn compute_z_component<P: MLDSAParams>(
-    s1: &Polynomial,
+    s1: Polynomial,
     rho_p_p: &[u8; 64],
     c_hat: &Polynomial,
     kappa: u16,
     col: usize,
-) -> Result<Option<Polynomial>, SignatureError> {
+    z_out: &mut Polynomial,
+) -> bool {
     let y = expand_mask_poly::<P>(rho_p_p, kappa + col as u16);
-    let mut s1_hat = s1.clone();
-    s1_hat.ntt();
-    s1_hat.multiply_ntt(c_hat);
-    let mut cs1 = s1_hat; // rename
-    cs1.inv_ntt();
-    let mut z = cs1;
-    z.add_ntt(&y);
-
-    if z.check_norm(P::gamma1_minus_beta) { Ok(None) } else { Ok(Some(z)) }
+
+    // 𝑐𝐬1 ← NTT−1(𝑐_hat ∘ NTT(𝐬1)), built in place in the caller's buffer.
+    *z_out = s1;
+    z_out.ntt();
+    z_out.multiply_ntt(c_hat);
+    z_out.inv_ntt();
+
+    // 𝐳 ← 𝐲 + 𝑐𝐬1
+    z_out.add_ntt(&y);
+
+    !z_out.check_norm(P::gamma1_minus_beta)
 }
 
 pub(crate) fn compute_w0cs2_component<P: MLDSAParams>(
-    s2: &Polynomial,
+    s2: Polynomial,
     w: &Polynomial,
     c_hat: &Polynomial,
-) -> Option<Polynomial> {
-    let mut s2_hat = s2.clone();
-    s2_hat.ntt();
-    s2_hat.multiply_ntt(c_hat);
-    let mut cs2 = s2_hat; // rename
+    w0cs2_out: &mut Polynomial,
+) -> bool {
+    let mut cs2 = s2;
+    cs2.ntt();
+    cs2.multiply_ntt(c_hat);
     cs2.inv_ntt();
 
     //  Note: this could be further optimized by using the optimization described in
@@ -125,23 +129,25 @@ pub(crate) fn compute_w0cs2_component<P: MLDSAParams>(
     //      and checking whether ‖r0‖∞ < γ2 − β and r1 = w1, it is equivalent to just check that
     //      ‖w0 − cs2‖∞ < γ2 − β, where w0 is the low part of w. If this check passes, w0 − cs2
     //      is the low part of w − cs2."
-    let mut w0cs2 = w.clone();
-    w0cs2.low_bits::<P>();
-    w0cs2.sub(&cs2);
-    if w0cs2.check_norm(P::gamma2_minus_beta) { None } else { Some(w0cs2) }
+    // `w` is still needed by the caller, so its low half is taken in the out-buffer.
+    *w0cs2_out = *w;
+    w0cs2_out.low_bits::<P>();
+    w0cs2_out.sub(&cs2);
+
+    !w0cs2_out.check_norm(P::gamma2_minus_beta)
 }
 
 pub(crate) fn compute_ct0_component<P: MLDSAParams>(
-    t0_row: &Polynomial,
+    t0_row: Polynomial,
     c_hat: &Polynomial,
-) -> Option<Polynomial> {
-    let mut t0_hat = t0_row.clone();
-    t0_hat.ntt();
-    t0_hat.multiply_ntt(c_hat);
-    let mut ct0 = t0_hat; // rename
-    ct0.inv_ntt();
-
-    if ct0.check_norm(P::gamma2) { None } else { Some(ct0) }
+    ct0_out: &mut Polynomial,
+) -> bool {
+    *ct0_out = t0_row;
+    ct0_out.ntt();
+    ct0_out.multiply_ntt(c_hat);
+    ct0_out.inv_ntt();
+
+    !ct0_out.check_norm(P::gamma2)
 }
 
 /// Unpack a single s value from the packed representation.