Formatting, benchmarks and slight improvement to chacha_rounds function

src/main.rs

@@ -1,11 +1,14 @@
 #![feature(portable_simd)]
 #![feature(slice_as_array)]
+#![feature(test)]
 
-use std::fs::File;
-use std::io::{stdin, stdout, Read, Write};
-use std::simd::u32x16;
 use clap::Parser;
 use clap_num::maybe_hex;
+use std::fs::File;
+use std::io::{Read, Write, stdin, stdout};
+use std::simd::u32x16;
+
+extern crate test;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
@@ -30,7 +33,8 @@ fn main() {
 
     let mut input = Vec::new();
     let mut inp = stdin.lock();
-    inp.read_to_end(&mut input).expect("failed to read from stdin");
+    inp.read_to_end(&mut input)
+        .expect("failed to read from stdin");
 
     let mut f = File::open(args.key_file).expect("failed to open key file");
     let key_bytes = &mut [0u8; 32];
@@ -41,7 +45,6 @@ fn main() {
         key[i] = u32::from_le_bytes(key_bytes[i * 4..(i + 1) * 4].try_into().unwrap());
     }
 
-
     let mut bytes: Vec<u8> = Vec::new();
     for (i, c) in input.chunks(64).enumerate() {
         let mut x = [0u8; 64];
@@ -85,28 +88,47 @@ fn make_chacha_block(key: &[u32; 8], counter: u64, nonce: u64) -> u32x16 {
     let [k1, k2, k3, k4, k5, k6, k7, k8] = key;
 
     u32x16::from([
-        u32::from_ne_bytes(*b"expa"), u32::from_ne_bytes(*b"nd 3"), u32::from_ne_bytes(*b"2-by"), u32::from_ne_bytes(*b"te k"),
-        *k1,                                *k2,                                *k3,                                *k4,
-        *k5,                                *k6,                                *k7,                                *k8,
-        (counter >> 32) as u32,             (counter & 0xFFFFFFFF) as u32,      (nonce >> 32) as u32,               (nonce & 0xFFFFFFFF) as u32
+        u32::from_ne_bytes(*b"expa"),
+        u32::from_ne_bytes(*b"nd 3"),
+        u32::from_ne_bytes(*b"2-by"),
+        u32::from_ne_bytes(*b"te k"),
+        *k1,
+        *k2,
+        *k3,
+        *k4,
+        *k5,
+        *k6,
+        *k7,
+        *k8,
+        (counter >> 32) as u32,
+        (counter & 0xFFFFFFFF) as u32,
+        (nonce >> 32) as u32,
+        (nonce & 0xFFFFFFFF) as u32,
     ])
 }
 
 fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
     let mut x: u32x16 = input;
 
-    for _ in (0..rounds).step_by(2) {
+    for i in (0..rounds).step_by(2) {
+        match i & 1 {
+            0 => {
                 // odd rounds
                 chacha_quarter_round(&mut x, 0, 4, 8, 12);
                 chacha_quarter_round(&mut x, 1, 5, 9, 13);
                 chacha_quarter_round(&mut x, 2, 6, 10, 14);
                 chacha_quarter_round(&mut x, 3, 7, 11, 15);
+            },
+            1 => {
                 // even rounds
                 chacha_quarter_round(&mut x, 0, 5, 10, 15);
                 chacha_quarter_round(&mut x, 1, 6, 11, 12);
                 chacha_quarter_round(&mut x, 2, 7, 8, 13);
                 chacha_quarter_round(&mut x, 3, 4, 9, 14);
             }
+            _ => unreachable!()
+        }
+    }
 
     x + input
 }
@@ -139,8 +161,9 @@ fn chacha_quarter_round(x: &mut u32x16, a: usize, b: usize, c: usize, d: usize)
 
 #[cfg(test)]
 mod tests {
+    use test::{Bencher, black_box};
+    use crate::{chacha_quarter_round, chacha_rounds, make_chacha_block};
     use std::simd::u32x16;
-    use crate::chacha_quarter_round;
 
     #[test]
     fn test_chacha_quarter_round() {
@@ -159,4 +182,20 @@ mod tests {
         assert_eq!(arr[2], 0x4581472e);
         assert_eq!(arr[3], 0x5881c4bb);
     }
+
+    #[bench]
+    fn bench_make_chacha_block(b: &mut Bencher) {
+        b.iter(|| {
+            black_box(make_chacha_block(&[1u32; 8], 0, 0))
+        })
+    }
+
+    #[bench]
+    fn bench_chacha_rounds(b: &mut Bencher) {
+        let data = u32x16::from([69u32; 16]);
+
+        b.iter(|| {
+            black_box(chacha_rounds(data, 20));
+        })
+    }
 }