Cosmetic refactoring to MLDSA / MLKEM Polynomial and Matrix

crypto/mldsa-lowmemory/src/lib.rs

@@ -266,6 +266,3 @@ pub use mldsa::{MLDSA65_PK_LEN, MLDSA65_SIG_LEN, MLDSA65_SK_LEN};
 pub use mldsa::{MLDSA87_PK_LEN, MLDSA87_SIG_LEN, MLDSA87_SK_LEN};
 
 pub use mldsa::SUSPENDED_MU_BUILDER_STATE_LEN;
-
-// re-export just so it's visible to unit tests
-pub use polynomial::Polynomial;

crypto/mldsa-lowmemory/src/polynomial.rs

@@ -6,19 +6,13 @@ use core::ops::{Index, IndexMut};
 
 /// A polynomial over the ML-DSA ring.
 ///
-/// Dev note: The following structure does not necessarily need to be declared as public.
-/// There is no real scenario where this function needs to be called directly.
-/// However, in order to test the Debug and Display traits, it is necessary to use STD, so those
-/// can't be tested from inline tests in this file and the real unit tests are in a different crate.
-/// That's the reason why pub is used.
-///
 /// # 🚨 Security 🚨
 /// Polynomials themselves are not inherently secret since sometimes they are part of public keys
 /// and sometimes private keys.
 /// It is the responsibility of the caller to wrap sensitive instances in `Secret<Polynomial>`.
 /// Note: at the moment, nothing in this crate uses `Secret<Polynomial>`, so I have left the `impl ZeroizablePrimitive` commented-out.
 #[derive(Clone, Copy)]
