From 8e72c8549144fdcc2916ba65fb52c264cc95bedb Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Fri, 13 Dec 2024 06:32:01 +0000
Subject: [PATCH] CBMC: Add custom reordering loop into CBMC scope

`gen_matrix()` concludes with an entry-wise permutation of polynomial
coefficients, if required by the native backend.

Previously, the loop over the matrix coefficients was conditional on
a custom permutation being present, which would exclude it from the
scope of CBMC.

With this commit, the loop is unconditional, and a contract used for
the custom permutation. If no custom permutation is needed, it is
defined as a no-op which the compiler will automate away.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 cbmc/proofs/gen_matrix/Makefile |  6 +++---
 mlkem/indcpa.c                  | 13 +++++++++++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/cbmc/proofs/gen_matrix/Makefile b/cbmc/proofs/gen_matrix/Makefile
index eb3d76228..129bded75 100644
--- a/cbmc/proofs/gen_matrix/Makefile
+++ b/cbmc/proofs/gen_matrix/Makefile
@@ -13,13 +13,13 @@ DEFINES +=
 INCLUDES +=
 
 REMOVE_FUNCTION_BODY +=
-UNWINDSET += $(MLKEM_NAMESPACE)gen_matrix.0:4 $(MLKEM_NAMESPACE)gen_matrix.1:4 $(MLKEM_NAMESPACE)gen_matrix.2:4
+UNWINDSET += $(MLKEM_NAMESPACE)gen_matrix.0:4 $(MLKEM_NAMESPACE)gen_matrix.1:4 $(MLKEM_NAMESPACE)gen_matrix.2:4 $(MLKEM_NAMESPACE)gen_matrix.3:4 $(MLKEM_NAMESPACE)gen_matrix.4:4
 
 PROOF_SOURCES += $(PROOFDIR)/$(HARNESS_FILE).c
 PROJECT_SOURCES += $(SRCDIR)/mlkem/indcpa.c
 
 CHECK_FUNCTION_CONTRACTS=$(MLKEM_NAMESPACE)gen_matrix
-USE_FUNCTION_CONTRACTS=gen_matrix_entry gen_matrix_entry_x4
+USE_FUNCTION_CONTRACTS=gen_matrix_entry gen_matrix_entry_x4 poly_permute_bitrev_to_custom
 APPLY_LOOP_CONTRACTS=on
 USE_DYNAMIC_FRAMES=1
 
@@ -36,7 +36,7 @@ FUNCTION_NAME = $(MLKEM_NAMESPACE)gen_matrix
 # EXPENSIVE = true
 
 # This function is large enough to need...
-CBMC_OBJECT_BITS = 8
+CBMC_OBJECT_BITS = 10
 
 # 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/mlkem/indcpa.c b/mlkem/indcpa.c
index af4a5dc7f..bde832c18 100644
--- a/mlkem/indcpa.c
+++ b/mlkem/indcpa.c
@@ -250,6 +250,17 @@ __contract__(
   xof_release(&state);
 }
 
+#if !defined(MLKEM_USE_NATIVE_NTT_CUSTOM_ORDER)
+STATIC_INLINE_TESTABLE
+void poly_permute_bitrev_to_custom(poly *data)
+__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 gen_matrix. */
+  requires(memory_no_alias(data, sizeof(poly)))
+  requires(array_bound(data->coeffs, 0, MLKEM_N - 1, 0, MLKEM_Q - 1))
+  assigns(memory_slice(data, sizeof(poly)))
+  ensures(array_bound(data->coeffs, 0, MLKEM_N - 1, 0, MLKEM_Q - 1))) { ((void)data); }
+#endif /* MLKEM_USE_NATIVE_NTT_CUSTOM_ORDER */
 
 /* Not static for benchmarking */
 void gen_matrix(polyvec *a, const uint8_t seed[MLKEM_SYMBYTES], int transposed)
@@ -330,7 +341,6 @@ void gen_matrix(polyvec *a, const uint8_t seed[MLKEM_SYMBYTES], int transposed)
   cassert(i == MLKEM_K * MLKEM_K,
           "gen_matrix: failed to generate whole matrix");
 
-#if defined(MLKEM_USE_NATIVE_NTT_CUSTOM_ORDER)
   /*
    * The public matrix is generated in NTT domain. If the native backend
    * uses a custom order in NTT domain, permute A accordingly.
@@ -342,7 +352,6 @@ void gen_matrix(polyvec *a, const uint8_t seed[MLKEM_SYMBYTES], int transposed)
       poly_permute_bitrev_to_custom(&a[i].vec[j]);
     }
   }
-#endif /* MLKEM_USE_NATIVE_NTT_CUSTOM_ORDER */
 }
 
 /*************************************************