Improve key generation stack usage

libcrux/libcrux-ml-kem/src/hash_functions.rs

@@ -14,59 +14,7 @@ pub(crate) const BLOCK_SIZE: usize = 168;
 /// The size of 3 SHA3 blocks.
 pub(crate) const THREE_BLOCKS: usize = BLOCK_SIZE * 3;
 
-/// Abstraction for the hashing, to pick the fastest version depending on the
-/// platform features available.
-///
-/// In libcrux-iot we currently support a portable instantiations of
-/// this trait right now, whereas mainline libcrux supports additional
-/// SIMD platform.
-#[hax_lib::attributes]
-pub(crate) trait Hash {
-    /// G aka SHA3 512
-    #[requires(true)]
-    #[ensures(|result|
-        fstar!(r#"$result == Spec.Utils.v_G $input"#))
-    ]
-    fn G(input: &[u8], output: &mut [u8]);
-
-    /// H aka SHA3 256
-    #[requires(true)]
-    #[ensures(|result|
-        fstar!(r#"$result == Spec.Utils.v_H $input"#))
-    ]
-    fn H(input: &[u8], output: &mut [u8]);
-
-    /// PRF aka SHAKE256
-    #[requires(fstar!(r#"v $LEN < pow2 32"#))]
-    #[ensures(|result|
-        // We need to repeat the pre-condition here because of https://github.com/hacspec/hax/issues/784
-        fstar!(r#"v $LEN < pow2 32 ==> $result == Spec.Utils.v_PRF $LEN $input"#))
-    ]
-    fn PRF<const LEN: usize>(input: &[u8], out: &mut [u8]);
-
-    /// PRFxN aka N SHAKE256
-    #[requires(fstar!(r#"v $LEN < pow2 32 /\ (v $K == 2 \/ v $K == 3 \/ v $K == 4)"#))]
-    #[ensures(|result|
-        // We need to repeat the pre-condition here because of https://github.com/hacspec/hax/issues/784
-        fstar!(r#"(v $LEN < pow2 32 /\ (v $K == 2 \/ v $K == 3 \/ v $K == 4)) ==>
-            $result == Spec.Utils.v_PRFxN $K $LEN $input"#))
-    ]
-    fn PRFxN(input: &[[u8; 33]], outputs: &mut [u8], out_len: usize);
-
-    /// Create a SHAKE128 state and absorb the input.
-    #[requires(true)]
-    fn shake128_init_absorb_final(input: &[[u8; 34]]) -> Self;
-
-    /// Squeeze 3 blocks out of the SHAKE128 state.
-    #[requires(true)]
-    fn shake128_squeeze_first_three_blocks(&mut self, output: &mut [[u8; THREE_BLOCKS]]);
-
-    /// Squeeze 1 block out of the SHAKE128 state.
-    #[requires(true)]
-    fn shake128_squeeze_next_block(&mut self, output: &mut [[u8; BLOCK_SIZE]]);
-}
-
-/// A portable implementation of [`Hash`]
+/// Portable hash function implementations as needed for ML-KEM.
 pub(crate) mod portable {
     use super::*;
     use libcrux_sha3::portable::{self, incremental, KeccakState};
@@ -77,7 +25,7 @@ pub(crate) mod portable {
     /// All other functions don't actually use any members.
     #[cfg_attr(hax, hax_lib::opaque)]
     pub(crate) struct PortableHash {
-        shake128_state: [KeccakState; 4],
+        shake128_state: KeccakState,
     }
 
     #[hax_lib::ensures(|result|
@@ -105,65 +53,38 @@ pub(crate) mod portable {
         portable::shake256(out, input);
     }
 
-    #[hax_lib::requires(fstar!(r#"v $LEN < pow2 32 /\ (v $K == 2 \/ v $K == 3 \/ v $K == 4)"#))]
-    #[hax_lib::ensures(|result|
-        fstar!(r#"$result == Spec.Utils.v_PRFxN $K $LEN $input"#))
-    ]
-    #[inline(always)]
-    fn PRFxN(input: &[[u8; 33]], outputs: &mut [u8], out_len: usize) {
-        for i in 0..input.len() {
-            portable::shake256(&mut outputs[i * out_len..(i + 1) * out_len], &input[i]);
-        }
-    }
-
     #[inline(always)]
-    fn shake128_init_absorb_final(input: &[[u8; 34]]) -> PortableHash {
-        // debug_assert!(input.len() == 2 || input.len() == 3 || input.len() == 4);
-
-        let mut shake128_state = [incremental::shake128_init(); 4];
-        for i in 0..input.len() {
-            incremental::shake128_absorb_final(&mut shake128_state[i], &input[i]);
-        }
+    fn shake128_init_absorb_final(input: &[u8; 34]) -> PortableHash {
+        let mut shake128_state = incremental::shake128_init();
+        incremental::shake128_absorb_final(&mut shake128_state, &input[..]);
         PortableHash { shake128_state }
     }
 
     #[inline(always)]
-    fn shake128_squeeze_first_three_blocks(
-        st: &mut PortableHash,
-        outputs: &mut [[u8; THREE_BLOCKS]],
-    ) {
-        // debug_assert!(outputs.len() == 2 || outputs.len() == 3 || outputs.len() == 4);
-
-        for i in 0..outputs.len() {
-            incremental::shake128_squeeze_first_three_blocks(
-                &mut st.shake128_state[i],
-                &mut outputs[i],
-            );
-        }
+    fn shake128_squeeze_first_three_blocks(st: &mut PortableHash, output: &mut [u8; THREE_BLOCKS]) {
+        incremental::shake128_squeeze_first_three_blocks(&mut st.shake128_state, &mut output[..]);
     }
 
     #[inline(always)]
-    fn shake128_squeeze_next_block(st: &mut PortableHash, outputs: &mut [[u8; BLOCK_SIZE]]) {
-        for i in 0..outputs.len() {
-            incremental::shake128_squeeze_next_block(&mut st.shake128_state[i], &mut outputs[i]);
-        }
+    fn shake128_squeeze_next_block(st: &mut PortableHash, output: &mut [u8; BLOCK_SIZE]) {
+        incremental::shake128_squeeze_next_block(&mut st.shake128_state, &mut output[..]);
     }
 
     #[hax_lib::attributes]
-    impl Hash for PortableHash {
+    impl PortableHash {
         #[ensures(|out|
             fstar!(r#"$out == Spec.Utils.v_G $input"#))
         ]
         #[inline(always)]
-        fn G(input: &[u8], output: &mut [u8]) {
+        pub(crate) fn G(input: &[u8], output: &mut [u8]) {
             G(input, output)
         }
 
         #[ensures(|out|
             fstar!(r#"$out == Spec.Utils.v_H $input"#))
         ]
         #[inline(always)]
-        fn H(input: &[u8], output: &mut [u8]) {
+        pub(crate) fn H(input: &[u8], output: &mut [u8]) {
             H(input, output)
         }
 
@@ -173,32 +94,25 @@ pub(crate) mod portable {
             fstar!(r#"v $LEN < pow2 32 ==> $out == Spec.Utils.v_PRF $LEN $input"#))
         ]
         #[inline(always)]
-        fn PRF<const LEN: usize>(input: &[u8], out: &mut [u8]) {
+        pub(crate) fn PRF<const LEN: usize>(input: &[u8], out: &mut [u8]) {
             PRF::<LEN>(input, out)
         }
 
-        #[requires(fstar!(r#"v $LEN < pow2 32 /\ (v $K == 2 \/ v $K == 3 \/ v $K == 4)"#))]
-        #[ensures(|out|
-            fstar!(r#"(v $LEN < pow2 32 /\ (v $K == 2 \/ v $K == 3 \/ v $K == 4)) ==>
-                $out == Spec.Utils.v_PRFxN $K $LEN $input"#))
-        ]
-        #[inline(always)]
-        fn PRFxN(input: &[[u8; 33]], outputs: &mut [u8], out_len: usize) {
-            PRFxN(input, outputs, out_len)
-        }
-
         #[inline(always)]
-        fn shake128_init_absorb_final(input: &[[u8; 34]]) -> Self {
+        pub(crate) fn shake128_init_absorb_final(input: &[u8; 34]) -> Self {
             shake128_init_absorb_final(input)
         }
 
         #[inline(always)]
-        fn shake128_squeeze_first_three_blocks(&mut self, output: &mut [[u8; THREE_BLOCKS]]) {
+        pub(crate) fn shake128_squeeze_first_three_blocks(
+            &mut self,
+            output: &mut [u8; THREE_BLOCKS],
+        ) {
             shake128_squeeze_first_three_blocks(self, output)
         }
 
         #[inline(always)]
-        fn shake128_squeeze_next_block(&mut self, output: &mut [[u8; BLOCK_SIZE]]) {
+        pub(crate) fn shake128_squeeze_next_block(&mut self, output: &mut [u8; BLOCK_SIZE]) {
             shake128_squeeze_next_block(self, output)
         }
     }