Gotta inline 'em all.

cryspen · Jun 18, 2024 · 6693cd9 · 6693cd9
1 parent 1045b22
commit 6693cd9
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 7 deletions.
diff --git a/libcrux-ml-dsa/src/arithmetic.rs b/libcrux-ml-dsa/src/arithmetic.rs
@@ -11,6 +11,7 @@ impl PolynomialRingElement {
         coefficients: [0i32; 256],
     };
 
+    #[inline(always)]
     pub(crate) fn add(&self, rhs: &Self) -> Self {
         let mut sum = Self::ZERO;
 
@@ -21,6 +22,7 @@ impl PolynomialRingElement {
         sum
     }
 
+    #[inline(always)]
     pub(crate) fn sub(&self, rhs: &Self) -> Self {
         let mut difference = Self::ZERO;
 
@@ -31,6 +33,7 @@ impl PolynomialRingElement {
         difference
     }
 
+    #[inline(always)]
     pub(crate) fn infinity_norm_exceeds(&self, value: i32) -> bool {
         if value > (FIELD_MODULUS - 1) / 8 {
             return true;
@@ -53,6 +56,7 @@ impl PolynomialRingElement {
     }
 }
 
+#[inline(always)]
 pub(crate) fn vector_infinity_norm_exceeds<const DIMENSION: usize>(
     vector: [PolynomialRingElement; DIMENSION],
     value: i32,
@@ -66,6 +70,7 @@ pub(crate) fn vector_infinity_norm_exceeds<const DIMENSION: usize>(
     false
 }
 
+#[inline(always)]
 pub(crate) fn get_n_least_significant_bits(n: u8, value: u64) -> u64 {
     value & ((1 << n) - 1)
 }
@@ -86,6 +91,7 @@ pub(crate) type FieldElementTimesMontgomeryR = i32;
 const MONTGOMERY_SHIFT: u8 = 32;
 const INVERSE_OF_MODULUS_MOD_MONTGOMERY_R: u64 = 58_728_449; // FIELD_MODULUS^{-1} mod 2^32
 
+#[inline(always)]
 pub(crate) fn montgomery_reduce(value: i64) -> MontgomeryFieldElement {
     let t = get_n_least_significant_bits(MONTGOMERY_SHIFT, value as u64)
         * INVERSE_OF_MODULUS_MOD_MONTGOMERY_R;
@@ -115,6 +121,7 @@ pub(crate) fn montgomery_multiply_fe_by_fer(
 //
 // We assume the input t is in the signed representative range and convert it
 // to the standard unsigned range.
+#[inline(always)]
 fn power2round(t: i32) -> (i32, i32) {
     debug_assert!(t > -FIELD_MODULUS && t < FIELD_MODULUS, "t is {}", t);
 
@@ -131,6 +138,8 @@ fn power2round(t: i32) -> (i32, i32) {
 
     (t0, t1)
 }
+
+#[inline(always)]
 pub(crate) fn power2round_vector<const DIMENSION: usize>(
     t: [PolynomialRingElement; DIMENSION],
 ) -> (
@@ -166,6 +175,7 @@ pub(crate) fn power2round_vector<const DIMENSION: usize>(
 // - α/2 ≤ r₀ < 0.
 //
 // Note that 0 ≤ r₁ < (q-1)/α.
+#[inline(always)]
 fn decompose<const GAMMA2: i32>(r: i32) -> (i32, i32) {
     debug_assert!(
         r > -FIELD_MODULUS && r < FIELD_MODULUS,
@@ -212,6 +222,8 @@ fn decompose<const GAMMA2: i32>(r: i32) -> (i32, i32) {
 
     (r0, r1)
 }
+
+#[inline(always)]
 pub(crate) fn decompose_vector<const DIMENSION: usize, const GAMMA2: i32>(
     t: [PolynomialRingElement; DIMENSION],
 ) -> (
@@ -233,10 +245,12 @@ pub(crate) fn decompose_vector<const DIMENSION: usize, const GAMMA2: i32>(
     (vector_low, vector_high)
 }
 
+#[inline(always)]
 fn make_hint<const GAMMA2: i32>(low: i32, high: i32) -> bool {
     (low > GAMMA2) || (low < -GAMMA2) || (low == -GAMMA2 && high != 0)
 }
 
+#[inline(always)]
 pub(crate) fn make_hint_vector<const DIMENSION: usize, const GAMMA2: i32>(
     low: [PolynomialRingElement; DIMENSION],
     high: [PolynomialRingElement; DIMENSION],

diff --git a/libcrux-ml-dsa/src/encoding/commitment.rs b/libcrux-ml-dsa/src/encoding/commitment.rs
@@ -39,6 +39,7 @@ fn serialize<const OUTPUT_SIZE: usize>(re: PolynomialRingElement) -> [u8; OUTPUT
     }
 }
 
+#[inline(always)]
 pub(crate) fn serialize_vector<
     const DIMENSION: usize,
     const RING_ELEMENT_SIZE: usize,

diff --git a/libcrux-ml-dsa/src/encoding/error.rs b/libcrux-ml-dsa/src/encoding/error.rs
@@ -1,6 +1,4 @@
-// ---------------------------------------------------------------------------
 // Functions for serializing and deserializing an error ring element.
-// ---------------------------------------------------------------------------
 
 use crate::{arithmetic::PolynomialRingElement, ntt::ntt};
 

diff --git a/libcrux-ml-dsa/src/hash_functions.rs b/libcrux-ml-dsa/src/hash_functions.rs
@@ -20,13 +20,15 @@ pub(crate) mod H_128 {
         state
     }
 
+    #[inline(always)]
     pub(crate) fn squeeze_first_five_blocks(state: &mut KeccakState1) -> [u8; FIVE_BLOCKS_SIZE] {
         let mut out = [0u8; FIVE_BLOCKS_SIZE];
         incremental::shake128_squeeze_first_five_blocks(state, &mut out);
 
         out
     }
 
+    #[inline(always)]
     pub(crate) fn squeeze_next_block(state: &mut KeccakState1) -> [u8; BLOCK_SIZE] {
         let mut out = [0u8; BLOCK_SIZE];
         incremental::shake128_squeeze_next_block(state, &mut out);

diff --git a/libcrux-ml-dsa/src/matrix.rs b/libcrux-ml-dsa/src/matrix.rs
@@ -24,8 +24,8 @@ pub(crate) fn expand_to_A<const ROWS_IN_A: usize, const COLUMNS_IN_A: usize>(
 }
 
 /// Compute InvertNTT(Â ◦ ŝ₁) + s₂
-#[inline(always)]
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn compute_As1_plus_s2<const ROWS_IN_A: usize, const COLUMNS_IN_A: usize>(
     A_as_ntt: &[[PolynomialRingElement; COLUMNS_IN_A]; ROWS_IN_A],
     s1: &[PolynomialRingElement; COLUMNS_IN_A],
@@ -47,8 +47,8 @@ pub(crate) fn compute_As1_plus_s2<const ROWS_IN_A: usize, const COLUMNS_IN_A: us
 }
 
 /// Compute InvertNTT(Â ◦ ŷ)
-#[inline(always)]
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn compute_A_times_mask<const ROWS_IN_A: usize, const COLUMNS_IN_A: usize>(
     A_as_ntt: &[[PolynomialRingElement; COLUMNS_IN_A]; ROWS_IN_A],
     mask: &[PolynomialRingElement; COLUMNS_IN_A],
@@ -67,8 +67,8 @@ pub(crate) fn compute_A_times_mask<const ROWS_IN_A: usize, const COLUMNS_IN_A: u
     result
 }
 
-#[inline(always)]
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn vector_times_ring_element<const DIMENSION: usize>(
     vector: &[PolynomialRingElement; DIMENSION],
     ring_element: &PolynomialRingElement,
@@ -82,8 +82,8 @@ pub(crate) fn vector_times_ring_element<const DIMENSION: usize>(
     result
 }
 
-#[inline(always)]
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn add_vectors<const DIMENSION: usize>(
     lhs: &[PolynomialRingElement; DIMENSION],
     rhs: &[PolynomialRingElement; DIMENSION],
@@ -97,8 +97,8 @@ pub(crate) fn add_vectors<const DIMENSION: usize>(
     result
 }
 
-#[inline(always)]
 #[allow(non_snake_case)]
+#[inline(always)]
 pub(crate) fn subtract_vectors<const DIMENSION: usize>(
     lhs: &[PolynomialRingElement; DIMENSION],
     rhs: &[PolynomialRingElement; DIMENSION],

diff --git a/libcrux-ml-dsa/src/ntt.rs b/libcrux-ml-dsa/src/ntt.rs
@@ -101,6 +101,7 @@ fn invert_ntt_at_layer(
 
     re
 }
+
 #[inline(always)]
 pub(crate) fn invert_ntt_montgomery(mut re: PolynomialRingElement) -> PolynomialRingElement {
     let mut zeta_i = COEFFICIENTS_IN_RING_ELEMENT;
@@ -125,6 +126,7 @@ pub(crate) fn invert_ntt_montgomery(mut re: PolynomialRingElement) -> Polynomial
     re
 }
 
+#[inline(always)]
 pub(crate) fn ntt_multiply_montgomery(
     lhs: &PolynomialRingElement,
     rhs: &PolynomialRingElement,

diff --git a/libcrux-ml-dsa/src/sample.rs b/libcrux-ml-dsa/src/sample.rs
@@ -5,6 +5,7 @@ use crate::{
     hash_functions::{H, H_128},
 };
 
+#[inline(always)]
 fn rejection_sample_less_than_field_modulus(
     randomness: &[u8],
     sampled: &mut usize,
@@ -33,6 +34,8 @@ fn rejection_sample_less_than_field_modulus(
 
     done
 }
+
+#[inline(always)]
 pub(crate) fn sample_ring_element_uniform(seed: [u8; 34]) -> PolynomialRingElement {
     let mut state = H_128::new(seed);
     let randomness = H_128::squeeze_first_five_blocks(&mut state);
@@ -50,6 +53,7 @@ pub(crate) fn sample_ring_element_uniform(seed: [u8; 34]) -> PolynomialRingEleme
     out
 }
 
+#[inline(always)]
 fn rejection_sample_less_than_eta_equals_2(
     randomness: &[u8],
     sampled: &mut usize,
@@ -90,6 +94,8 @@ fn rejection_sample_less_than_eta_equals_2(
 
     done
 }
+
+#[inline(always)]
 fn rejection_sample_less_than_eta_equals_4(
     randomness: &[u8],
     sampled: &mut usize,
@@ -121,6 +127,7 @@ fn rejection_sample_less_than_eta_equals_4(
     done
 }
 
+#[inline(always)]
 pub(crate) fn rejection_sample_less_than_eta<const ETA: usize>(
     randomness: &[u8],
     sampled: &mut usize,
@@ -134,6 +141,7 @@ pub(crate) fn rejection_sample_less_than_eta<const ETA: usize>(
 }
 
 #[allow(non_snake_case)]
+#[inline(always)]
 fn sample_error_ring_element<const ETA: usize>(seed: [u8; 66]) -> PolynomialRingElement {
     // TODO: Use incremental API to squeeze one block at a time.
     let randomness = H::<272>(&seed);
@@ -150,6 +158,7 @@ fn sample_error_ring_element<const ETA: usize>(seed: [u8; 66]) -> PolynomialRing
 
     out
 }
+
 #[inline(always)]
 pub(crate) fn sample_error_vector<const DIMENSION: usize, const ETA: usize>(
     mut seed: [u8; 66],
@@ -194,6 +203,7 @@ pub(crate) fn sample_mask_vector<const DIMENSION: usize, const GAMMA1_EXPONENT:
     error
 }
 
+#[inline(always)]
 pub(crate) fn sample_challenge_ring_element<const NUMBER_OF_ONES: usize>(
     seed: [u8; 32],
 ) -> PolynomialRingElement {