-pub struct Polynomial {
+pub(crate) struct Polynomial {
     pub(crate) coeffs: [i32; N],
 }
 

crypto/mldsa/src/aux_functions.rs

@@ -407,7 +407,7 @@ pub(crate) fn sig_encode<
 
     for i in 0..l {
         output[pos..pos + POLY_Z_PACKED_LEN]
-            .copy_from_slice(&bitpack_gamma1::<POLY_Z_PACKED_LEN, GAMMA1>(&z.vec[i]));
+            .copy_from_slice(&bitpack_gamma1::<POLY_Z_PACKED_LEN, GAMMA1>(&z.elems[i]));
         pos += POLY_Z_PACKED_LEN;
     }
 
@@ -416,7 +416,7 @@ pub(crate) fn sig_encode<
     let mut m: usize = 0;
     for i in 0..k {
         for j in 0..N {
-            if h.vec[i][j] != 0 {
+            if h.elems[i][j] != 0 {
                 output[pos + m] = j as u8;
                 m += 1;
             }
@@ -453,7 +453,7 @@ pub(crate) fn sig_decode<
     pos += LAMBDA_over_4;
 
     for i in 0..l {
-        z.vec[i] = bit_unpack_gamma1::<GAMMA1>(&sig[pos..pos + POLY_Z_PACKED_LEN]);
+        z.elems[i] = bit_unpack_gamma1::<GAMMA1>(&sig[pos..pos + POLY_Z_PACKED_LEN]);
         pos += POLY_Z_PACKED_LEN;
     }
 
@@ -486,7 +486,7 @@ pub(crate) fn sig_decode<
                 return Err(());
             }
             // 12: 𝐡[𝑖]_𝑦[Index] ← 1
-            h.vec[i][sig[pos + j] as usize] = 1;
+            h.elems[i][sig[pos + j] as usize] = 1;
 
             // 13: Index ← Index + 1
             //  > done by for loop
@@ -672,7 +672,7 @@ pub(crate) fn expandA<const k: usize, const l: usize>(rho: &[u8; 32]) -> Matrix<
 
     for r in 0..k {
         for s in 0..l {
-            A_hat[r][s] = rej_ntt_poly(rho, &[s as u8, r as u8]);
+            A_hat.elems[r][s] = rej_ntt_poly(rho, &[s as u8, r as u8]);
         }
     }
 
@@ -692,11 +692,11 @@ pub(crate) fn expandS<const k: usize, const l: usize, const ETA: usize>(
     let mut s2: Secret<Vector<k>> = Secret::new();
 
     for r in 0..l {
-        s1.vec[r] = rej_bounded_poly::<ETA>(rho, &(r as u16).to_le_bytes());
+        s1.elems[r] = rej_bounded_poly::<ETA>(rho, &(r as u16).to_le_bytes());
     }
 
     for r in 0..k {
-        s2.vec[r] = rej_bounded_poly::<ETA>(rho, &(r as u16 + l as u16).to_le_bytes());
+        s2.elems[r] = rej_bounded_poly::<ETA>(rho, &(r as u16 + l as u16).to_le_bytes());
     }
 
     (s1, s2)
@@ -710,7 +710,7 @@ pub(crate) fn power_2_round_vec<const LEN: usize>(v: &Vector<LEN>) -> (Vector<LE
 
     for i in 0..LEN {
         for j in 0..N {
-            (r1.vec[i][j], r0.vec[i][j]) = power_2_round(v.vec[i][j]);
+            (r1.elems[i][j], r0.elems[i][j]) = power_2_round(v.elems[i][j]);
         }
     }
 
@@ -745,7 +745,7 @@ pub(crate) fn expand_mask<const l: usize, const GAMMA1: i32, const GAMMA1_MASK_L
         };
 
         // 5: 𝐲[𝑟] ← BitUnpack(𝑣, 𝛾1 − 1, 𝛾1)
-        y.vec[r] = bit_unpack_gamma1::<GAMMA1>(&v);
+        y.elems[r] = bit_unpack_gamma1::<GAMMA1>(&v);
     }
 
     y
@@ -881,8 +881,8 @@ pub(crate) fn make_hint_vecs<const k: usize, const GAMMA2: i32>(
     let mut count = 0i32;
 
     for i in 0..k {
-        let (w, c) = r.vec[i].make_hint::<GAMMA2>(&s.vec[i]);
-        out.vec[i] = w;
+        let (w, c) = r.elems[i].make_hint::<GAMMA2>(&s.elems[i]);
+        out.elems[i] = w;
 
         // mutants note: this chains up to hint_hamming_weight > OMEGA and there is no test KAT that triggers this branch
         count += c;
@@ -942,7 +942,7 @@ pub(crate) fn use_hint_vecs<const k: usize, const GAMMA2: i32>(
 ) -> Vector<k> {
     let mut out = Vector::<k>::new();
     for i in 0..k {
-        use_hint_polys::<GAMMA2>(&wp_approx.vec[i], &h.vec[i], &mut out.vec[i]);
+        use_hint_polys::<GAMMA2>(&wp_approx.elems[i], &h.elems[i], &mut out.elems[i]);
     }
 
     out

crypto/mldsa/src/lib.rs

@@ -188,6 +188,3 @@ pub use mldsa::{MLDSA87_PK_LEN, MLDSA87_SIG_LEN, MLDSA87_SK_LEN};
 pub use mldsa::SUSPENDED_MU_BUILDER_STATE_LEN;
 
 pub use matrix::Matrix;
-
-// re-export just so it's visible to unit tests
-pub use polynomial::Polynomial;

crypto/mldsa/src/matrix.rs

@@ -10,26 +10,14 @@ use core::ops::{Index, IndexMut};
 
 /// A matrix over the ML-DSA ring.
 #[derive(Clone)]
-pub struct Matrix<const k: usize, const l: usize>(/*pub(crate)*/ [[Polynomial; l]; k]);
-
-/// Convenience function to avoid ".0" all over the place.
-impl<const k: usize, const l: usize> Index<usize> for Matrix<k, l> {
-    type Output = [Polynomial; l];
-
-    fn index(&self, index: usize) -> &Self::Output {
-        &self.0[index]
-    }
-}
-/// Convenience function to avoid ".0" all over the place.
-impl<const k: usize, const l: usize> IndexMut<usize> for Matrix<k, l> {
-    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
-        &mut self.0[index]
-    }
+pub struct Matrix<const k: usize, const l: usize> {
+    /// Indexed `elems[row][col]`
+    pub(crate) elems: [[Polynomial; l]; k],
 }
 
 impl<const k: usize, const l: usize> Matrix<k, l> {
     pub(crate) fn new() -> Self {
-        Self { 0: [[(); l]; k].map(|_| [(); l].map(|_| Polynomial::new())) }
+        Self { elems: [[(); l]; k].map(|_| [(); l].map(|_| Polynomial::new())) }
     }
 
     /// Algorithm 48 MatrixVectorNTT(𝐌, 𝐯)
@@ -39,18 +27,18 @@ impl<const k: usize, const l: usize> Matrix<k, l> {
     /// Performs dot product multiplication of this matrix by a vector
     /// Input: vector of length l
     /// Output: vector of length k
-    pub fn matrix_vector_ntt(&self, v: &Vector<l>) -> Vector<k> {
+    pub(crate) fn matrix_vector_ntt(&self, v: &Vector<l>) -> Vector<k> {
         let mut w = Vector::<k>::new();
         for i in 0..k {
             // split out the 0 case to skip a no-op add_ntt()
-            w[i].coeffs.copy_from_slice(&multiply_ntt(&self[i][0], &v[0]).coeffs);
+            w[i].coeffs.copy_from_slice(&multiply_ntt(&self.elems[i][0], &v[0]).coeffs);
 
             let mut w1: Polynomial;
             for j in 1..l {
                 // dot product a vector into a matrix: multiply the input vector
                 // into each row of the matrix, then sum the results to produce a vector of
                 // length k.
-                w1 = multiply_ntt(&self[i][j], &v[j]);
+                w1 = multiply_ntt(&self.elems[i][j], &v[j]);
                 w[i].add_ntt(&w1);
             }
         }
@@ -61,21 +49,21 @@ impl<const k: usize, const l: usize> Matrix<k, l> {
 
 #[derive(Clone, Copy)]
 pub(crate) struct Vector<const LEN: usize> {
-    pub(crate) vec: [Polynomial; LEN],
+    pub(crate) elems: [Polynomial; LEN],
 }
 
 /// Convenience function to avoid ".0" all over the place.
 impl<const LEN: usize> Index<usize> for Vector<LEN> {
     type Output = Polynomial;
 
     fn index(&self, index: usize) -> &Self::Output {
-        &self.vec[index]
+        &self.elems[index]
     }
 }
 /// Convenience function to avoid ".0" all over the place.
 impl<const LEN: usize> IndexMut<usize> for Vector<LEN> {
     fn index_mut(&mut self, index: usize) -> &mut Self::Output {
-        &mut self.vec[index]
+        &mut self.elems[index]
     }
 }
 
@@ -85,7 +73,7 @@ impl<const LEN: usize> ZeroizablePrimitive for Vector<LEN> {
 
 impl<const LEN: usize> Vector<LEN> {
     pub(crate) const fn new() -> Self {
-        Self { vec: [Polynomial::new(); LEN] }
+        Self { elems: [Polynomial::new(); LEN] }
     }
 
     /// Algorithm 46 AddVectorNTT(𝐯, 𝐰)̂
@@ -176,7 +164,7 @@ impl<const LEN: usize> Vector<LEN> {
 
     pub(crate) fn check_norm<const BOUND: i32>(&self) -> bool {
         // Fine that this is not constant-time because it is used in a rejection loop -- the early quit leads to rejection.
-        for x in self.vec.iter() {
+        for x in self.elems.iter() {
             if x.check_norm::<BOUND>() {
                 return true;
             }
@@ -196,7 +184,7 @@ impl<const LEN: usize> Vector<LEN> {
         // 2: for 𝑖 from 0 to 𝑘 − 1 do
         // 3:   𝐰̃1 ← 𝐰̃1 || SimpleBitPack (𝐰1[𝑖], (𝑞 − 1)/(2𝛾2) − 1)
         // 4: end for
-        for w in self.vec.iter() {
+        for w in self.elems.iter() {
             h.absorb(&w.w1_encode::<POLY_W1_PACKED_LEN>())
                 .expect("absorb before squeeze is infallible");
         }

crypto/mldsa/src/mldsa_keys.rs

@@ -105,7 +105,7 @@ impl<const k: usize, const l: usize, const PK_LEN: usize> MLDSAPublicKey<k, l, P
         debug_assert_eq!(pk_chunks.len(), k);
         debug_assert_eq!(last_chunk.len(), 0);
 
-        for (pk_chunk, t1_i) in pk_chunks.into_iter().zip(&self.t1.vec) {
+        for (pk_chunk, t1_i) in pk_chunks.into_iter().zip(&self.t1.elems) {
             pk_chunk.copy_from_slice(&simple_bit_pack_t1(&t1_i));
         }
 
@@ -160,7 +160,7 @@ impl<const k: usize, const l: usize, const PK_LEN: usize> MLDSAPublicKeyTrait<k,
         debug_assert_eq!(pk_chunks.len(), k);
         debug_assert_eq!(last_chunk.len(), 0);
 
-        for (t1_i, pk_chunk) in t1.vec.iter_mut().zip(pk_chunks) {
+        for (t1_i, pk_chunk) in t1.elems.iter_mut().zip(pk_chunks) {
             // 3: 𝐭1[𝑖] ← SimpleBitUnpack(𝑧𝑖, 2bitlen (𝑞−1)−𝑑 − 1)
             //  ▷ This is always in the correct range
             //  Therefore, we don't need to check that the coeeffs are in range
@@ -407,7 +407,7 @@ impl<
 
 /// An ML-DSA private key.
 ///
-/// This will automatically inherit the [`Secret`] protections because [`Polynomial`] wraps the underlying data with [`Secret`].
+// Dev note: This will automatically inherit the [`Secret`] protections because [`Polynomial`] wraps the underlying data with [`Secret`].
 #[derive(Clone)]
 pub struct MLDSAPrivateKey<
     const k: usize,
@@ -456,7 +456,7 @@ impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, cons
 
         let sk_chunks = out[off..off + l * bitlen_eta(eta)].chunks_mut(bitlen_eta(eta));
         debug_assert_eq!(sk_chunks.len(), l);
-        for (sk_chunk, s1_hat_i) in sk_chunks.into_iter().zip(&self.s1_hat.vec) {
+        for (sk_chunk, s1_hat_i) in sk_chunks.into_iter().zip(&self.s1_hat.elems) {
             // Deviation from the FIPS:
             //   We are holding these in ntt form, so need to convert back to standard form
             let mut s1_hat_i = s1_hat_i.clone();
@@ -471,7 +471,7 @@ impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, cons
 
         let sk_chunks = out[off..off + k * bitlen_eta(eta)].chunks_mut(bitlen_eta(eta));
         debug_assert_eq!(sk_chunks.len(), k);
-        for (sk_chunk, s2_hat_i) in sk_chunks.into_iter().zip(&self.s2_hat.vec) {
+        for (sk_chunk, s2_hat_i) in sk_chunks.into_iter().zip(&self.s2_hat.elems) {
             // Deviation from the FIPS:
             //   We are holding these in ntt form, so need to convert back to standard form
             let mut s2_hat_i = s2_hat_i.clone();
@@ -486,7 +486,7 @@ impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, cons
 
         let sk_chunks = out[off..off + k * POLY_T0PACKED_LEN].chunks_mut(POLY_T0PACKED_LEN);
         debug_assert_eq!(sk_chunks.len(), k);
-        for (sk_chunk, t0_hat_i) in sk_chunks.into_iter().zip(&self.t0_hat.vec) {
+        for (sk_chunk, t0_hat_i) in sk_chunks.into_iter().zip(&self.t0_hat.elems) {
             // Deviation from the FIPS:
             //   We are holding these in ntt form, so need to convert back to standard form
             let mut t0_hat_i = t0_hat_i.clone();
@@ -630,7 +630,7 @@ impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, cons
         // unpack s1 directly into key.s1_hat so that we don't make additional non-Secret copies.
         let sk_chunks = sk[off..off + (l * bitlen_eta(eta))].chunks(bitlen_eta(eta));
         debug_assert_eq!(sk_chunks.len(), l);
-        for (s1_i, sk_chunk) in key.s1_hat.vec.iter_mut().zip(sk_chunks) {
+        for (s1_i, sk_chunk) in key.s1_hat.elems.iter_mut().zip(sk_chunks) {
             // 3: 𝐬1[𝑖] ← BitUnpack(𝑦𝑖, 𝜂, 𝜂)
             //  ▷ this may lie outside [−𝜂, 𝜂] if input is malformed
             s1_i.coeffs.copy_from_slice(&bit_unpack_eta::<eta>(&sk_chunk).coeffs);
@@ -650,7 +650,7 @@ impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, cons
         // unpack s2 directly into key.s2_hat so that we don't make additional non-Secret copies.
         let sk_chunks = sk[off..off + (k * bitlen_eta(eta))].chunks(bitlen_eta(eta));
         debug_assert_eq!(sk_chunks.len(), k);
-        for (s2_i, sk_chunk) in key.s2_hat.vec.iter_mut().zip(sk_chunks) {
+        for (s2_i, sk_chunk) in key.s2_hat.elems.iter_mut().zip(sk_chunks) {
             // 6: 𝐬2[𝑖] ← BitUnpack(𝑧𝑖, 𝜂, 𝜂)
             //  ▷ this may lie outside [−𝜂, 𝜂] if input is malformed
             s2_i.coeffs.copy_from_slice(&bit_unpack_eta::<eta>(&sk_chunk).coeffs);
@@ -675,7 +675,7 @@ impl<const k: usize, const l: usize, const eta: usize, const SK_LEN: usize, cons
         debug_assert_eq!(sk_chunks.len(), k);
         debug_assert_eq!(last_chunk.len(), 0);
 
-        for (t0_i, sk_chunk) in key.t0_hat.vec.iter_mut().zip(sk_chunks) {
+        for (t0_i, sk_chunk) in key.t0_hat.elems.iter_mut().zip(sk_chunks) {
             t0_i.coeffs.copy_from_slice(&bit_unpack_t0(sk_chunk).coeffs);
         }
         // Deviation from the FIPS:

crypto/mldsa/src/polynomial.rs

@@ -8,18 +8,12 @@ use core::ops::{Index, IndexMut};
 
 /// A polynomial over the ML-DSA ring.
 ///
-/// Dev note: The following structure does not necessarily need to be declared as public.
-/// There is no real scenario where this function needs to be called directly.
-/// However, in order to test the Debug and Display traits, it is necessary to use STD, so those
-/// can't be tested from inline tests in this file and the real unit tests are in a different crate.
-/// That's the reason why pub is used.
-///
 /// # 🚨 Security 🚨
 /// Polynomials themselves are not inherently secret since sometimes they are part of public keys
 /// and sometimes private keys.
 /// It is the responsibility of the caller to wrap sensitive instances in `Secret<Polynomial>`.
 #[derive(Clone, Copy)]
-pub struct Polynomial {
+pub(crate) struct Polynomial {
     pub(crate) coeffs: [i32; N],
 }
 

crypto/mlkem-lowmemory/src/lib.rs

@@ -264,6 +264,3 @@ pub use mlkem::{MLKEM_RND_LEN, MLKEM_SEED_LEN, MLKEM_SS_LEN};
 pub use mlkem::{MLKEM512_CT_LEN, MLKEM512_PK_LEN, MLKEM512_SK_LEN};
 pub use mlkem::{MLKEM768_CT_LEN, MLKEM768_PK_LEN, MLKEM768_SK_LEN};
 pub use mlkem::{MLKEM1024_CT_LEN, MLKEM1024_PK_LEN, MLKEM1024_SK_LEN};
-
-// re-export just so it is visible to unit tests
-pub use polynomial::Polynomial;

crypto/mlkem-lowmemory/src/polynomial.rs

@@ -7,17 +7,14 @@ use crate::mlkem::{N, q};
 use core::ops::{Index, IndexMut};
 
 /// A polynomial over the ML-KEM ring.
-/// Dev note: this doesn't strictly need to be pub ... ie there's no good reason for a caller to use this class directly,
-/// but in order to test the Debug and Display traits, you need STD, so those can't be tested from inline tests in this file
-/// and the real unit tests are in a different crate, so here we are.
 ///
 /// # 🚨 Security 🚨
 /// Polynomials themselves are not inherently secret since sometimes they are part of public keys
 /// and sometimes private keys.
 /// It is the responsibility of the caller to wrap sensitive instances in `Secret<Polynomial>`.
 /// Note: at the moment, nothing in this crate uses `Secret<Polynomial>`, so I have left the `impl ZeroizablePrimitive` commented-out.
 #[derive(Clone, Copy)]
-pub struct Polynomial {
+pub(crate) struct Polynomial {
     pub(crate) coeffs: [i16; N],
 }
 

crypto/mlkem/src/aux_functions.rs

@@ -1,8 +1,8 @@
 //! Implements auxiliary functions for ML-DSA as defined in Section 7 of FIPS 204.
 
-use crate::matrix::Vector;
+use crate::matrix::{Matrix, Vector};
 use crate::mlkem::{N, q, q_inv};
-use crate::{Matrix, Polynomial};
+use crate::polynomial::Polynomial;
 use bouncycastle_core::traits::XOF;
 use bouncycastle_sha3::{SHAKE128, SHAKE256};
 
@@ -13,7 +13,7 @@ pub(crate) fn expandA<const k: usize>(rho: &[u8; 32]) -> Matrix<k, k> {
         for j in 0..k {
             // 6: 𝐀[𝑖, 𝑗] ← SampleNTT(𝜌‖𝑗‖𝑖)
             //  ▷ 𝑗 and 𝑖 are bytes 33 and 34 of the input
-            A_hat[i][j] = sample_ntt(rho, &[j as u8, i as u8]);
+            A_hat.elems[i][j] = sample_ntt(rho, &[j as u8, i as u8]);
         }
     }