grpc 1.84.0 → 1.85.0.pre1

data/third_party/boringssl-with-bazel/crypto/fipsmodule/mlkem/mlkem.cc.inc

@@ -26,10 +26,19 @@
 
 #include "../../internal.h"
 #include "../bcm_interface.h"
-#include "../delocate.h"
 #include "../keccak/internal.h"
 
 
+using namespace bssl;
+
+#if defined(BORINGSSL_FIPS)
+
+DEFINE_STATIC_ONCE(g_mlkem_keygen_self_test_once)
+DEFINE_STATIC_ONCE(g_mlkem_encap_self_test_once)
+DEFINE_STATIC_ONCE(g_mlkem_decap_self_test_once)
+
+#endif
+
 namespace mlkem {
 namespace {
 
@@ -42,10 +51,17 @@ void ensure_decap_self_test();
 // See
 // https://csrc.nist.gov/pubs/fips/203/final
 
-static void prf(uint8_t *out, size_t out_len, const uint8_t in[33]) {
+inline void prf(uint8_t *out, size_t out_len, const uint8_t in[33]) {
   BORINGSSL_keccak(out, out_len, in, 33, boringssl_shake256);
 }
 
+#if defined(HAVE_KECCAK_X2)
+inline void prf_x2(uint8_t *outs[2], size_t out_len,
+                   const uint8_t *ins[2] /* 33 each */) {
+  BORINGSSL_keccak_short_x2(outs, out_len, ins, 33, boringssl_shake256);
+}
+#endif
+
 // Section 4.1
 void hash_h(uint8_t out[32], const uint8_t *in, size_t len) {
   BORINGSSL_keccak(out, 32, in, len, boringssl_sha3_256);
@@ -56,21 +72,21 @@ void hash_g(uint8_t out[64], const uint8_t *in, size_t len) {
 }
 
 // This is called `J` in the spec.
-void kdf(uint8_t out[BCM_MLKEM_SHARED_SECRET_BYTES],
+void kdf(uint8_t out[MLKEM_SHARED_SECRET_BYTES],
          const uint8_t failure_secret[32], const uint8_t *ciphertext,
          size_t ciphertext_len) {
-  struct BORINGSSL_keccak_st st;
+  BORINGSSL_keccak_st st;
   BORINGSSL_keccak_init(&st, boringssl_shake256);
   BORINGSSL_keccak_absorb(&st, failure_secret, 32);
   BORINGSSL_keccak_absorb(&st, ciphertext, ciphertext_len);
-  BORINGSSL_keccak_squeeze(&st, out, BCM_MLKEM_SHARED_SECRET_BYTES);
+  BORINGSSL_keccak_squeeze(&st, out, MLKEM_SHARED_SECRET_BYTES);
 }
 
 // Constants that are common across all sizes.
 #define DEGREE 256
 const size_t kBarrettMultiplier = 5039;
 const unsigned kBarrettShift = 24;
-static const uint16_t kPrime = 3329;
+const uint16_t kPrime = 3329;
 const int kLog2Prime = 12;
 const uint16_t kHalfPrime = (/*kPrime=*/3329 - 1) / 2;
 // kInverseDegree is 128^-1 mod 3329; 128 because kPrime does not have a 512th
@@ -79,10 +95,10 @@ const uint16_t kInverseDegree = 3303;
 
 // Rank-specific constants.
 #define RANK768 3
-static const int kDU768 = 10;
+const int kDU768 = 10;
 const int kDV768 = 4;
 #define RANK1024 4
-static const int kDU1024 = 11;
+const int kDU1024 = 11;
 const int kDV1024 = 5;
 
 constexpr size_t encoded_vector_size(int rank) {
@@ -93,11 +109,8 @@ constexpr size_t encoded_public_key_size(int rank) {
   return encoded_vector_size(rank) + /*sizeof(rho)=*/32;
 }
 
-static_assert(encoded_public_key_size(RANK768) == BCM_MLKEM768_PUBLIC_KEY_BYTES,
-              "");
-static_assert(encoded_public_key_size(RANK1024) ==
-                  BCM_MLKEM1024_PUBLIC_KEY_BYTES,
-              "");
+static_assert(encoded_public_key_size(RANK768) == MLKEM768_PUBLIC_KEY_BYTES);
+static_assert(encoded_public_key_size(RANK1024) == MLKEM1024_PUBLIC_KEY_BYTES);
 
 constexpr size_t compressed_vector_size(int rank) {
   // `if constexpr` isn't available in C++17.
@@ -110,13 +123,13 @@ constexpr size_t ciphertext_size(int rank) {
          (rank == RANK768 ? kDV768 : kDV1024) * DEGREE / 8;
 }
 
-static_assert(ciphertext_size(RANK768) == BCM_MLKEM768_CIPHERTEXT_BYTES, "");
-static_assert(ciphertext_size(RANK1024) == BCM_MLKEM1024_CIPHERTEXT_BYTES, "");
+static_assert(ciphertext_size(RANK768) == MLKEM768_CIPHERTEXT_BYTES);
+static_assert(ciphertext_size(RANK1024) == MLKEM1024_CIPHERTEXT_BYTES);
 
-typedef struct scalar {
+struct scalar {
   // On every function entry and exit, 0 <= c < kPrime.
   uint16_t c[DEGREE];
-} scalar;
+};
 
 template <int RANK>
 struct vector {
@@ -187,7 +200,7 @@ const uint16_t kModRoots[128] = {
 };
 
 // reduce_once reduces 0 <= x < 2*kPrime, mod kPrime.
-uint16_t reduce_once(uint16_t x) {
+inline uint16_t reduce_once(uint16_t x) {
   declassify_assert(x < 2 * kPrime);
   const uint16_t subtracted = x - kPrime;
   uint16_t mask = 0u - (subtracted >> 15);
@@ -199,13 +212,13 @@ uint16_t reduce_once(uint16_t x) {
   // We usually add value barriers to selects because Clang turns consecutive
   // selects with the same condition into a branch instead of CMOV/CSEL. This
   // condition does not occur in ML-KEM, so omitting it seems to be safe so far,
-  // but see |scalar_centered_binomial_distribution_eta_2_with_prf|.
+  // but see `scalar_centered_binomial_distribution_eta_2_with_prf`.
   return (mask & x) | (~mask & subtracted);
 }
 
 // constant time reduce x mod kPrime using Barrett reduction. x must be less
 // than kPrime + 2×kPrime².
-static uint16_t reduce(uint32_t x) {
+inline uint16_t reduce(uint32_t x) {
   declassify_assert(x < kPrime + 2u * kPrime * kPrime);
   uint64_t product = (uint64_t)x * kBarrettMultiplier;
   uint32_t quotient = (uint32_t)(product >> kBarrettShift);
@@ -213,41 +226,49 @@ static uint16_t reduce(uint32_t x) {
   return reduce_once(remainder);
 }
 
-void scalar_zero(scalar *out) { OPENSSL_memset(out, 0, sizeof(*out)); }
+inline void scalar_zero(scalar *out) { OPENSSL_memset(out, 0, sizeof(*out)); }
 
 template <int RANK>
-void vector_zero(vector<RANK> *out) {
+inline void vector_zero(vector<RANK> *out) {
   OPENSSL_memset(out->v, 0, sizeof(scalar) * RANK);
 }
 
 // In place number theoretic transform of a given scalar.
 // Note that MLKEM's kPrime 3329 does not have a 512th root of unity, so this
 // transform leaves off the last iteration of the usual FFT code, with the 128
-// relevant roots of unity being stored in |kNTTRoots|. This means the output
+// relevant roots of unity being stored in `kNTTRoots`. This means the output
 // should be seen as 128 elements in GF(3329^2), with the coefficients of the
-// elements being consecutive entries in |s->c|.
-static void scalar_ntt(scalar *s) {
-  int offset = DEGREE;
-  // `int` is used here because using `size_t` throughout caused a ~5% slowdown
-  // with Clang 14 on Aarch64.
-  for (int step = 1; step < DEGREE / 2; step <<= 1) {
-    offset >>= 1;
-    int k = 0;
-    for (int i = 0; i < step; i++) {
-      const uint32_t step_root = kNTTRoots[i + step];
-      for (int j = k; j < k + offset; j++) {
-        uint16_t odd = reduce(step_root * s->c[j + offset]);
-        uint16_t even = s->c[j];
-        s->c[j] = reduce_once(odd + even);
-        s->c[j + offset] = reduce_once(even - odd + kPrime);
-      }
-      k += 2 * offset;
-    }
+// elements being consecutive entries in `s->c`.
+inline void scalar_ntt(scalar *s) {
+  // Manually unrolled loop to maximize vectorization.
+#define ITER(step, offset)                                   \
+  {                                                          \
+    int k = 0;                                               \
+    for (int i = 0; i < step; i++) {                         \
+      const uint32_t step_root = kNTTRoots[i + step];        \
+      for (int j = k; j < k + offset; j++) {                 \
+        uint16_t odd = reduce(step_root * s->c[j + offset]); \
+        uint16_t even = s->c[j];                             \
+        s->c[j] = reduce_once(odd + even);                   \
+        s->c[j + offset] = reduce_once(even - odd + kPrime); \
+      }                                                      \
+      k += 2 * offset;                                       \
+    }                                                        \
   }
+  // for (int step = 1; step < DEGREE / 2; step <<= 1)
+  ITER(1, 128)
+  ITER(2, 64)
+  ITER(4, 32)
+  ITER(8, 16)
+  ITER(16, 8)
+  ITER(32, 4)
+  ITER(64, 2)
+  static_assert(DEGREE == 256);
+#undef ITER
 }
 
 template <int RANK>
-static void vector_ntt(vector<RANK> *a) {
+inline void vector_ntt(vector<RANK> *a) {
   for (int i = 0; i < RANK; i++) {
     scalar_ntt(&a->v[i]);
   }
@@ -257,25 +278,34 @@ static void vector_ntt(vector<RANK> *a) {
 // entries of s->v being interpreted as elements of GF(3329^2). Just as with the
 // number theoretic transform, this leaves off the first step of the normal iFFT
 // to account for the fact that 3329 does not have a 512th root of unity, using
-// the precomputed 128 roots of unity stored in |kInverseNTTRoots|.
+// the precomputed 128 roots of unity stored in `kInverseNTTRoots`.
 void scalar_inverse_ntt(scalar *s) {
-  int step = DEGREE / 2;
-  // `int` is used here because using `size_t` throughout caused a ~5% slowdown
-  // with Clang 14 on Aarch64.
-  for (int offset = 2; offset < DEGREE; offset <<= 1) {
-    step >>= 1;
-    int k = 0;
-    for (int i = 0; i < step; i++) {
-      uint32_t step_root = kInverseNTTRoots[i + step];
-      for (int j = k; j < k + offset; j++) {
-        uint16_t odd = s->c[j + offset];
-        uint16_t even = s->c[j];
-        s->c[j] = reduce_once(odd + even);
-        s->c[j + offset] = reduce(step_root * (even - odd + kPrime));
-      }
-      k += 2 * offset;
-    }
+  // Manually unrolled loop to maximize vectorization.
+#define ITER(step, offset)                                            \
+  {                                                                   \
+    int k = 0;                                                        \
+    for (int i = 0; i < step; i++) {                                  \
+      uint32_t step_root = kInverseNTTRoots[i + step];                \
+      for (int j = k; j < k + offset; j++) {                          \
+        uint16_t odd = s->c[j + offset];                              \
+        uint16_t even = s->c[j];                                      \
+        s->c[j] = reduce_once(odd + even);                            \
+        s->c[j + offset] = reduce(step_root * (even - odd + kPrime)); \
+      }                                                               \
+      k += 2 * offset;                                                \
+    }                                                                 \
   }
+  // for (int offset = 2; offset < DEGREE; offset <<= 1)
+  ITER(64, 2)
+  ITER(32, 4)
+  ITER(16, 8)
+  ITER(8, 16)
+  ITER(4, 32)
+  ITER(2, 64)
+  ITER(1, 128)
+  static_assert(DEGREE == 256);
+#undef ITER
+
   for (int i = 0; i < DEGREE; i++) {
     s->c[i] = reduce(s->c[i] * kInverseDegree);
   }
@@ -288,13 +318,13 @@ void vector_inverse_ntt(vector<RANK> *a) {
   }
 }
 
-void scalar_add(scalar *lhs, const scalar *rhs) {
+inline void scalar_add(scalar *lhs, const scalar *rhs) {
   for (int i = 0; i < DEGREE; i++) {
     lhs->c[i] = reduce_once(lhs->c[i] + rhs->c[i]);
   }
 }
 
-void scalar_sub(scalar *lhs, const scalar *rhs) {
+inline void scalar_sub(scalar *lhs, const scalar *rhs) {
   for (int i = 0; i < DEGREE; i++) {
     lhs->c[i] = reduce_once(lhs->c[i] - rhs->c[i] + kPrime);
   }
@@ -304,11 +334,11 @@ void scalar_sub(scalar *lhs, const scalar *rhs) {
 // 3329 does not have a 512th root of unity, this means we have to interpret
 // the 2*ith and (2*i+1)th entries of the scalar as elements of GF(3329)[X]/(X^2
 // - 17^(2*bitreverse(i)+1)) The value of 17^(2*bitreverse(i)+1) mod 3329 is
-// stored in the precomputed |kModRoots| table. Note that our Barrett transform
-// only allows us to multipy two reduced numbers together, so we need some
+// stored in the precomputed `kModRoots` table. Note that our Barrett transform
+// only allows us to multiply two reduced numbers together, so we need some
 // intermediate reduction steps, even if an uint64_t could hold 3 multiplied
 // numbers.
-void scalar_mult(scalar *out, const scalar *lhs, const scalar *rhs) {
+inline void scalar_mult(scalar *out, const scalar *lhs, const scalar *rhs) {
   for (int i = 0; i < DEGREE / 2; i++) {
     uint32_t real_real = (uint32_t)lhs->c[2 * i] * rhs->c[2 * i];
     uint32_t img_img = (uint32_t)lhs->c[2 * i + 1] * rhs->c[2 * i + 1];
@@ -321,14 +351,14 @@ void scalar_mult(scalar *out, const scalar *lhs, const scalar *rhs) {
 }
 
 template <int RANK>
-void vector_add(vector<RANK> *lhs, const vector<RANK> *rhs) {
+inline void vector_add(vector<RANK> *lhs, const vector<RANK> *rhs) {
   for (int i = 0; i < RANK; i++) {
     scalar_add(&lhs->v[i], &rhs->v[i]);
   }
 }
 
 template <int RANK>
-static void matrix_mult(vector<RANK> *out, const matrix<RANK> *m,
+inline void matrix_mult(vector<RANK> *out, const matrix<RANK> *m,
                         const vector<RANK> *a) {
   vector_zero(out);
   for (int i = 0; i < RANK; i++) {
@@ -341,7 +371,7 @@ static void matrix_mult(vector<RANK> *out, const matrix<RANK> *m,
 }
 
 template <int RANK>
-void matrix_mult_transpose(vector<RANK> *out, const matrix<RANK> *m,
+inline void matrix_mult_transpose(vector<RANK> *out, const matrix<RANK> *m,
                            const vector<RANK> *a) {
   vector_zero(out);
   for (int i = 0; i < RANK; i++) {
@@ -354,7 +384,7 @@ void matrix_mult_transpose(vector<RANK> *out, const matrix<RANK> *m,
 }
 
 template <int RANK>
-void scalar_inner_product(scalar *out, const vector<RANK> *lhs,
+inline void scalar_inner_product(scalar *out, const vector<RANK> *lhs,
                           const vector<RANK> *rhs) {
   scalar_zero(out);
   for (int i = 0; i < RANK; i++) {
@@ -364,11 +394,25 @@ void scalar_inner_product(scalar *out, const vector<RANK> *lhs,
   }
 }
 
-// Algorithm 6 from the spec. Rejection samples a Keccak stream to get
+inline void scalar_from_keccak_block_vartime(scalar *out, int *done,
+                                             const uint8_t block[168]) {
+  for (size_t i = 0; i < 168 && *done < DEGREE; i += 3) {
+    uint16_t d1 = block[i] + 256 * (block[i + 1] % 16);
+    uint16_t d2 = block[i + 1] / 16 + 16 * block[i + 2];
+    if (d1 < kPrime) {
+      out->c[(*done)++] = d1;
+    }
+    if (d2 < kPrime && *done < DEGREE) {
+      out->c[(*done)++] = d2;
+    }
+  }
+}
+
+// Algorithm 7 from the spec. Rejection samples a Keccak stream to get
 // uniformly distributed elements. This is used for matrix expansion and only
 // operates on public inputs.
-static void scalar_from_keccak_vartime(scalar *out,
-                                       struct BORINGSSL_keccak_st *keccak_ctx) {
+inline void scalar_from_keccak_vartime(scalar *out,
+                                       BORINGSSL_keccak_st *keccak_ctx) {
   assert(keccak_ctx->squeeze_offset == 0);
   assert(keccak_ctx->rate_bytes == 168);
   static_assert(168 % 3 == 0, "block and coefficient boundaries do not align");
@@ -377,40 +421,41 @@ static void scalar_from_keccak_vartime(scalar *out,
   while (done < DEGREE) {
     uint8_t block[168];
     BORINGSSL_keccak_squeeze(keccak_ctx, block, sizeof(block));
-    for (size_t i = 0; i < sizeof(block) && done < DEGREE; i += 3) {
-      uint16_t d1 = block[i] + 256 * (block[i + 1] % 16);
-      uint16_t d2 = block[i + 1] / 16 + 16 * block[i + 2];
-      if (d1 < kPrime) {
-        out->c[done++] = d1;
-      }
-      if (d2 < kPrime && done < DEGREE) {
-        out->c[done++] = d2;
-      }
-    }
+    scalar_from_keccak_block_vartime(out, &done, block);
   }
 }
 
-// Algorithm 7 from the spec, with eta fixed to two and the PRF call
-// included. Creates binominally distributed elements by sampling 2*|eta| bits,
-// and setting the coefficient to the count of the first bits minus the count of
-// the second bits, resulting in a centered binomial distribution. Since eta is
-// two this gives -2/2 with a probability of 1/16, -1/1 with probability 1/4,
-// and 0 with probability 3/8.
-void scalar_centered_binomial_distribution_eta_2_with_prf(
-    scalar *out, const uint8_t input[33]) {
-  uint8_t entropy[128];
-  static_assert(sizeof(entropy) == 2 * /*kEta=*/2 * DEGREE / 8, "");
-  prf(entropy, sizeof(entropy), input);
+#if defined(HAVE_KECCAK_X2)
+inline void scalar_from_keccak_vartime_x2(scalar *out[2],
+                                          BORINGSSL_keccak_st keccak_ctx[2]) {
+  assert(keccak_ctx[0].squeeze_offset == 0);
+  assert(keccak_ctx[0].rate_bytes == 168);
+  assert(keccak_ctx[1].squeeze_offset == 0);
+  assert(keccak_ctx[1].rate_bytes == 168);
+  static_assert(168 % 3 == 0, "block and coefficient boundaries do not align");
+
+  int done[2] = {0, 0};
+  while (done[0] < DEGREE || done[1] < DEGREE) {
+    uint8_t block[2][168];
+    uint8_t *blocks[] = {block[0], block[1]};
+    BORINGSSL_keccak_squeeze_x2(keccak_ctx, blocks, sizeof(block[0]));
+    scalar_from_keccak_block_vartime(out[0], &done[0], block[0]);
+    scalar_from_keccak_block_vartime(out[1], &done[1], block[1]);
+  }
+}
+#endif
 
+static void scalar_centered_binomial_distribution_eta_2(
+    scalar *out, const uint8_t entropy[128]) {
   for (int i = 0; i < DEGREE; i += 2) {
     uint8_t byte = entropy[i / 2];
 
     uint16_t value = (byte & 1) + ((byte >> 1) & 1);
     value -= ((byte >> 2) & 1) + ((byte >> 3) & 1);
-    // Add |kPrime| if |value| underflowed. See |reduce_once| for a discussion
+    // Add `kPrime` if `value` underflowed. See `reduce_once` for a discussion
     // on why the value barrier is omitted. While this could have been written
     // reduce_once(value + kPrime), this is one extra addition and small range
-    // of |value| tempts some versions of Clang to emit a branch.
+    // of `value` tempts some versions of Clang to emit a branch.
     uint16_t mask = 0u - (value >> 15);
     out->c[i] = ((value + kPrime) & mask) | (value & ~mask);
 
@@ -423,170 +468,388 @@ void scalar_centered_binomial_distribution_eta_2_with_prf(
   }
 }
 
+// Algorithm 8 from the spec, with eta fixed to two and the PRF call
+// included. Creates binominally distributed elements by sampling 2*`eta` bits,
+// and setting the coefficient to the count of the first bits minus the count of
+// the second bits, resulting in a centered binomial distribution. Since eta is
+// two this gives -2/2 with a probability of 1/16, -1/1 with probability 1/4,
+// and 0 with probability 3/8.
+void scalar_centered_binomial_distribution_eta_2_with_prf(
+    scalar *out, const uint8_t input[33]) {
+  uint8_t entropy[128];
+  static_assert(sizeof(entropy) == 2 * /*kEta=*/2 * DEGREE / 8);
+  prf(entropy, sizeof(entropy), input);
+  scalar_centered_binomial_distribution_eta_2(out, entropy);
+}
+
+#if defined(HAVE_KECCAK_X2)
+// scalar_centered_binomial_distribution_eta_2_with_prf_x2 computes two blocks
+// of scalar_centered_binomial_distribution_eta_2_with_prf.
+static void scalar_centered_binomial_distribution_eta_2_with_prf_x2(
+    scalar out[2], uint8_t input[33]) {
+  uint8_t entropy[2][128];
+  static_assert(sizeof(entropy[0]) == 2 * /*kEta=*/2 * DEGREE / 8);
+
+  uint8_t input1[33];
+  OPENSSL_memcpy(input1, input, 33);
+  ++input1[32];  // Counter.
+
+  uint8_t *entropies[] = {entropy[0], entropy[1]};
+  const uint8_t *inputs[2] = {input, input1};
+  prf_x2(entropies, sizeof(entropy[0]), inputs);
+
+  for (int k = 0; k < 2; k++) {
+    scalar_centered_binomial_distribution_eta_2(&out[k], entropy[k]);
+  }
+}
+#endif
+
 // Generates a secret vector by using
-// |scalar_centered_binomial_distribution_eta_2_with_prf|, using the given seed
-// appending and incrementing |counter| for entry of the vector.
+// `scalar_centered_binomial_distribution_eta_2_with_prf`, using the given seed
+// appending and incrementing `counter` for entry of the vector.
 template <int RANK>
 void vector_generate_secret_eta_2(vector<RANK> *out, uint8_t *counter,
                                   const uint8_t seed[32]) {
   uint8_t input[33];
   OPENSSL_memcpy(input, seed, 32);
-  for (int i = 0; i < RANK; i++) {
+
+  int i = 0;
+#if defined(HAVE_KECCAK_X2)
+  for (; i + 2 <= RANK; i += 2) {
+    input[32] = *counter;
+    *counter += 2;
+    scalar_centered_binomial_distribution_eta_2_with_prf_x2(&out->v[i], input);
+  }
+#endif
+
+  for (; i < RANK; i++) {
     input[32] = (*counter)++;
     scalar_centered_binomial_distribution_eta_2_with_prf(&out->v[i], input);
   }
 }
 
+template <int RANK>
+void matrix_expand_step(scalar *out, uint8_t input[34], int i) {
+  input[32] = i / RANK;
+  input[33] = i % RANK;
+  BORINGSSL_keccak_st keccak_ctx;
+  BORINGSSL_keccak_init(&keccak_ctx, boringssl_shake128);
+  BORINGSSL_keccak_absorb(&keccak_ctx, input, 34);
+  scalar_from_keccak_vartime(out, &keccak_ctx);
+}
+
+#if defined(HAVE_KECCAK_X2)
+template <int RANK>
+void matrix_expand_step_x2(matrix<RANK> out[2], uint8_t input[34], int i) {
+  input[32] = i / RANK;
+  input[33] = i % RANK;
+  uint8_t input1[34];
+  OPENSSL_memcpy(input1, input, 32);
+  input1[32] = (i + 1) / RANK;
+  input1[33] = (i + 1) % RANK;
+  BORINGSSL_keccak_st keccak_ctx[2];
+  BORINGSSL_keccak_init(&keccak_ctx[0], boringssl_shake128);
+  BORINGSSL_keccak_init(&keccak_ctx[1], boringssl_shake128);
+  BORINGSSL_keccak_absorb(&keccak_ctx[0], input, 34);
+  BORINGSSL_keccak_absorb(&keccak_ctx[1], input1, 34);
+  scalar *outs[] = {&out->v[i / RANK][i % RANK],
+                    &out->v[(i + 1) / RANK][(i + 1) % RANK]};
+  scalar_from_keccak_vartime_x2(outs, keccak_ctx);
+}
+#endif
+
 // Expands the matrix of a seed for key generation and for encaps-CPA.
 template <int RANK>
 void matrix_expand(matrix<RANK> *out, const uint8_t rho[32]) {
   uint8_t input[34];
   OPENSSL_memcpy(input, rho, 32);
-  for (int i = 0; i < RANK; i++) {
-    for (int j = 0; j < RANK; j++) {
-      input[32] = i;
-      input[33] = j;
-      struct BORINGSSL_keccak_st keccak_ctx;
-      BORINGSSL_keccak_init(&keccak_ctx, boringssl_shake128);
-      BORINGSSL_keccak_absorb(&keccak_ctx, input, sizeof(input));
-      scalar_from_keccak_vartime(&out->v[i][j], &keccak_ctx);
-    }
+
+  int i = 0;
+#if defined(HAVE_KECCAK_X2)
+  for (; i + 2 <= RANK * RANK; i += 2) {
+    matrix_expand_step_x2(out, input, i);
   }
-}
+#endif
 
-const uint8_t kMasks[8] = {0x01, 0x03, 0x07, 0x0f, 0x1f, 0x3f, 0x7f, 0xff};
+  for (; i < RANK * RANK; i++) {
+    matrix_expand_step<RANK>(&out->v[i / RANK][i % RANK], input, i);
+  }
+}
 
-void scalar_encode(uint8_t *out, const scalar *s, int bits) {
-  assert(bits <= (int)sizeof(*s->c) * 8 && bits != 1);
+// Encodes a scalar of 256 `BITS`-bit words into 32*`BITS` bytes by splitting
+// and joining into bytes using LSB-first bit order (i.e. opposite to standard
+// reading order). See below for examples. If an input is >= 1 << `BITS`, the
+// result is undefined.
+template <int BITS>
+void scalar_encode(uint8_t *out, const scalar *s);
+
+// Encodes a scalar of 256 10-bit words into 320 bytes as follows:
+// 000000Aaaaaaaaaa 000000Bbbbbbbbbb 000000Cccccccccc 000000Dddddddddd ...
+// -> aaaaaaaa bbbbbbAa ccccBbbb ddCccccc Dddddddd ...
+template <>
+void scalar_encode<10>(uint8_t out[320], const scalar *s) {
+  for (int i = 0; i < DEGREE; i += 4) {
+    uint16_t s0 = s->c[i];
+    uint16_t s1 = s->c[i + 1];
+    uint16_t s2 = s->c[i + 2];
+    uint16_t s3 = s->c[i + 3];
+    declassify_assert((s0 | s1 | s2 | s3) < (1 << 10));
+    out[0] = (uint8_t)s0;
+    out[1] = (uint8_t)((s0 >> 8) | (s1 << 2));
+    out[2] = (uint8_t)((s1 >> 6) | (s2 << 4));
+    out[3] = (uint8_t)((s2 >> 4) | (s3 << 6));
+    out[4] = (uint8_t)(s3 >> 2);
+    out += 5;
+  }
+}
 
-  uint8_t out_byte = 0;
-  int out_byte_bits = 0;
+// Encodes a scalar of 256 12-bit words into 384 bytes as follows:
+// 0000Aaaaaaaaaaaa 0000Bbbbbbbbbbbb 0000Cccccccccccc 0000Dddddddddddd ...
+// -> aaaaaaaa bbbbAaaa Bbbbbbbb cccccccc ddddCccc Dddddddd ....
+template <>
+void scalar_encode<12>(uint8_t out[384], const scalar *s) {
+  for (int i = 0; i < DEGREE; i += 2) {
+    uint16_t s0 = s->c[i];
+    uint16_t s1 = s->c[i + 1];
+    declassify_assert((s0 | s1) < (1 << 12));
+    out[0] = (uint8_t)s0;
+    out[1] = (uint8_t)((s0 >> 8) | (s1 << 4));
+    out[2] = (uint8_t)(s1 >> 4);
+    out += 3;
+  }
+}
 
-  for (int i = 0; i < DEGREE; i++) {
-    uint16_t element = s->c[i];
-    int element_bits_done = 0;
-
-    while (element_bits_done < bits) {
-      int chunk_bits = bits - element_bits_done;
-      int out_bits_remaining = 8 - out_byte_bits;
-      if (chunk_bits >= out_bits_remaining) {
-        chunk_bits = out_bits_remaining;
-        out_byte |= (element & kMasks[chunk_bits - 1]) << out_byte_bits;
-        *out = out_byte;
-        out++;
-        out_byte_bits = 0;
-        out_byte = 0;
-      } else {
-        out_byte |= (element & kMasks[chunk_bits - 1]) << out_byte_bits;
-        out_byte_bits += chunk_bits;
-      }
+// Encodes a scalar of 256 4-bit words into 128 bytes as follows:
+// 000000000000Aaaa 00000000000Bbbb 000000000000Cccc 000000000000Dddd ...
+// -> BbbbAaaa DdddCccc ...
+template <>
+void scalar_encode<4>(uint8_t out[128], const scalar *s) {
+  for (int i = 0; i < DEGREE; i += 2) {
+    uint16_t s0 = s->c[i];
+    uint16_t s1 = s->c[i + 1];
+    declassify_assert((s0 | s1) < (1 << 4));
+    out[0] = (uint8_t)(s0 | (s1 << 4));
+    out += 1;
+  }
+}
 
-      element_bits_done += chunk_bits;
-      element >>= chunk_bits;
-    }
+// Encodes a scalar of 256 11-bit words into 352 bytes as follows:
+// 00000Aaaaaaaaaaa 00000Bbbbbbbbbbb 00000Ccccccccccc 00000Ddddddddddd ...
+// -> aaaaaaaa bbbbbAaa ccBbbbbb cccccccc dddddddC eeeeDddd fEeeeeee ...
+template <>
+void scalar_encode<11>(uint8_t out[352], const scalar *s) {
+  for (int i = 0; i < DEGREE; i += 8) {
+    uint16_t s0 = s->c[i];
+    uint16_t s1 = s->c[i + 1];
+    uint16_t s2 = s->c[i + 2];
+    uint16_t s3 = s->c[i + 3];
+    uint16_t s4 = s->c[i + 4];
+    uint16_t s5 = s->c[i + 5];
+    uint16_t s6 = s->c[i + 6];
+    uint16_t s7 = s->c[i + 7];
+    declassify_assert((s0 | s1 | s2 | s3 | s4 | s5 | s6 | s7) < (1 << 11));
+    out[0] = (uint8_t)s0;
+    out[1] = (uint8_t)((s0 >> 8) | (s1 << 3));
+    out[2] = (uint8_t)((s1 >> 5) | (s2 << 6));
+    out[3] = (uint8_t)(s2 >> 2);
+    out[4] = (uint8_t)((s2 >> 10) | (s3 << 1));
+    out[5] = (uint8_t)((s3 >> 7) | (s4 << 4));
+    out[6] = (uint8_t)((s4 >> 4) | (s5 << 7));
+    out[7] = (uint8_t)(s5 >> 1);
+    out[8] = (uint8_t)((s5 >> 9) | (s6 << 2));
+    out[9] = (uint8_t)((s6 >> 6) | (s7 << 5));
+    out[10] = (uint8_t)(s7 >> 3);
+    out += 11;
   }
+}
 
-  if (out_byte_bits > 0) {
-    *out = out_byte;
+// Encodes a scalar of 256 5-bit words into 160 bytes as follows:
+// 00000000000Aaaaa 00000000000Bbbbb 00000000000Ccccc 00000000000Ddddd ...
+// -> bbbAaaaa dCccccBb eeeeDddd ggFffffE HhhhhGgg ...
+template <>
+void scalar_encode<5>(uint8_t out[160], const scalar *s) {
+  for (int i = 0; i < DEGREE; i += 8) {
+    uint16_t s0 = s->c[i];
+    uint16_t s1 = s->c[i + 1];
+    uint16_t s2 = s->c[i + 2];
+    uint16_t s3 = s->c[i + 3];
+    uint16_t s4 = s->c[i + 4];
+    uint16_t s5 = s->c[i + 5];
+    uint16_t s6 = s->c[i + 6];
+    uint16_t s7 = s->c[i + 7];
+    declassify_assert((s0 | s1 | s2 | s3 | s4 | s5 | s6 | s7) < (1 << 5));
+    out[0] = (uint8_t)(s0 | (s1 << 5));
+    out[1] = (uint8_t)((s1 >> 3) | (s2 << 2) | (s3 << 7));
+    out[2] = (uint8_t)((s3 >> 1) | (s4 << 4));
+    out[3] = (uint8_t)((s4 >> 4) | (s5 << 1) | (s6 << 6));
+    out[4] = (uint8_t)((s6 >> 2) | (s7 << 3));
+    out += 5;
   }
 }
 
-// scalar_encode_1 is |scalar_encode| specialised for |bits| == 1.
-void scalar_encode_1(uint8_t out[32], const scalar *s) {
+// Encodes a scalar of 256 1-bit "words" into 32 bytes as follows:
+// 0000000000000A 000000000000000B 000000000000000C 000000000000000D ...
+// -> HGFEDCBA PONMLKJI XWVUTSRQ ...
+// This order is best understood as the natural way of joining into bytes
+// assuming LSB-first bit order.
+template <>
+void scalar_encode<1>(uint8_t out[32], const scalar *s) {
   for (int i = 0; i < DEGREE; i += 8) {
     uint8_t out_byte = 0;
     for (int j = 0; j < 8; j++) {
-      out_byte |= (s->c[i + j] & 1) << j;
+      declassify_assert(s->c[i + j] <= 1);
+      out_byte |= s->c[i + j] << j;
     }
-    *out = out_byte;
-    out++;
+    out[i / 8] = out_byte;
   }
 }
 
-// Encodes an entire vector into 32*|RANK|*|bits| bytes. Note that since 256
+// Encodes an entire vector into 32*`RANK`*`bits` bytes. Note that since 256
 // (DEGREE) is divisible by 8, the individual vector entries will always fill a
 // whole number of bytes, so we do not need to worry about bit packing here.
-template <int RANK>
-void vector_encode(uint8_t *out, const vector<RANK> *a, int bits) {
+template <int bits, int RANK>
+void vector_encode(uint8_t *out, const vector<RANK> *a) {
   for (int i = 0; i < RANK; i++) {
-    scalar_encode(out + i * bits * DEGREE / 8, &a->v[i], bits);
+    scalar_encode<bits>(out + i * bits * DEGREE / 8, &a->v[i]);
   }
 }
 
-// scalar_decode parses |DEGREE * bits| bits from |in| into |DEGREE| values in
-// |out|. It returns one on success and zero if any parsed value is >=
-// |kPrime|.
-int scalar_decode(scalar *out, const uint8_t *in, int bits) {
-  assert(bits <= (int)sizeof(*out->c) * 8 && bits != 1);
-
-  uint8_t in_byte = 0;
-  int in_byte_bits_left = 0;
-
-  for (int i = 0; i < DEGREE; i++) {
-    uint16_t element = 0;
-    int element_bits_done = 0;
-
-    while (element_bits_done < bits) {
-      if (in_byte_bits_left == 0) {
-        in_byte = *in;
-        in++;
-        in_byte_bits_left = 8;
-      }
+// The inverse of `scalar_encode`. Returns 1 iff the encoded scalar is valid,
+// i.e. all components are < `kPrime`. Otherwise, returns 0 and the value of
+// `out` is undefined.
+template <int BITS>
+int scalar_decode(scalar *out, const uint8_t *in);
+
+template <>
+int scalar_decode<10>(scalar *out, const uint8_t in[320]) {
+  for (int i = 0; i < DEGREE; i += 4) {
+    uint16_t s0 = (uint16_t)(in[0] | ((in[1] & 0x03) << 8));
+    uint16_t s1 = (uint16_t)((in[1] >> 2) | ((in[2] & 0x0f) << 6));
+    uint16_t s2 = (uint16_t)((in[2] >> 4) | ((in[3] & 0x3f) << 4));
+    uint16_t s3 = (uint16_t)((in[3] >> 6) | (in[4] << 2));
+    out->c[i] = s0;
+    out->c[i + 1] = s1;
+    out->c[i + 2] = s2;
+    out->c[i + 3] = s3;
+    in += 5;
+  }
+  return 1;
+}
 
-      int chunk_bits = bits - element_bits_done;
-      if (chunk_bits > in_byte_bits_left) {
-        chunk_bits = in_byte_bits_left;
+template <>
+int scalar_decode<12>(scalar *out, const uint8_t in[384]) {
+  for (int i = 0; i < DEGREE; i += 2) {
+    uint16_t s0 = (uint16_t)(in[0] | ((in[1] & 0x0f) << 8));
+    uint16_t s1 = (uint16_t)((in[1] >> 4) | (in[2] << 4));
+    if (constant_time_declassify_int((s0 | s1) >= kPrime)) {
+      if (s0 >= kPrime || s1 >= kPrime) {
+        return 0;
       }
-
-      element |= (in_byte & kMasks[chunk_bits - 1]) << element_bits_done;
-      in_byte_bits_left -= chunk_bits;
-      in_byte >>= chunk_bits;
-
-      element_bits_done += chunk_bits;
     }
+    out->c[i] = s0;
+    out->c[i + 1] = s1;
+    in += 3;
+  }
+  return 1;
+}
 
-    // An element is only out of range in the case of invalid input, in which
-    // case it is okay to leak the comparison.
-    if (constant_time_declassify_int(element >= kPrime)) {
-      return 0;
-    }
-    out->c[i] = element;
+template <>
+int scalar_decode<4>(scalar *out, const uint8_t in[128]) {
+  for (int i = 0; i < DEGREE; i += 2) {
+    uint16_t s0 = (uint16_t)(in[0] & 0x0f);
+    uint16_t s1 = (uint16_t)(in[0] >> 4);
+    // kPrime is 3329, so 4-bit values are always < kPrime.
+    out->c[i] = s0;
+    out->c[i + 1] = s1;
+    in += 1;
+  }
+  return 1;
+}
+
+// scalar_decode parses `DEGREE * bits` bits from `in` into `DEGREE` values in
+// `out`. It returns one on success and zero if any parsed value is >=
+// `kPrime`.
+template <>
+int scalar_decode<11>(scalar *out, const uint8_t in[352]) {
+  for (int i = 0; i < DEGREE; i += 8) {
+    uint16_t s0 = (uint16_t)(in[0] | ((in[1] & 0x07) << 8));
+    uint16_t s1 = (uint16_t)((in[1] >> 3) | ((in[2] & 0x3f) << 5));
+    uint16_t s2 =
+        (uint16_t)((in[2] >> 6) | (in[3] << 2) | ((in[4] & 0x01) << 10));
+    uint16_t s3 = (uint16_t)((in[4] >> 1) | ((in[5] & 0x0f) << 7));
+    uint16_t s4 = (uint16_t)((in[5] >> 4) | ((in[6] & 0x7f) << 4));
+    uint16_t s5 =
+        (uint16_t)((in[6] >> 7) | (in[7] << 1) | ((in[8] & 0x03) << 9));
+    uint16_t s6 = (uint16_t)((in[8] >> 2) | ((in[9] & 0x1f) << 6));
+    uint16_t s7 = (uint16_t)((in[9] >> 5) | (in[10] << 3));
+    out->c[i] = s0;
+    out->c[i + 1] = s1;
+    out->c[i + 2] = s2;
+    out->c[i + 3] = s3;
+    out->c[i + 4] = s4;
+    out->c[i + 5] = s5;
+    out->c[i + 6] = s6;
+    out->c[i + 7] = s7;
+    in += 11;
   }
+  return 1;
+}
 
+template <>
+int scalar_decode<5>(scalar *out, const uint8_t in[160]) {
+  for (int i = 0; i < DEGREE; i += 8) {
+    uint16_t s0 = (uint16_t)(in[0] & 0x1f);
+    uint16_t s1 = (uint16_t)((in[0] >> 5) | ((in[1] & 0x03) << 3));
+    uint16_t s2 = (uint16_t)((in[1] >> 2) & 0x1f);
+    uint16_t s3 = (uint16_t)((in[1] >> 7) | ((in[2] & 0x0f) << 1));
+    uint16_t s4 = (uint16_t)((in[2] >> 4) | ((in[3] & 0x01) << 4));
+    uint16_t s5 = (uint16_t)((in[3] >> 1) & 0x1f);
+    uint16_t s6 = (uint16_t)((in[3] >> 6) | ((in[4] & 0x07) << 2));
+    uint16_t s7 = (uint16_t)(in[4] >> 3);
+    // kPrime is 3329, so 5-bit values are always < kPrime.
+    out->c[i] = s0;
+    out->c[i + 1] = s1;
+    out->c[i + 2] = s2;
+    out->c[i + 3] = s3;
+    out->c[i + 4] = s4;
+    out->c[i + 5] = s5;
+    out->c[i + 6] = s6;
+    out->c[i + 7] = s7;
+    in += 5;
+  }
   return 1;
 }
 
-// scalar_decode_1 is |scalar_decode| specialised for |bits| == 1.
-void scalar_decode_1(scalar *out, const uint8_t in[32]) {
+template <>
+int scalar_decode<1>(scalar *out, const uint8_t in[32]) {
   for (int i = 0; i < DEGREE; i += 8) {
-    uint8_t in_byte = *in;
-    in++;
+    uint8_t in_byte = in[i / 8];
     for (int j = 0; j < 8; j++) {
-      out->c[i + j] = in_byte & 1;
-      in_byte >>= 1;
+      out->c[i + j] = (in_byte >> j) & 1;
     }
   }
+  return 1;
 }
 
-// Decodes 32*|RANK|*|bits| bytes from |in| into |out|. It returns one on
-// success or zero if any parsed value is >= |kPrime|.
-template <int RANK>
-static int vector_decode(vector<RANK> *out, const uint8_t *in, int bits) {
+// Decodes 32*`RANK`*`bits` bytes from `in` into `out`. It returns one on
+// success or zero if any parsed value is >= `kPrime`.
+template <int bits, int RANK>
+inline int vector_decode(vector<RANK> *out, const uint8_t *in) {
   for (int i = 0; i < RANK; i++) {
-    if (!scalar_decode(&out->v[i], in + i * bits * DEGREE / 8, bits)) {
+    if (!scalar_decode<bits>(&out->v[i], in + i * bits * DEGREE / 8)) {
       return 0;
     }
   }
   return 1;
 }
 
-// Compresses (lossily) an input |x| mod 3329 into |bits| many bits by grouping
+// Compresses (lossily) an input `x` mod 3329 into `bits` many bits by grouping
 // numbers close to each other together. The formula used is
-// round(2^|bits|/kPrime*x) mod 2^|bits|.
+// round(2^`bits`/kPrime*x) mod 2^`bits`.
 // Uses Barrett reduction to achieve constant time. Since we need both the
 // remainder (for rounding) and the quotient (as the result), we cannot use
-// |reduce| here, but need to do the Barrett reduction directly.
-static uint16_t compress(uint16_t x, int bits) {
+// `reduce` here, but need to do the Barrett reduction directly.
+inline uint16_t compress(uint16_t x, int bits) {
   uint32_t shifted = (uint32_t)x << bits;
   uint64_t product = (uint64_t)shifted * kBarrettMultiplier;
   uint32_t quotient = (uint32_t)(product >> kBarrettShift);
@@ -602,30 +865,30 @@ static uint16_t compress(uint16_t x, int bits) {
   return quotient & ((1 << bits) - 1);
 }
 
-// Decompresses |x| by using an equi-distant representative. The formula is
-// round(kPrime/2^|bits|*x). Note that 2^|bits| being the divisor allows us to
+// Decompresses `x` by using an equi-distant representative. The formula is
+// round(kPrime/2^`bits`*x). Note that 2^`bits` being the divisor allows us to
 // implement this logic using only bit operations.
-uint16_t decompress(uint16_t x, int bits) {
+inline uint16_t decompress(uint16_t x, int bits) {
   uint32_t product = (uint32_t)x * kPrime;
   uint32_t power = 1 << bits;
-  // This is |product| % power, since |power| is a power of 2.
+  // This is `product` % power, since `power` is a power of 2.
   uint32_t remainder = product & (power - 1);
-  // This is |product| / power, since |power| is a power of 2.
+  // This is `product` / power, since `power` is a power of 2.
   uint32_t lower = product >> bits;
-  // The rounding logic works since the first half of numbers mod |power| have a
-  // 0 as first bit, and the second half has a 1 as first bit, since |power| is
-  // a power of 2. As a 12 bit number, |remainder| is always positive, so we
+  // The rounding logic works since the first half of numbers mod `power` have a
+  // 0 as first bit, and the second half has a 1 as first bit, since `power` is
+  // a power of 2. As a 12 bit number, `remainder` is always positive, so we
   // will shift in 0s for a right shift.
   return lower + (remainder >> (bits - 1));
 }
 
-static void scalar_compress(scalar *s, int bits) {
+inline void scalar_compress(scalar *s, int bits) {
   for (int i = 0; i < DEGREE; i++) {
     s->c[i] = compress(s->c[i], bits);
   }
 }
 
-static void scalar_decompress(scalar *s, int bits) {
+inline void scalar_decompress(scalar *s, int bits) {
   for (int i = 0; i < DEGREE; i++) {
     s->c[i] = decompress(s->c[i], bits);
   }
@@ -655,51 +918,56 @@ struct public_key {
 
 template <int RANK>
 struct private_key {
-  struct public_key<RANK> pub;
+  public_key<RANK> pub;
   vector<RANK> s;
   uint8_t fo_failure_secret[32];
 };
 
 template <int RANK>
-static void decrypt_cpa(
-    uint8_t out[32], const struct private_key<RANK> *priv,
-    const uint8_t ciphertext[BCM_MLKEM768_CIPHERTEXT_BYTES]) {
+inline void decrypt_cpa(uint8_t out[32], const private_key<RANK> *priv,
+                        const uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES]) {
   constexpr int du = RANK == RANK768 ? kDU768 : kDU1024;
   constexpr int dv = RANK == RANK768 ? kDV768 : kDV1024;
 
   vector<RANK> u;
-  vector_decode(&u, ciphertext, du);
+  vector_decode<du>(&u, ciphertext);
   vector_decompress(&u, du);
   vector_ntt(&u);
   scalar v;
-  scalar_decode(&v, ciphertext + compressed_vector_size(RANK), dv);
+  scalar_decode<dv>(&v, ciphertext + compressed_vector_size(RANK));
   scalar_decompress(&v, dv);
   scalar mask;
   scalar_inner_product(&mask, &priv->s, &u);
   scalar_inverse_ntt(&mask);
   scalar_sub(&v, &mask);
   scalar_compress(&v, 1);
-  scalar_encode_1(out, &v);
+  scalar_encode<1>(out, &v);
 }
 
 template <int RANK>
-static bcm_status mlkem_marshal_public_key(CBB *out,
-                                           const struct public_key<RANK> *pub) {
+inline bcm_status mlkem_marshal_public_key(CBB *out,
+                                           const public_key<RANK> *pub) {
   uint8_t *vector_output;
   if (!CBB_add_space(out, &vector_output, encoded_vector_size(RANK))) {
     return bcm_status::failure;
   }
-  vector_encode(vector_output, &pub->t, kLog2Prime);
+  vector_encode<kLog2Prime>(vector_output, &pub->t);
   if (!CBB_add_bytes(out, pub->rho, sizeof(pub->rho))) {
     return bcm_status::failure;
   }
   return bcm_status::approved;
 }
 
+template <int RANK>
+inline bool mlkem_public_keys_equal(const public_key<RANK> *a,
+                                    const public_key<RANK> *b) {
+  return OPENSSL_memcmp(a->public_key_hash, b->public_key_hash, 32) == 0;
+}
+
 template <int RANK>
 void mlkem_generate_key_external_seed_no_self_test(
     uint8_t *out_encoded_public_key, private_key<RANK> *priv,
-    const uint8_t seed[BCM_MLKEM_SEED_BYTES]) {
+    const uint8_t seed[MLKEM_SEED_BYTES]) {
   uint8_t augmented_seed[33];
   OPENSSL_memcpy(augmented_seed, seed, 32);
   augmented_seed[32] = RANK;
@@ -735,20 +1003,20 @@ void mlkem_generate_key_external_seed_no_self_test(
 }
 
 template <int RANK>
-void mlkem_generate_key_external_seed(
-    uint8_t *out_encoded_public_key, private_key<RANK> *priv,
-    const uint8_t seed[BCM_MLKEM_SEED_BYTES]) {
+void mlkem_generate_key_external_seed(uint8_t *out_encoded_public_key,
+                                      private_key<RANK> *priv,
+                                      const uint8_t seed[MLKEM_SEED_BYTES]) {
   fips::ensure_keygen_self_test();
   mlkem_generate_key_external_seed_no_self_test(out_encoded_public_key, priv,
                                                 seed);
 }
 
 // Encrypts a message with given randomness to
-// the ciphertext in |out|. Without applying the Fujisaki-Okamoto transform this
+// the ciphertext in `out`. Without applying the Fujisaki-Okamoto transform this
 // would not result in a CCA secure scheme, since lattice schemes are vulnerable
 // to decryption failure oracles.
 template <int RANK>
-void encrypt_cpa(uint8_t *out, const struct mlkem::public_key<RANK> *pub,
+void encrypt_cpa(uint8_t *out, const mlkem::public_key<RANK> *pub,
                  const uint8_t message[32], const uint8_t randomness[32]) {
   constexpr int du = RANK == RANK768 ? mlkem::kDU768 : mlkem::kDU1024;
   constexpr int dv = RANK == RANK768 ? mlkem::kDV768 : mlkem::kDV1024;
@@ -773,20 +1041,20 @@ void encrypt_cpa(uint8_t *out, const struct mlkem::public_key<RANK> *pub,
   scalar_inverse_ntt(&v);
   scalar_add(&v, &scalar_error);
   mlkem::scalar expanded_message;
-  scalar_decode_1(&expanded_message, message);
+  scalar_decode<1>(&expanded_message, message);
   scalar_decompress(&expanded_message, 1);
   scalar_add(&v, &expanded_message);
   vector_compress(&u, du);
-  vector_encode(out, &u, du);
+  vector_encode<du>(out, &u);
   scalar_compress(&v, dv);
-  scalar_encode(out + mlkem::compressed_vector_size(RANK), &v, dv);
+  scalar_encode<dv>(out + mlkem::compressed_vector_size(RANK), &v);
 }
 
 // See section 6.3
 template <int RANK>
 void mlkem_decap_no_self_test(
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const uint8_t *ciphertext, const struct private_key<RANK> *priv) {
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const uint8_t *ciphertext, const private_key<RANK> *priv) {
   uint8_t decrypted[64];
   decrypt_cpa(decrypted, priv, ciphertext);
   OPENSSL_memcpy(decrypted + 32, priv->pub.public_key_hash,
@@ -794,8 +1062,8 @@ void mlkem_decap_no_self_test(
   uint8_t key_and_randomness[64];
   hash_g(key_and_randomness, decrypted, sizeof(decrypted));
   constexpr size_t ciphertext_len = ciphertext_size(RANK);
-  uint8_t expected_ciphertext[BCM_MLKEM1024_CIPHERTEXT_BYTES];
-  static_assert(ciphertext_len <= sizeof(expected_ciphertext), "");
+  uint8_t expected_ciphertext[MLKEM1024_CIPHERTEXT_BYTES];
+  static_assert(ciphertext_len <= sizeof(expected_ciphertext));
   encrypt_cpa(expected_ciphertext, &priv->pub, decrypted,
               key_and_randomness + 32);
 
@@ -804,53 +1072,56 @@ void mlkem_decap_no_self_test(
 
   uint8_t mask = constant_time_eq_int_8(
       CRYPTO_memcmp(ciphertext, expected_ciphertext, ciphertext_len), 0);
-  for (int i = 0; i < BCM_MLKEM_SHARED_SECRET_BYTES; i++) {
+  for (int i = 0; i < MLKEM_SHARED_SECRET_BYTES; i++) {
     out_shared_secret[i] =
         constant_time_select_8(mask, key_and_randomness[i], failure_key[i]);
   }
 }
 
 template <int RANK>
-void mlkem_decap(uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-                 const uint8_t *ciphertext,
-                 const struct private_key<RANK> *priv) {
+void mlkem_decap(uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+                 const uint8_t *ciphertext, const private_key<RANK> *priv) {
   fips::ensure_decap_self_test();
   mlkem_decap_no_self_test(out_shared_secret, ciphertext, priv);
 }
 
-// mlkem_parse_public_key_no_hash parses |in| into |pub| but doesn't calculate
-// the value of |pub->public_key_hash|.
+// mlkem_parse_public_key_with_trailing_data parses `in` into `pub` but leaves
+// trailing data in `in` for the caller.
 template <int RANK>
-int mlkem_parse_public_key_no_hash(struct public_key<RANK> *pub, CBS *in) {
+int mlkem_parse_public_key_with_trailing_data(public_key<RANK> *pub, CBS *in) {
+  CBS orig_in = *in;
   CBS t_bytes;
   if (!CBS_get_bytes(in, &t_bytes, encoded_vector_size(RANK)) ||
-      !vector_decode(&pub->t, CBS_data(&t_bytes), kLog2Prime) ||
+      !vector_decode<kLog2Prime>(&pub->t, CBS_data(&t_bytes)) ||
       !CBS_copy_bytes(in, pub->rho, sizeof(pub->rho))) {
     return 0;
   }
   matrix_expand(&pub->m, pub->rho);
+  size_t pub_key_len = CBS_len(&orig_in) - CBS_len(in);
+  assert(pub_key_len == encoded_public_key_size(RANK));
+  hash_h(pub->public_key_hash, CBS_data(&orig_in), pub_key_len);
   return 1;
 }
 
 template <int RANK>
-int mlkem_parse_public_key(struct public_key<RANK> *pub, CBS *in) {
-  CBS orig_in = *in;
-  if (!mlkem_parse_public_key_no_hash(pub, in) ||  //
+int mlkem_parse_public_key(public_key<RANK> *pub, CBS *in) {
+  if (!mlkem_parse_public_key_with_trailing_data(pub, in) ||  //
       CBS_len(in) != 0) {
     return 0;
   }
-  hash_h(pub->public_key_hash, CBS_data(&orig_in), CBS_len(&orig_in));
   return 1;
 }
 
 template <int RANK>
-int mlkem_parse_private_key(struct private_key<RANK> *priv, CBS *in) {
-  CBS s_bytes;
+int mlkem_parse_private_key(private_key<RANK> *priv, CBS *in) {
+  CBS s_bytes, public_key_hash;
   if (!CBS_get_bytes(in, &s_bytes, encoded_vector_size(RANK)) ||
-      !vector_decode(&priv->s, CBS_data(&s_bytes), kLog2Prime) ||
-      !mlkem_parse_public_key_no_hash(&priv->pub, in) ||
-      !CBS_copy_bytes(in, priv->pub.public_key_hash,
-                      sizeof(priv->pub.public_key_hash)) ||
+      !vector_decode<kLog2Prime>(&priv->s, CBS_data(&s_bytes)) ||
+      !mlkem_parse_public_key_with_trailing_data(&priv->pub, in) ||
+      // We compute the public key hash ourselves, but check it matched.
+      !CBS_get_bytes(in, &public_key_hash, sizeof(priv->pub.public_key_hash)) ||
+      !CBS_mem_equal(&public_key_hash, priv->pub.public_key_hash,
+                     sizeof(priv->pub.public_key_hash)) ||
       !CBS_copy_bytes(in, priv->fo_failure_secret,
                       sizeof(priv->fo_failure_secret)) ||
       CBS_len(in) != 0) {
@@ -860,12 +1131,12 @@ int mlkem_parse_private_key(struct private_key<RANK> *priv, CBS *in) {
 }
 
 template <int RANK>
-int mlkem_marshal_private_key(CBB *out, const struct private_key<RANK> *priv) {
+int mlkem_marshal_private_key(CBB *out, const private_key<RANK> *priv) {
   uint8_t *s_output;
   if (!CBB_add_space(out, &s_output, encoded_vector_size(RANK))) {
     return 0;
   }
-  vector_encode(s_output, &priv->s, kLog2Prime);
+  vector_encode<kLog2Prime>(s_output, &priv->s);
   if (!bcm_success(mlkem_marshal_public_key(out, &priv->pub)) ||
       !CBB_add_bytes(out, priv->pub.public_key_hash,
                      sizeof(priv->pub.public_key_hash)) ||
@@ -876,56 +1147,69 @@ int mlkem_marshal_private_key(CBB *out, const struct private_key<RANK> *priv) {
   return 1;
 }
 
-struct public_key<RANK768> *public_key_768_from_external(
-    const struct BCM_mlkem768_public_key *external) {
-  static_assert(sizeof(struct BCM_mlkem768_public_key) >=
-                    sizeof(struct public_key<RANK768>),
-                "MLKEM public key is too small");
-  static_assert(alignof(struct BCM_mlkem768_public_key) >=
-                    alignof(struct public_key<RANK768>),
-                "MLKEM public key alignment incorrect");
-  return (struct public_key<RANK768> *)external;
-}
-
-static struct public_key<RANK1024> *public_key_1024_from_external(
-    const struct BCM_mlkem1024_public_key *external) {
-  static_assert(sizeof(struct BCM_mlkem1024_public_key) >=
-                    sizeof(struct public_key<RANK1024>),
-                "MLKEM1024 public key is too small");
-  static_assert(alignof(struct BCM_mlkem1024_public_key) >=
-                    alignof(struct public_key<RANK1024>),
-                "MLKEM1024 public key alignment incorrect");
-  return (struct public_key<RANK1024> *)external;
-}
-
-struct private_key<RANK768> *private_key_768_from_external(
-    const struct BCM_mlkem768_private_key *external) {
-  static_assert(sizeof(struct BCM_mlkem768_private_key) >=
-                    sizeof(struct private_key<RANK768>),
-                "MLKEM private key too small");
-  static_assert(alignof(struct BCM_mlkem768_private_key) >=
-                    alignof(struct private_key<RANK768>),
-                "MLKEM private key alignment incorrect");
-  return (struct private_key<RANK768> *)external;
-}
-
-struct private_key<RANK1024> *private_key_1024_from_external(
-    const struct BCM_mlkem1024_private_key *external) {
-  static_assert(sizeof(struct BCM_mlkem1024_private_key) >=
-                    sizeof(struct private_key<RANK1024>),
-                "MLKEM1024 private key too small");
-  static_assert(alignof(struct BCM_mlkem1024_private_key) >=
-                    alignof(struct private_key<RANK1024>),
-                "MLKEM1024 private key alignment incorrect");
-  return (struct private_key<RANK1024> *)external;
+static_assert(sizeof(MLKEM768_public_key) == sizeof(public_key<RANK768>));
+static_assert(alignof(MLKEM768_public_key) == alignof(public_key<RANK768>));
+
+const public_key<RANK768> *public_key_768_from_external(
+    const MLKEM768_public_key *external) {
+  return reinterpret_cast<const public_key<RANK768> *>(external);
+}
+public_key<RANK768> *public_key_768_from_external(
+    MLKEM768_public_key *external) {
+  return reinterpret_cast<public_key<RANK768> *>(external);
+}
+const MLKEM768_public_key *public_key_768_to_external(
+    const public_key<RANK768> *public_key) {
+  return reinterpret_cast<const MLKEM768_public_key *>(public_key);
+}
+
+static_assert(sizeof(MLKEM1024_public_key) == sizeof(public_key<RANK1024>));
+static_assert(alignof(MLKEM1024_public_key) == alignof(public_key<RANK1024>));
+
+const public_key<RANK1024> *public_key_1024_from_external(
+    const MLKEM1024_public_key *external) {
+  return reinterpret_cast<const public_key<RANK1024> *>(external);
+}
+public_key<RANK1024> *public_key_1024_from_external(
+    MLKEM1024_public_key *external) {
+  return reinterpret_cast<public_key<RANK1024> *>(external);
+}
+
+const MLKEM1024_public_key *public_key_1024_to_external(
+    const public_key<RANK1024> *public_key) {
+  return reinterpret_cast<const MLKEM1024_public_key *>(public_key);
+}
+
+static_assert(sizeof(MLKEM768_private_key) >= sizeof(private_key<RANK768>));
+static_assert(alignof(MLKEM768_private_key) >= alignof(private_key<RANK768>));
+
+const private_key<RANK768> *private_key_768_from_external(
+    const MLKEM768_private_key *external) {
+  return reinterpret_cast<const private_key<RANK768> *>(external);
+}
+private_key<RANK768> *private_key_768_from_external(
+    MLKEM768_private_key *external) {
+  return reinterpret_cast<private_key<RANK768> *>(external);
+}
+
+static_assert(sizeof(MLKEM1024_private_key) >= sizeof(private_key<RANK1024>));
+static_assert(alignof(MLKEM1024_private_key) >= alignof(private_key<RANK1024>));
+
+const private_key<RANK1024> *private_key_1024_from_external(
+    const MLKEM1024_private_key *external) {
+  return reinterpret_cast<const private_key<RANK1024> *>(external);
+}
+private_key<RANK1024> *private_key_1024_from_external(
+    MLKEM1024_private_key *external) {
+  return reinterpret_cast<private_key<RANK1024> *>(external);
 }
 
 // See section 6.2.
 template <int RANK>
 void mlkem_encap_external_entropy_no_self_test(
     uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const struct mlkem::public_key<RANK> *pub,
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const mlkem::public_key<RANK> *pub,
     const uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY]) {
   uint8_t input[64];
   OPENSSL_memcpy(input, entropy, BCM_MLKEM_ENCAP_ENTROPY);
@@ -936,15 +1220,15 @@ void mlkem_encap_external_entropy_no_self_test(
   encrypt_cpa(out_ciphertext, pub, entropy, key_and_randomness + 32);
   // The ciphertext is public.
   CONSTTIME_DECLASSIFY(out_ciphertext, mlkem::ciphertext_size(RANK));
-  static_assert(BCM_MLKEM_SHARED_SECRET_BYTES == 32, "");
+  static_assert(MLKEM_SHARED_SECRET_BYTES == 32);
   memcpy(out_shared_secret, key_and_randomness, 32);
 }
 
 template <int RANK>
 void mlkem_encap_external_entropy(
     uint8_t *out_ciphertext,
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const struct mlkem::public_key<RANK> *pub,
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const mlkem::public_key<RANK> *pub,
     const uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY]) {
   fips::ensure_encap_self_test();
   mlkem_encap_external_entropy_no_self_test(out_ciphertext, out_shared_secret,
@@ -955,75 +1239,67 @@ namespace fips {
 
 #include "fips_known_values.inc"
 
-static int keygen_self_test() {
-  uint8_t pub_key[BCM_MLKEM768_PUBLIC_KEY_BYTES];
+inline int keygen_self_test() {
+  uint8_t pub_key[MLKEM768_PUBLIC_KEY_BYTES];
   private_key<RANK768> priv;
-  static_assert(sizeof(kTestEntropy) >= BCM_MLKEM_SEED_BYTES);
+  static_assert(sizeof(kTestEntropy) >= MLKEM_SEED_BYTES);
   mlkem_generate_key_external_seed_no_self_test(pub_key, &priv, kTestEntropy);
   CBB cbb;
   constexpr size_t kMarshaledPrivateKeySize = 2400;
   uint8_t priv_bytes[kMarshaledPrivateKeySize];
   CBB_init_fixed(&cbb, priv_bytes, sizeof(priv_bytes));
-  static_assert(sizeof(kExpectedPrivateKeyBytes) == kMarshaledPrivateKeySize);
-  static_assert(sizeof(kExpectedPublicKeyBytes) == sizeof(pub_key));
   if (!mlkem_marshal_private_key(&cbb, &priv) ||
       !BORINGSSL_check_test(kExpectedPrivateKeyBytes, priv_bytes,
-                            sizeof(priv_bytes), "ML-KEM keygen private key") ||
-      !BORINGSSL_check_test(kExpectedPublicKeyBytes, pub_key, sizeof(pub_key),
+                            "ML-KEM keygen private key") ||
+      !BORINGSSL_check_test(kExpectedPublicKeyBytes, pub_key,
                             "ML-KEM keygen public key")) {
     return 0;
   }
   return 1;
 }
 
-static int encap_self_test() {
+inline int encap_self_test() {
   CBS cbs;
   CBS_init(&cbs, kExpectedPublicKeyBytes, sizeof(kExpectedPublicKeyBytes));
   public_key<RANK768> pub;
   if (!mlkem_parse_public_key(&pub, &cbs)) {
     return 0;
   }
-  uint8_t ciphertext[BCM_MLKEM768_CIPHERTEXT_BYTES];
-  uint8_t shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES];
+  uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
   static_assert(sizeof(kTestEntropy) >= BCM_MLKEM_ENCAP_ENTROPY);
   mlkem_encap_external_entropy_no_self_test(ciphertext, shared_secret, &pub,
                                             kTestEntropy);
-  if (!BORINGSSL_check_test(ciphertext, kExpectedCiphertext, sizeof(ciphertext),
+  if (!BORINGSSL_check_test(ciphertext, kExpectedCiphertext,
                             "ML-KEM encap ciphertext") ||
       !BORINGSSL_check_test(kExpectedSharedSecret, shared_secret,
-                            sizeof(kExpectedSharedSecret),
                             "ML-KEM encap shared secret")) {
     return 0;
   }
   return 1;
 }
 
-static int decap_self_test() {
+inline int decap_self_test() {
   CBS cbs;
   CBS_init(&cbs, kExpectedPrivateKeyBytes, sizeof(kExpectedPrivateKeyBytes));
   private_key<RANK768> priv;
   if (!mlkem_parse_private_key(&priv, &cbs)) {
     return 0;
   }
-  uint8_t shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
   mlkem_decap_no_self_test(shared_secret, kExpectedCiphertext, &priv);
-  static_assert(sizeof(kExpectedSharedSecret) == sizeof(shared_secret));
   if (!BORINGSSL_check_test(kExpectedSharedSecret, shared_secret,
-                            sizeof(shared_secret),
                             "ML-KEM decap shared secret")) {
     return 0;
   }
 
-  uint8_t implicit_rejection_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t implicit_rejection_shared_secret[MLKEM_SHARED_SECRET_BYTES];
   static_assert(sizeof(kExpectedPrivateKeyBytes) >=
                 sizeof(kExpectedCiphertext));
   mlkem_decap_no_self_test(implicit_rejection_shared_secret,
                            kExpectedPrivateKeyBytes, &priv);
-  static_assert(sizeof(kExpectedImplicitRejectionSharedSecret) ==
-                sizeof(implicit_rejection_shared_secret));
   if (!BORINGSSL_check_test(kExpectedImplicitRejectionSharedSecret,
                             implicit_rejection_shared_secret,
-                            sizeof(implicit_rejection_shared_secret),
                             "ML-KEM decap implicit rejection shared secret")) {
     return 0;
   }
@@ -1032,9 +1308,7 @@ static int decap_self_test() {
 
 #if defined(BORINGSSL_FIPS)
 
-DEFINE_STATIC_ONCE(g_mlkem_keygen_self_test_once)
-
-void ensure_keygen_self_test(void) {
+void ensure_keygen_self_test() {
   CRYPTO_once(g_mlkem_keygen_self_test_once_bss_get(), []() {
     if (!keygen_self_test()) {
       BORINGSSL_FIPS_abort();
@@ -1042,9 +1316,7 @@ void ensure_keygen_self_test(void) {
   });
 }
 
-DEFINE_STATIC_ONCE(g_mlkem_encap_self_test_once)
-
-void ensure_encap_self_test(void) {
+void ensure_encap_self_test() {
   CRYPTO_once(g_mlkem_encap_self_test_once_bss_get(), []() {
     if (!encap_self_test()) {
       BORINGSSL_FIPS_abort();
@@ -1052,9 +1324,7 @@ void ensure_encap_self_test(void) {
   });
 }
 
-DEFINE_STATIC_ONCE(g_mlkem_decap_self_test_once)
-
-void ensure_decap_self_test(void) {
+void ensure_decap_self_test() {
   CRYPTO_once(g_mlkem_decap_self_test_once_bss_get(), []() {
     if (!decap_self_test()) {
       BORINGSSL_FIPS_abort();
@@ -1064,9 +1334,9 @@ void ensure_decap_self_test(void) {
 
 #else
 
-void ensure_keygen_self_test(void) {}
-void ensure_encap_self_test(void) {}
-void ensure_decap_self_test(void) {}
+void ensure_keygen_self_test() {}
+void ensure_encap_self_test() {}
+void ensure_decap_self_test() {}
 
 #endif
 }  // namespace fips
@@ -1074,14 +1344,14 @@ void ensure_decap_self_test(void) {}
 }  // namespace
 }  // namespace mlkem
 
-bcm_status BCM_mlkem768_check_fips(
-    const struct BCM_mlkem768_private_key *private_key) {
-  mlkem::private_key<RANK768> *priv =
+bcm_status bssl::BCM_mlkem768_check_fips(
+    const MLKEM768_private_key *private_key) {
+  const mlkem::private_key<RANK768> *priv =
       mlkem::private_key_768_from_external(private_key);
 
   const uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY] = {1, 2, 3, 4};
-  uint8_t ciphertext[BCM_MLKEM768_CIPHERTEXT_BYTES];
-  uint8_t shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t ciphertext[MLKEM768_CIPHERTEXT_BYTES];
+  uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
   mlkem_encap_external_entropy_no_self_test(ciphertext, shared_secret,
                                             &priv->pub, entropy);
 
@@ -1089,7 +1359,7 @@ bcm_status BCM_mlkem768_check_fips(
     shared_secret[0] ^= 1;
   }
 
-  uint8_t shared_secret2[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES];
   mlkem::mlkem_decap_no_self_test(shared_secret2, ciphertext, priv);
   if (CRYPTO_memcmp(shared_secret, shared_secret2, sizeof(shared_secret)) !=
       0) {
@@ -1098,10 +1368,10 @@ bcm_status BCM_mlkem768_check_fips(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem768_generate_key_fips(
-    uint8_t out_encoded_public_key[BCM_MLKEM768_PUBLIC_KEY_BYTES],
-    uint8_t optional_out_seed[BCM_MLKEM_SEED_BYTES],
-    struct BCM_mlkem768_private_key *out_private_key) {
+bcm_status bssl::BCM_mlkem768_generate_key_fips(
+    uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
+    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
+    MLKEM768_private_key *out_private_key) {
   if (out_encoded_public_key == nullptr || out_private_key == nullptr) {
     return bcm_status::failure;
   }
@@ -1110,11 +1380,11 @@ bcm_status BCM_mlkem768_generate_key_fips(
   return BCM_mlkem768_check_fips(out_private_key);
 }
 
-bcm_infallible BCM_mlkem768_generate_key(
-    uint8_t out_encoded_public_key[BCM_MLKEM768_PUBLIC_KEY_BYTES],
-    uint8_t optional_out_seed[BCM_MLKEM_SEED_BYTES],
-    struct BCM_mlkem768_private_key *out_private_key) {
-  uint8_t seed[BCM_MLKEM_SEED_BYTES];
+bcm_infallible bssl::BCM_mlkem768_generate_key(
+    uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
+    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
+    MLKEM768_private_key *out_private_key) {
+  uint8_t seed[MLKEM_SEED_BYTES];
   BCM_rand_bytes(seed, sizeof(seed));
   CONSTTIME_SECRET(seed, sizeof(seed));
   if (optional_out_seed) {
@@ -1125,27 +1395,27 @@ bcm_infallible BCM_mlkem768_generate_key(
   return bcm_infallible::not_approved;
 }
 
-bcm_status BCM_mlkem768_private_key_from_seed(
-    struct BCM_mlkem768_private_key *out_private_key, const uint8_t *seed,
+bcm_status bssl::BCM_mlkem768_private_key_from_seed(
+    MLKEM768_private_key *out_private_key, const uint8_t *seed,
     size_t seed_len) {
-  if (seed_len != BCM_MLKEM_SEED_BYTES) {
+  if (seed_len != MLKEM_SEED_BYTES) {
     return bcm_status::failure;
   }
 
-  uint8_t public_key_bytes[BCM_MLKEM768_PUBLIC_KEY_BYTES];
+  uint8_t public_key_bytes[MLKEM768_PUBLIC_KEY_BYTES];
   BCM_mlkem768_generate_key_external_seed(public_key_bytes, out_private_key,
                                           seed);
   return bcm_status::not_approved;
 }
 
-bcm_status BCM_mlkem1024_check_fips(
-    const struct BCM_mlkem1024_private_key *private_key) {
-  mlkem::private_key<RANK1024> *priv =
+bcm_status bssl::BCM_mlkem1024_check_fips(
+    const MLKEM1024_private_key *private_key) {
+  const mlkem::private_key<RANK1024> *priv =
       mlkem::private_key_1024_from_external(private_key);
 
   const uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY] = {1, 2, 3, 4};
-  uint8_t ciphertext[BCM_MLKEM1024_CIPHERTEXT_BYTES];
-  uint8_t shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t ciphertext[MLKEM1024_CIPHERTEXT_BYTES];
+  uint8_t shared_secret[MLKEM_SHARED_SECRET_BYTES];
   mlkem_encap_external_entropy_no_self_test(ciphertext, shared_secret,
                                             &priv->pub, entropy);
 
@@ -1153,7 +1423,7 @@ bcm_status BCM_mlkem1024_check_fips(
     shared_secret[0] ^= 1;
   }
 
-  uint8_t shared_secret2[BCM_MLKEM_SHARED_SECRET_BYTES];
+  uint8_t shared_secret2[MLKEM_SHARED_SECRET_BYTES];
   mlkem::mlkem_decap_no_self_test(shared_secret2, ciphertext, priv);
   if (CRYPTO_memcmp(shared_secret, shared_secret2, sizeof(shared_secret)) !=
       0) {
@@ -1162,10 +1432,10 @@ bcm_status BCM_mlkem1024_check_fips(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem1024_generate_key_fips(
-    uint8_t out_encoded_public_key[BCM_MLKEM1024_PUBLIC_KEY_BYTES],
-    uint8_t optional_out_seed[BCM_MLKEM_SEED_BYTES],
-    struct BCM_mlkem1024_private_key *out_private_key) {
+bcm_status bssl::BCM_mlkem1024_generate_key_fips(
+    uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
+    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
+    MLKEM1024_private_key *out_private_key) {
   if (out_encoded_public_key == nullptr || out_private_key == nullptr) {
     return bcm_status::failure;
   }
@@ -1174,11 +1444,11 @@ bcm_status BCM_mlkem1024_generate_key_fips(
   return BCM_mlkem1024_check_fips(out_private_key);
 }
 
-bcm_infallible BCM_mlkem1024_generate_key(
-    uint8_t out_encoded_public_key[BCM_MLKEM1024_PUBLIC_KEY_BYTES],
-    uint8_t optional_out_seed[BCM_MLKEM_SEED_BYTES],
-    struct BCM_mlkem1024_private_key *out_private_key) {
-  uint8_t seed[BCM_MLKEM_SEED_BYTES];
+bcm_infallible bssl::BCM_mlkem1024_generate_key(
+    uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
+    uint8_t optional_out_seed[MLKEM_SEED_BYTES],
+    MLKEM1024_private_key *out_private_key) {
+  uint8_t seed[MLKEM_SEED_BYTES];
   BCM_rand_bytes(seed, sizeof(seed));
   CONSTTIME_SECRET(seed, sizeof(seed));
   if (optional_out_seed) {
@@ -1189,66 +1459,80 @@ bcm_infallible BCM_mlkem1024_generate_key(
   return bcm_infallible::not_approved;
 }
 
-bcm_status BCM_mlkem1024_private_key_from_seed(
-    struct BCM_mlkem1024_private_key *out_private_key, const uint8_t *seed,
+bcm_status bssl::BCM_mlkem1024_private_key_from_seed(
+    MLKEM1024_private_key *out_private_key, const uint8_t *seed,
     size_t seed_len) {
-  if (seed_len != BCM_MLKEM_SEED_BYTES) {
+  if (seed_len != MLKEM_SEED_BYTES) {
     return bcm_status::failure;
   }
-  uint8_t public_key_bytes[BCM_MLKEM1024_PUBLIC_KEY_BYTES];
+  uint8_t public_key_bytes[MLKEM1024_PUBLIC_KEY_BYTES];
   BCM_mlkem1024_generate_key_external_seed(public_key_bytes, out_private_key,
                                            seed);
   return bcm_status::not_approved;
 }
 
-bcm_infallible BCM_mlkem768_generate_key_external_seed(
-    uint8_t out_encoded_public_key[BCM_MLKEM768_PUBLIC_KEY_BYTES],
-    struct BCM_mlkem768_private_key *out_private_key,
-    const uint8_t seed[BCM_MLKEM_SEED_BYTES]) {
+bcm_infallible bssl::BCM_mlkem768_generate_key_external_seed(
+    uint8_t out_encoded_public_key[MLKEM768_PUBLIC_KEY_BYTES],
+    MLKEM768_private_key *out_private_key,
+    const uint8_t seed[MLKEM_SEED_BYTES]) {
   mlkem::private_key<RANK768> *priv =
       mlkem::private_key_768_from_external(out_private_key);
   mlkem_generate_key_external_seed(out_encoded_public_key, priv, seed);
   return bcm_infallible::approved;
 }
 
-bcm_infallible BCM_mlkem1024_generate_key_external_seed(
-    uint8_t out_encoded_public_key[BCM_MLKEM1024_PUBLIC_KEY_BYTES],
-    struct BCM_mlkem1024_private_key *out_private_key,
-    const uint8_t seed[BCM_MLKEM_SEED_BYTES]) {
+bcm_infallible bssl::BCM_mlkem1024_generate_key_external_seed(
+    uint8_t out_encoded_public_key[MLKEM1024_PUBLIC_KEY_BYTES],
+    MLKEM1024_private_key *out_private_key,
+    const uint8_t seed[MLKEM_SEED_BYTES]) {
   mlkem::private_key<RANK1024> *priv =
       mlkem::private_key_1024_from_external(out_private_key);
   mlkem_generate_key_external_seed(out_encoded_public_key, priv, seed);
   return bcm_infallible::approved;
 }
 
-bcm_infallible BCM_mlkem768_public_from_private(
-    struct BCM_mlkem768_public_key *out_public_key,
-    const struct BCM_mlkem768_private_key *private_key) {
-  struct mlkem::public_key<RANK768> *const pub =
+bcm_infallible bssl::BCM_mlkem768_public_from_private(
+    MLKEM768_public_key *out_public_key,
+    const MLKEM768_private_key *private_key) {
+  mlkem::public_key<RANK768> *const pub =
       mlkem::public_key_768_from_external(out_public_key);
-  const struct mlkem::private_key<RANK768> *const priv =
+  const mlkem::private_key<RANK768> *const priv =
       mlkem::private_key_768_from_external(private_key);
   *pub = priv->pub;
   return bcm_infallible::approved;
 }
 
-bcm_infallible BCM_mlkem1024_public_from_private(
-    struct BCM_mlkem1024_public_key *out_public_key,
-    const struct BCM_mlkem1024_private_key *private_key) {
-  struct mlkem::public_key<RANK1024> *const pub =
+bcm_infallible bssl::BCM_mlkem1024_public_from_private(
+    MLKEM1024_public_key *out_public_key,
+    const MLKEM1024_private_key *private_key) {
+  mlkem::public_key<RANK1024> *const pub =
       mlkem::public_key_1024_from_external(out_public_key);
-  const struct mlkem::private_key<RANK1024> *const priv =
+  const mlkem::private_key<RANK1024> *const priv =
       mlkem::private_key_1024_from_external(private_key);
   *pub = priv->pub;
   return bcm_infallible::approved;
 }
 
-// Calls |MLKEM768_encap_external_entropy| with random bytes from
-// |BCM_rand_bytes|
-bcm_infallible BCM_mlkem768_encap(
-    uint8_t out_ciphertext[BCM_MLKEM768_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const struct BCM_mlkem768_public_key *public_key) {
+const MLKEM768_public_key *bssl::BCM_mlkem768_public_of_private(
+    const MLKEM768_private_key *private_key) {
+  const mlkem::private_key<RANK768> *const priv =
+      mlkem::private_key_768_from_external(private_key);
+  return mlkem::public_key_768_to_external(&priv->pub);
+}
+
+const MLKEM1024_public_key *bssl::BCM_mlkem1024_public_of_private(
+    const MLKEM1024_private_key *private_key) {
+  const mlkem::private_key<RANK1024> *const priv =
+      mlkem::private_key_1024_from_external(private_key);
+  return mlkem::public_key_1024_to_external(&priv->pub);
+}
+
+// Calls `MLKEM768_encap_external_entropy` with random bytes from
+// `BCM_rand_bytes`
+bcm_infallible bssl::BCM_mlkem768_encap(
+    uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const MLKEM768_public_key *public_key) {
   uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY];
   BCM_rand_bytes(entropy, BCM_MLKEM_ENCAP_ENTROPY);
   CONSTTIME_SECRET(entropy, BCM_MLKEM_ENCAP_ENTROPY);
@@ -1257,10 +1541,10 @@ bcm_infallible BCM_mlkem768_encap(
   return bcm_infallible::approved;
 }
 
-bcm_infallible BCM_mlkem1024_encap(
-    uint8_t out_ciphertext[BCM_MLKEM1024_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const struct BCM_mlkem1024_public_key *public_key) {
+bcm_infallible bssl::BCM_mlkem1024_encap(
+    uint8_t out_ciphertext[MLKEM1024_CIPHERTEXT_BYTES],
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const MLKEM1024_public_key *public_key) {
   uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY];
   BCM_rand_bytes(entropy, BCM_MLKEM_ENCAP_ENTROPY);
   CONSTTIME_SECRET(entropy, BCM_MLKEM_ENCAP_ENTROPY);
@@ -1269,71 +1553,91 @@ bcm_infallible BCM_mlkem1024_encap(
   return bcm_infallible::approved;
 }
 
-bcm_infallible BCM_mlkem768_encap_external_entropy(
-    uint8_t out_ciphertext[BCM_MLKEM768_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const struct BCM_mlkem768_public_key *public_key,
+bcm_infallible bssl::BCM_mlkem768_encap_external_entropy(
+    uint8_t out_ciphertext[MLKEM768_CIPHERTEXT_BYTES],
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const MLKEM768_public_key *public_key,
     const uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY]) {
-  const struct mlkem::public_key<RANK768> *pub =
+  const mlkem::public_key<RANK768> *pub =
       mlkem::public_key_768_from_external(public_key);
   mlkem_encap_external_entropy(out_ciphertext, out_shared_secret, pub, entropy);
   return bcm_infallible::approved;
 }
 
-bcm_infallible BCM_mlkem1024_encap_external_entropy(
-    uint8_t out_ciphertext[BCM_MLKEM1024_CIPHERTEXT_BYTES],
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
-    const struct BCM_mlkem1024_public_key *public_key,
+bcm_infallible bssl::BCM_mlkem1024_encap_external_entropy(
+    uint8_t out_ciphertext[MLKEM1024_CIPHERTEXT_BYTES],
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
+    const MLKEM1024_public_key *public_key,
     const uint8_t entropy[BCM_MLKEM_ENCAP_ENTROPY]) {
-  const struct mlkem::public_key<RANK1024> *pub =
+  const mlkem::public_key<RANK1024> *pub =
       mlkem::public_key_1024_from_external(public_key);
   mlkem_encap_external_entropy(out_ciphertext, out_shared_secret, pub, entropy);
   return bcm_infallible::approved;
 }
 
-bcm_status BCM_mlkem768_decap(
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
+bcm_status bssl::BCM_mlkem768_decap(
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
     const uint8_t *ciphertext, size_t ciphertext_len,
-    const struct BCM_mlkem768_private_key *private_key) {
-  if (ciphertext_len != BCM_MLKEM768_CIPHERTEXT_BYTES) {
-    BCM_rand_bytes(out_shared_secret, BCM_MLKEM_SHARED_SECRET_BYTES);
+    const MLKEM768_private_key *private_key) {
+  if (ciphertext_len != MLKEM768_CIPHERTEXT_BYTES) {
+    BCM_rand_bytes(out_shared_secret, MLKEM_SHARED_SECRET_BYTES);
     return bcm_status::failure;
   }
-  const struct mlkem::private_key<RANK768> *priv =
+  const mlkem::private_key<RANK768> *priv =
       mlkem::private_key_768_from_external(private_key);
   mlkem_decap(out_shared_secret, ciphertext, priv);
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem1024_decap(
-    uint8_t out_shared_secret[BCM_MLKEM_SHARED_SECRET_BYTES],
+bcm_status bssl::BCM_mlkem1024_decap(
+    uint8_t out_shared_secret[MLKEM_SHARED_SECRET_BYTES],
     const uint8_t *ciphertext, size_t ciphertext_len,
-    const struct BCM_mlkem1024_private_key *private_key) {
-  if (ciphertext_len != BCM_MLKEM1024_CIPHERTEXT_BYTES) {
-    BCM_rand_bytes(out_shared_secret, BCM_MLKEM_SHARED_SECRET_BYTES);
+    const MLKEM1024_private_key *private_key) {
+  if (ciphertext_len != MLKEM1024_CIPHERTEXT_BYTES) {
+    BCM_rand_bytes(out_shared_secret, MLKEM_SHARED_SECRET_BYTES);
     return bcm_status::failure;
   }
-  const struct mlkem::private_key<RANK1024> *priv =
+  const mlkem::private_key<RANK1024> *priv =
       mlkem::private_key_1024_from_external(private_key);
   mlkem_decap(out_shared_secret, ciphertext, priv);
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem768_marshal_public_key(
-    CBB *out, const struct BCM_mlkem768_public_key *public_key) {
+bcm_status bssl::BCM_mlkem768_marshal_public_key(
+    CBB *out, const MLKEM768_public_key *public_key) {
   return mlkem_marshal_public_key(
       out, mlkem::public_key_768_from_external(public_key));
 }
 
-bcm_status BCM_mlkem1024_marshal_public_key(
-    CBB *out, const struct BCM_mlkem1024_public_key *public_key) {
+bcm_status bssl::BCM_mlkem1024_marshal_public_key(
+    CBB *out, const MLKEM1024_public_key *public_key) {
   return mlkem_marshal_public_key(
       out, mlkem::public_key_1024_from_external(public_key));
 }
 
-bcm_status BCM_mlkem768_parse_public_key(
-    struct BCM_mlkem768_public_key *public_key, CBS *in) {
-  struct mlkem::public_key<RANK768> *pub =
+int bssl::BCM_mlkem768_public_keys_equal(const MLKEM768_public_key *a,
+                                         const MLKEM768_public_key *b) {
+  const auto *pub_a = mlkem::public_key_768_from_external(a);
+  const auto *pub_b = mlkem::public_key_768_from_external(b);
+  if (mlkem_public_keys_equal(pub_a, pub_b)) {
+    return 1;
+  }
+  return 0;
+}
+
+int bssl::BCM_mlkem1024_public_keys_equal(const MLKEM1024_public_key *a,
+                                          const MLKEM1024_public_key *b) {
+  const auto *pub_a = mlkem::public_key_1024_from_external(a);
+  const auto *pub_b = mlkem::public_key_1024_from_external(b);
+  if (mlkem_public_keys_equal(pub_a, pub_b)) {
+    return 1;
+  }
+  return 0;
+}
+
+bcm_status bssl::BCM_mlkem768_parse_public_key(MLKEM768_public_key *public_key,
+                                               CBS *in) {
+  mlkem::public_key<RANK768> *pub =
       mlkem::public_key_768_from_external(public_key);
   if (!mlkem_parse_public_key(pub, in)) {
     return bcm_status::failure;
@@ -1341,9 +1645,9 @@ bcm_status BCM_mlkem768_parse_public_key(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem1024_parse_public_key(
-    struct BCM_mlkem1024_public_key *public_key, CBS *in) {
-  struct mlkem::public_key<RANK1024> *pub =
+bcm_status bssl::BCM_mlkem1024_parse_public_key(
+    MLKEM1024_public_key *public_key, CBS *in) {
+  mlkem::public_key<RANK1024> *pub =
       mlkem::public_key_1024_from_external(public_key);
   if (!mlkem_parse_public_key(pub, in)) {
     return bcm_status::failure;
@@ -1351,9 +1655,9 @@ bcm_status BCM_mlkem1024_parse_public_key(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem768_marshal_private_key(
-    CBB *out, const struct BCM_mlkem768_private_key *private_key) {
-  const struct mlkem::private_key<RANK768> *const priv =
+bcm_status bssl::BCM_mlkem768_marshal_private_key(
+    CBB *out, const MLKEM768_private_key *private_key) {
+  const mlkem::private_key<RANK768> *const priv =
       mlkem::private_key_768_from_external(private_key);
   if (!mlkem_marshal_private_key(out, priv)) {
     return bcm_status::failure;
@@ -1361,9 +1665,9 @@ bcm_status BCM_mlkem768_marshal_private_key(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem1024_marshal_private_key(
-    CBB *out, const struct BCM_mlkem1024_private_key *private_key) {
-  const struct mlkem::private_key<RANK1024> *const priv =
+bcm_status bssl::BCM_mlkem1024_marshal_private_key(
+    CBB *out, const MLKEM1024_private_key *private_key) {
+  const mlkem::private_key<RANK1024> *const priv =
       mlkem::private_key_1024_from_external(private_key);
   if (!mlkem_marshal_private_key(out, priv)) {
     return bcm_status::failure;
@@ -1371,9 +1675,9 @@ bcm_status BCM_mlkem1024_marshal_private_key(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem768_parse_private_key(
-    struct BCM_mlkem768_private_key *out_private_key, CBS *in) {
-  struct mlkem::private_key<RANK768> *const priv =
+bcm_status bssl::BCM_mlkem768_parse_private_key(
+    MLKEM768_private_key *out_private_key, CBS *in) {
+  mlkem::private_key<RANK768> *const priv =
       mlkem::private_key_768_from_external(out_private_key);
   if (!mlkem_parse_private_key(priv, in)) {
     return bcm_status::failure;
@@ -1381,9 +1685,9 @@ bcm_status BCM_mlkem768_parse_private_key(
   return bcm_status::approved;
 }
 
-bcm_status BCM_mlkem1024_parse_private_key(
-    struct BCM_mlkem1024_private_key *out_private_key, CBS *in) {
-  struct mlkem::private_key<RANK1024> *const priv =
+bcm_status bssl::BCM_mlkem1024_parse_private_key(
+    MLKEM1024_private_key *out_private_key, CBS *in) {
+  mlkem::private_key<RANK1024> *const priv =
       mlkem::private_key_1024_from_external(out_private_key);
   if (!mlkem_parse_private_key(priv, in)) {
     return bcm_status::failure;
@@ -1391,7 +1695,7 @@ bcm_status BCM_mlkem1024_parse_private_key(
   return bcm_status::approved;
 }
 
-int boringssl_self_test_mlkem() {
+int bssl::boringssl_self_test_mlkem() {
   return mlkem::fips::keygen_self_test() && mlkem::fips::encap_self_test() &&
          mlkem::fips::decap_self_test();
 }