perf: eliminate bounds checks in concat_kdf

codahale · Nov 18, 2023 · 01f526e · 01f526e
1 parent bdeac03
commit 01f526e
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -250,17 +250,28 @@ impl Protocol {
 fn concat_kdf(ikm: &[u8], kdf_key: &mut [u8; 32], out: &mut [u8]) {
     let mut counter = 1u32;
     let mut kdf = Sha256::new();
+
     kdf.update(counter.to_be_bytes());
     kdf.update(ikm);
     kdf.update(b"lockstitch");
     kdf.finalize_into_reset(kdf_key.into());
 
-    for chunk in out.chunks_mut(32) {
+    let mut chunks = out.chunks_exact_mut(32);
+    for chunk in chunks.by_ref() {
+        counter += 1;
+        kdf.update(counter.to_be_bytes());
+        kdf.update(ikm);
+        kdf.update(b"lockstitch");
+        chunk.copy_from_slice(&kdf.finalize_reset());
+    }
+
+    let remainder = chunks.into_remainder();
+    if !remainder.is_empty() {
         counter += 1;
         kdf.update(counter.to_be_bytes());
         kdf.update(ikm);
         kdf.update(b"lockstitch");
-        chunk.copy_from_slice(&kdf.finalize_reset()[..chunk.len()]);
+        remainder.copy_from_slice(&kdf.finalize_reset()[..remainder.len()]);
     }
 }