Refactor mlk_polymat_permute_bitrev_to_custom

willieyz · willieyz · commit 4f7ba6367b5a · 2026-01-06T18:26:42.000+08:00
This commit refactors mlk_polymat_permute_bitrev_to_custom to not require the helper function mlk_polyvec_permute_bitrev_to_custom. The function was only needed due CBMC limitations. Also updates CBMC to trial new build (CBMC PR[#8705](https://github.com/pq-code-package/mlkem-native/issues/8705)) that solves performance problem for this proof. Simplify polymat_permite_bitrev_to_customer() (native version) The code structure now mimics the data structure to make proof tractable. Also updates Makefile for this proof in line with the similar function in mldsa-native. Signed-off-by: Matthias J. Kannwischer <matthias@kannwischer.eu> Signed-off-by: Rod Chapman <rodchap@amazon.com> Signed-off-by: willieyz <willie.zhao@chelpis.com>
diff --git a/mlkem/src/indcpa.c b/mlkem/src/indcpa.c
@@ -39,8 +39,6 @@
 #define mlk_pack_ciphertext MLK_ADD_PARAM_SET(mlk_pack_ciphertext)
 #define mlk_unpack_ciphertext MLK_ADD_PARAM_SET(mlk_unpack_ciphertext)
 #define mlk_matvec_mul MLK_ADD_PARAM_SET(mlk_matvec_mul)
-#define mlk_polyvec_permute_bitrev_to_custom \
-  MLK_ADD_PARAM_SET(mlk_polyvec_permute_bitrev_to_custom)
 #define mlk_polymat_permute_bitrev_to_custom \
   MLK_ADD_PARAM_SET(mlk_polymat_permute_bitrev_to_custom)
 /* End of parameter set namespacing */
@@ -181,40 +179,6 @@ static void mlk_unpack_ciphertext(mlk_polyvec *b, mlk_poly *v,
   mlk_poly_decompress_dv(v, c + MLKEM_POLYVECCOMPRESSEDBYTES_DU);
 }
 
-/* Helper function to ensure that the polynomial entries in the output
- * of gen_matrix use the standard (bitreversed) ordering of coefficients.
- * No-op unless a native backend with a custom ordering is used.
- *
- * We don't inline this into gen_matrix to avoid having to split the CBMC
- * proof for gen_matrix based on MLK_USE_NATIVE_NTT_CUSTOM_ORDER. */
-static void mlk_polyvec_permute_bitrev_to_custom(mlk_polyvec *v)
-__contract__(
-  /* We don't specify that this should be a permutation, but only
-   * that it does not change the bound established at the end of mlk_gen_matrix. */
-  requires(memory_no_alias(v, sizeof(mlk_polyvec)))
-  requires(forall(x, 0, MLKEM_K,
-    array_bound(v->vec[x].coeffs, 0, MLKEM_N, 0, MLKEM_Q)))
-  assigns(memory_slice(v, sizeof(mlk_polyvec)))
-  ensures(forall(x, 0, MLKEM_K,
-    array_bound(v->vec[x].coeffs, 0, MLKEM_N, 0, MLKEM_Q))))
-{
-#if defined(MLK_USE_NATIVE_NTT_CUSTOM_ORDER)
-  unsigned i;
-  for (i = 0; i < MLKEM_K; i++)
-  __loop__(
-     assigns(i, memory_slice(v, sizeof(mlk_polyvec)))
-     invariant(i <= MLKEM_K)
-     invariant(forall(x, 0, MLKEM_K,
-       array_bound(v->vec[x].coeffs, 0, MLKEM_N, 0, MLKEM_Q))))
-  {
-    mlk_poly_permute_bitrev_to_custom(v->vec[i].coeffs);
-  }
-#else  /* MLK_USE_NATIVE_NTT_CUSTOM_ORDER */
-  /* Nothing to do */
-  (void)v;
-#endif /* !MLK_USE_NATIVE_NTT_CUSTOM_ORDER */
-}
-
 static void mlk_polymat_permute_bitrev_to_custom(mlk_polymat *a)
 __contract__(
   /* We don't specify that this should be a permutation, but only
@@ -226,16 +190,31 @@ __contract__(
   ensures(forall(x, 0, MLKEM_K, forall(y, 0, MLKEM_K,
     array_bound(a->vec[x].vec[y].coeffs, 0, MLKEM_N, 0, MLKEM_Q)))))
 {
-  unsigned i;
+#if defined(MLK_USE_NATIVE_NTT_CUSTOM_ORDER)
+  unsigned i, j;
   for (i = 0; i < MLKEM_K; i++)
   __loop__(
-     assigns(i, memory_slice(a, sizeof(mlk_polymat)))
-     invariant(i <= MLKEM_K)
-     invariant(forall(x, 0, MLKEM_K, forall(y, 0, MLKEM_K,
-       array_bound(a->vec[x].vec[y].coeffs, 0, MLKEM_N, 0, MLKEM_Q)))))
+    assigns(i, j, memory_slice(a, sizeof(mlk_polymat)))
+    invariant(i <= MLKEM_K)
+    invariant(forall(x2, 0, MLKEM_K, forall(y2, 0, MLKEM_K,
+      array_bound(a->vec[x2].vec[y2].coeffs, 0, MLKEM_N, 0, MLKEM_Q)))))
   {
-    mlk_polyvec_permute_bitrev_to_custom(&a->vec[i]);
+    for (j = 0; j < MLKEM_K; j++)
+    __loop__(
+      assigns(j, memory_slice(a, sizeof(mlk_polymat)))
+      invariant(i <= MLKEM_K)
+      invariant(j <= MLKEM_K)
+      invariant(forall(x3, 0, MLKEM_K, forall(y3, 0, MLKEM_K,
+        array_bound(a->vec[x3].vec[y3].coeffs, 0, MLKEM_N, 0, MLKEM_Q))))
+    )
+    {
+      mlk_poly_permute_bitrev_to_custom(a->vec[i].vec[j].coeffs);
+    }
   }
+#else  /* MLK_USE_NATIVE_NTT_CUSTOM_ORDER */
+  /* Nothing to do */
+  (void)a;
+#endif /* !MLK_USE_NATIVE_NTT_CUSTOM_ORDER */
 }
 
 /* Reference: `gen_matrix()` in the reference implementation @[REF].
@@ -617,5 +596,4 @@ int mlk_indcpa_dec(uint8_t m[MLKEM_INDCPA_MSGBYTES],
 #undef mlk_pack_ciphertext
 #undef mlk_unpack_ciphertext
 #undef mlk_matvec_mul
-#undef mlk_polyvec_permute_bitrev_to_custom
 #undef mlk_polymat_permute_bitrev_to_custom
diff --git a/proofs/cbmc/polymat_permute_bitrev_to_custom/Makefile b/proofs/cbmc/polymat_permute_bitrev_to_custom/Makefile
@@ -19,7 +19,7 @@ PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mlkem/src/indcpa.c
 
 CHECK_FUNCTION_CONTRACTS=mlk_polymat_permute_bitrev_to_custom
-USE_FUNCTION_CONTRACTS=mlk_polyvec_permute_bitrev_to_custom
+USE_FUNCTION_CONTRACTS=
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 
diff --git a/proofs/cbmc/polymat_permute_bitrev_to_custom_native/Makefile b/proofs/cbmc/polymat_permute_bitrev_to_custom_native/Makefile
@@ -4,11 +4,11 @@
 include ../Makefile_params.common
 
 HARNESS_ENTRY = harness
-HARNESS_FILE = polyvec_permute_bitrev_to_custom_native_harness
+HARNESS_FILE = polymat_permute_bitrev_to_custom_native_harness
 
 # This should be a unique identifier for this proof, and will appear on the
 # Litani dashboard. It can be human-readable and contain spaces if you wish.
-PROOF_UID = mlk_polyvec_permute_bitrev_to_custom_native
+PROOF_UID = mlk_polymat_permute_bitrev_to_custom_native
 
 DEFINES += -DMLK_CONFIG_USE_NATIVE_BACKEND_ARITH -DMLK_CONFIG_ARITH_BACKEND_FILE="\"dummy_backend.h\""
 INCLUDES +=
@@ -19,23 +19,23 @@ REMOVE_FUNCTION_BODY +=
 PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mlkem/src/indcpa.c
 
-CHECK_FUNCTION_CONTRACTS=mlk_polyvec_permute_bitrev_to_custom
-USE_FUNCTION_CONTRACTS= mlk_poly_permute_bitrev_to_custom
+CHECK_FUNCTION_CONTRACTS=mlk_polymat_permute_bitrev_to_custom
+USE_FUNCTION_CONTRACTS=mlk_poly_permute_bitrev_to_custom
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 
 # Disable any setting of EXTERNAL_SAT_SOLVER, and choose SMT backend instead
 EXTERNAL_SAT_SOLVER=
-CBMCFLAGS=--smt2
+CBMCFLAGS=--smt2 --slice-formula --no-array-field-sensitivity
 
-FUNCTION_NAME = mlk_polyvec_permute_bitrev_to_custom_native
+FUNCTION_NAME = mlk_polymat_permute_bitrev_to_custom_native
 
 # If this proof is found to consume huge amounts of RAM, you can set the
 # EXPENSIVE variable. With new enough versions of the proof tools, this will
 # restrict the number of EXPENSIVE CBMC jobs running at once. See the
 # documentation in Makefile.common under the "Job Pools" heading for details.
 # EXPENSIVE = true
-CBMC_OBJECT_BITS = 10
+CBMC_OBJECT_BITS = 12
 
 # If you require access to a file-local ("static") function or object to conduct
 # your proof, set the following (and do not include the original source file
diff --git a/proofs/cbmc/polymat_permute_bitrev_to_custom_native/polymat_permute_bitrev_to_custom_native_harness.c b/proofs/cbmc/polymat_permute_bitrev_to_custom_native/polymat_permute_bitrev_to_custom_native_harness.c
@@ -5,15 +5,10 @@
 #include <stdint.h>
 #include "poly_k.h"
 
-void mlk_polyvec_permute_bitrev_to_custom(mlk_polyvec *v);
+void mlk_polymat_permute_bitrev_to_custom(mlk_polymat *a);
 
 void harness(void)
 {
-  {
-    /* Dummy use of `free` to work around CBMC issue #8814. */
-    free(NULL);
-  }
-
-  mlk_polyvec *v;
-  mlk_polyvec_permute_bitrev_to_custom(v);
+  mlk_polymat *a;
+  mlk_polymat_permute_bitrev_to_custom(a);
 }
diff --git a/proofs/cbmc/polyvec_permute_bitrev_to_custom/Makefile b/proofs/cbmc/polyvec_permute_bitrev_to_custom/Makefile
diff --git a/proofs/cbmc/polyvec_permute_bitrev_to_custom_native/polyvec_permute_bitrev_to_custom_native_harness.c b/proofs/cbmc/polyvec_permute_bitrev_to_custom_native/polyvec_permute_bitrev_to_custom_native_harness.c
