Simplify gen_matrix() and poly_rej_uniform_x4()

1. Simplify poly_rej_uniform_x4() API to explicitly
   take 4 objects to write, rather than a slice
   of a larger array.
2. Simplify gen_matrix() in light of that, and also
   remove redudant transposition of the final poly
   sampling when MLKEM_K=3

Signed-off-by: Rod Chapman <rodchap@amazon.com>

Author: rod-chapman

mlkem/src/indcpa.c

@@ -221,10 +221,9 @@ void mlk_gen_matrix(mlk_polymat a, const uint8_t seed[MLKEM_SYMBYTES],
   /* Sample 4 matrix entries a time. */
   for (i = 0; i < (MLKEM_K * MLKEM_K / 4) * 4; i += 4)
   {
-    uint8_t x, y;
-
     for (j = 0; j < 4; j++)
     {
+      uint8_t x, y;
       x = (i + j) / MLKEM_K;
       y = (i + j) % MLKEM_K;
       if (transposed)
@@ -239,31 +238,25 @@ void mlk_gen_matrix(mlk_polymat a, const uint8_t seed[MLKEM_SYMBYTES],
       }
     }
 
-    /*
-     * This call writes across mlk_polyvec boundaries for K=2 and K=3.
-     * This is intentional and safe.
-     */
-    mlk_poly_rej_uniform_x4(&a[i], seed_ext);
+    mlk_poly_rej_uniform_x4(&a[i], &a[i + 1], &a[i + 2], &a[i + 3], seed_ext);
   }
 
-  /* For MLKEM_K == 3, sample the last entry individually. */
-  if (i < MLKEM_K * MLKEM_K)
-  {
-    uint8_t x, y;
-    x = i / MLKEM_K;
-    y = i % MLKEM_K;
-
-    if (transposed)
-    {
-      seed_ext[0][MLKEM_SYMBYTES + 0] = x;
-      seed_ext[0][MLKEM_SYMBYTES + 1] = y;
-    }
-    else
-    {
-      seed_ext[0][MLKEM_SYMBYTES + 0] = y;
-      seed_ext[0][MLKEM_SYMBYTES + 1] = x;
-    }
+  /* A ==> B is written as !A || B here so can be subject */
+  /* to both static and dynamic verification              */
+  mlk_assert((MLKEM_K != 2) || (i == 4));
+  mlk_assert((MLKEM_K != 3) || (i == 8));
+  mlk_assert((MLKEM_K != 4) || (i == 16));
 
+  /* For MLKEM_K == 2 or 4, we're done.                             */
+  /* For MLKEM_K == 3, we have one more poly to sample with i == 8. */
+  if (MLKEM_K == 3)
+  {
+    mlk_assert(i == 8 && i / MLKEM_K == 2 && i % MLKEM_K == 2);
+    /* The final poly is on the leading diagonal of the matrix */
+    /* at coordinates [2,2], so the "transpose" flag here is   */
+    /* irrelevant.                                             */
+    seed_ext[0][MLKEM_SYMBYTES + 0] = 2;
+    seed_ext[0][MLKEM_SYMBYTES + 1] = 2;
     mlk_poly_rej_uniform(&a[i], seed_ext[0]);
     i++;
   }

mlkem/src/sampling.c

@@ -146,7 +146,8 @@ __contract__(
  *            - x4-batched version of `rej_uniform()` from the
  *              reference implementation, leveraging x4-batched Keccak-f1600. */
 MLK_INTERNAL_API
-void mlk_poly_rej_uniform_x4(mlk_poly *vec,
+void mlk_poly_rej_uniform_x4(mlk_poly *vec0, mlk_poly *vec1, mlk_poly *vec2,
+                             mlk_poly *vec3,
                              uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)])
 {
   /* Temporary buffers for XOF output before rejection sampling */
@@ -167,10 +168,10 @@ void mlk_poly_rej_uniform_x4(mlk_poly *vec,
    */
   mlk_xof_x4_squeezeblocks(buf, MLKEM_GEN_MATRIX_NBLOCKS, &statex);
   buflen = MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE;
-  ctr[0] = mlk_rej_uniform(vec[0].coeffs, MLKEM_N, 0, buf[0], buflen);
-  ctr[1] = mlk_rej_uniform(vec[1].coeffs, MLKEM_N, 0, buf[1], buflen);
-  ctr[2] = mlk_rej_uniform(vec[2].coeffs, MLKEM_N, 0, buf[2], buflen);
-  ctr[3] = mlk_rej_uniform(vec[3].coeffs, MLKEM_N, 0, buf[3], buflen);
+  ctr[0] = mlk_rej_uniform(vec0->coeffs, MLKEM_N, 0, buf[0], buflen);
+  ctr[1] = mlk_rej_uniform(vec1->coeffs, MLKEM_N, 0, buf[1], buflen);
+  ctr[2] = mlk_rej_uniform(vec2->coeffs, MLKEM_N, 0, buf[2], buflen);
+  ctr[3] = mlk_rej_uniform(vec3->coeffs, MLKEM_N, 0, buf[3], buflen);
 
   /*
    * So long as not all matrix entries have been generated, squeeze
@@ -180,20 +181,27 @@ void mlk_poly_rej_uniform_x4(mlk_poly *vec,
   while (ctr[0] < MLKEM_N || ctr[1] < MLKEM_N || ctr[2] < MLKEM_N ||
          ctr[3] < MLKEM_N)
   __loop__(
-    assigns(ctr, statex, memory_slice(vec, sizeof(mlk_poly) * 4), object_whole(buf[0]),
-       object_whole(buf[1]), object_whole(buf[2]), object_whole(buf[3]))
+    assigns(ctr, statex,
+            memory_slice(vec0, sizeof(mlk_poly)),
+            memory_slice(vec1, sizeof(mlk_poly)),
+            memory_slice(vec2, sizeof(mlk_poly)),
+            memory_slice(vec3, sizeof(mlk_poly)),
+            object_whole(buf[0]),
+            object_whole(buf[1]),
+            object_whole(buf[2]),
+            object_whole(buf[3]))
     invariant(ctr[0] <= MLKEM_N && ctr[1] <= MLKEM_N)
     invariant(ctr[2] <= MLKEM_N && ctr[3] <= MLKEM_N)
-    invariant(array_bound(vec[0].coeffs, 0, ctr[0], 0, MLKEM_Q))
-    invariant(array_bound(vec[1].coeffs, 0, ctr[1], 0, MLKEM_Q))
-    invariant(array_bound(vec[2].coeffs, 0, ctr[2], 0, MLKEM_Q))
-    invariant(array_bound(vec[3].coeffs, 0, ctr[3], 0, MLKEM_Q)))
+    invariant(array_bound(vec0->coeffs, 0, ctr[0], 0, MLKEM_Q))
+    invariant(array_bound(vec1->coeffs, 0, ctr[1], 0, MLKEM_Q))
+    invariant(array_bound(vec2->coeffs, 0, ctr[2], 0, MLKEM_Q))
+    invariant(array_bound(vec3->coeffs, 0, ctr[3], 0, MLKEM_Q)))
   {
     mlk_xof_x4_squeezeblocks(buf, 1, &statex);
-    ctr[0] = mlk_rej_uniform(vec[0].coeffs, MLKEM_N, ctr[0], buf[0], buflen);
-    ctr[1] = mlk_rej_uniform(vec[1].coeffs, MLKEM_N, ctr[1], buf[1], buflen);
-    ctr[2] = mlk_rej_uniform(vec[2].coeffs, MLKEM_N, ctr[2], buf[2], buflen);
-    ctr[3] = mlk_rej_uniform(vec[3].coeffs, MLKEM_N, ctr[3], buf[3], buflen);
+    ctr[0] = mlk_rej_uniform(vec0->coeffs, MLKEM_N, ctr[0], buf[0], buflen);
+    ctr[1] = mlk_rej_uniform(vec1->coeffs, MLKEM_N, ctr[1], buf[1], buflen);
+    ctr[2] = mlk_rej_uniform(vec2->coeffs, MLKEM_N, ctr[2], buf[2], buflen);
+    ctr[3] = mlk_rej_uniform(vec3->coeffs, MLKEM_N, ctr[3], buf[3], buflen);
   }
 
   mlk_xof_x4_release(&statex);

mlkem/src/sampling.h

@@ -65,8 +65,8 @@ void mlk_poly_cbd3(mlk_poly *r, const uint8_t buf[3 * MLKEM_N / 4]);
  * Description: Generate four polynomials using rejection sampling
  *              on (pseudo-)uniformly random bytes sampled from a seed.
  *
- * Arguments:   - mlk_poly *vec:
- *                Pointer to an array of 4 polynomials to be sampled.
+ * Arguments:   - mlk_poly *vec0, *vec1, *vec2, *vec3:
+ *                Pointers to 4 polynomials to be sampled.
  *              - uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)]:
  *                Pointer consecutive array of seed buffers of size
  *                MLKEM_SYMBYTES + 2 each, plus padding for alignment.
@@ -75,16 +75,23 @@ void mlk_poly_cbd3(mlk_poly *r, const uint8_t buf[3 * MLKEM_N / 4]);
  *
  **************************************************/
 MLK_INTERNAL_API
-void mlk_poly_rej_uniform_x4(mlk_poly *vec,
+void mlk_poly_rej_uniform_x4(mlk_poly *vec0, mlk_poly *vec1, mlk_poly *vec2,
+                             mlk_poly *vec3,
                              uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)])
 __contract__(
-  requires(memory_no_alias(vec, sizeof(mlk_poly) * 4))
+  requires(memory_no_alias(vec0, sizeof(mlk_poly)))
+  requires(memory_no_alias(vec1, sizeof(mlk_poly)))
+  requires(memory_no_alias(vec2, sizeof(mlk_poly)))
+  requires(memory_no_alias(vec3, sizeof(mlk_poly)))
   requires(memory_no_alias(seed, 4 * MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)))
-  assigns(memory_slice(vec, sizeof(mlk_poly) * 4))
-  ensures(array_bound(vec[0].coeffs, 0, MLKEM_N, 0, MLKEM_Q))
-  ensures(array_bound(vec[1].coeffs, 0, MLKEM_N, 0, MLKEM_Q))
-  ensures(array_bound(vec[2].coeffs, 0, MLKEM_N, 0, MLKEM_Q))
-  ensures(array_bound(vec[3].coeffs, 0, MLKEM_N, 0, MLKEM_Q)));
+  assigns(memory_slice(vec0, sizeof(mlk_poly)))
+  assigns(memory_slice(vec1, sizeof(mlk_poly)))
+  assigns(memory_slice(vec2, sizeof(mlk_poly)))
+  assigns(memory_slice(vec3, sizeof(mlk_poly)))
+  ensures(array_bound(vec0->coeffs, 0, MLKEM_N, 0, MLKEM_Q))
+  ensures(array_bound(vec1->coeffs, 0, MLKEM_N, 0, MLKEM_Q))
+  ensures(array_bound(vec2->coeffs, 0, MLKEM_N, 0, MLKEM_Q))
+  ensures(array_bound(vec3->coeffs, 0, MLKEM_N, 0, MLKEM_Q)));
 
 #define mlk_poly_rej_uniform MLK_NAMESPACE(poly_rej_uniform)
 /*************************************************

proofs/cbmc/poly_rej_uniform_x4/poly_rej_uniform_x4_harness.c

@@ -7,7 +7,10 @@
 
 void harness(void)
 {
-  mlk_poly *out;
+  mlk_poly *out0;
+  mlk_poly *out1;
+  mlk_poly *out2;
+  mlk_poly *out3;
   uint8_t(*seed)[MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)];
-  mlk_poly_rej_uniform_x4(out, seed);
+  mlk_poly_rej_uniform_x4(out0, out1, out2, out3, seed);
 }