Add some benchmarks, move most to lib.rs, add do_chacha function, update

Cargo.lock

@@ -185,9 +185,9 @@ dependencies = [
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.18"
+version = "1.0.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
+checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d"
 
 [[package]]
 name = "utf8parse"

src/lib.rs

@@ -0,0 +1,177 @@
+#![feature(portable_simd)]
+#![feature(slice_as_array)]
+#![feature(test)]
+
+extern crate test;
+
+use std::simd::u32x16;
+
+#[inline]
+fn many_u8_to_few_u32(data: [u8; 64]) -> [u32; 16] {
+    let mut out = [0u32; 16];
+
+    for (i, c) in data.chunks_exact(4).enumerate() {
+        out[i] = u32::from_ne_bytes(*c.as_array().unwrap());
+    }
+
+    out
+}
+
+#[inline]
+fn few_u32_to_many_u8(data: &[u32; 16]) -> [u8; 64] {
+    let mut out = [0u8; 64];
+
+    for (i, c) in data.iter().enumerate() {
+        out[i * 4..][..4].copy_from_slice(&u32::to_ne_bytes(*c));
+    }
+
+    out
+}
+
+#[inline]
+fn make_chacha_block(key: &[u32; 8], counter: u64, nonce: u64) -> u32x16 {
+    let [k1, k2, k3, k4, k5, k6, k7, k8] = key;
+
+    u32x16::from_array([
+        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,
+    ])
+}
+
+pub fn do_chacha(input: &[u8], key_bytes: &[u8; 32], nonce: u64, initial_counter: u64, rounds: usize) -> Vec<u8> {
+    let mut key: [u32; 8] = [0u32; 8];
+    for i in 0..8 {
+        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];
+        x[..c.len()].copy_from_slice(c);
+
+        let cipher = chacha_block(&key, i as u64 + initial_counter, nonce, rounds);
+        let data = u32x16::from_array(many_u8_to_few_u32(x));
+        let new = cipher ^ data;
+        bytes.extend_from_slice(&few_u32_to_many_u8(new.as_array())[..c.len()]);
+    }
+
+    bytes
+}
+
+fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
+    let mut x: u32x16 = input;
+
+    for i in 0..rounds {
+        match (i & 1) == 0 {
+            true => {
+                // 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);
+            },
+            false => {
+                // 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);
+            }
+        }
+    }
+
+    x + input
+}
+
+#[inline]
+fn chacha_block(key: &[u32; 8], counter: u64, nonce: u64, rounds: usize) -> u32x16 {
+    let block = make_chacha_block(key, counter, nonce);
+
+    chacha_rounds(block, rounds)
+}
+
+#[inline]
+fn chacha_quarter_round(x: &mut u32x16, a: usize, b: usize, c: usize, d: usize) {
+    x[a] = x[a].wrapping_add(x[b]);
+    x[d] ^= x[a];
+    x[d] = x[d].rotate_left(16);
+
+    x[c] = x[c].wrapping_add(x[d]);
+    x[b] ^= x[c];
+    x[b] = x[b].rotate_left(12);
+
+    x[a] = x[a].wrapping_add(x[b]);
+    x[d] ^= x[a];
+    x[d] = x[d].rotate_left(8);
+
+    x[c] = x[c].wrapping_add(x[d]);
+    x[b] ^= x[c];
+    x[b] = x[b].rotate_left(7);
+}
+
+#[cfg(test)]
+mod tests {
+    use test::{Bencher, black_box};
+    use crate::{chacha_quarter_round, chacha_rounds, do_chacha, few_u32_to_many_u8, make_chacha_block};
+    use std::simd::u32x16;
+
+    #[test]
+    fn test_chacha_quarter_round() {
+        let mut data = [0u32; 16];
+        data[0] = 0x11111111;
+        data[1] = 0x01020304;
+        data[2] = 0x9b8d6f43;
+        data[3] = 0x01234567;
+
+        let mut arr = u32x16::from_array(data);
+
+        chacha_quarter_round(&mut arr, 0, 1, 2, 3);
+
+        assert_eq!(arr[0], 0xea2a92f4);
+        assert_eq!(arr[1], 0xcb1cf8ce);
+        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_array([69u32; 16]);
+
+        b.iter(|| {
+            black_box(chacha_rounds(data, 20));
+        })
+    }
+
+    #[bench]
+    fn bench_chacha_operation(b: &mut Bencher) {
+        // somewhat random data
+        let data = few_u32_to_many_u8(chacha_rounds(u32x16::from_array([69u32; 16]), 20).as_array());
+
+        let key = &[0u8; 32];
+
+        b.iter(|| {
+            black_box(do_chacha(data.as_slice(), key, 0, 0, 20));
+        })
+    }
+}

src/main.rs

@@ -1,14 +1,8 @@
-#![feature(portable_simd)]
-#![feature(slice_as_array)]
-#![feature(test)]
-
 use clap::Parser;
 use clap_num::maybe_hex;
 use std::fs::File;
 use std::io::{Read, Write, stdin, stdout};
-use std::simd::u32x16;
+use chacha::do_chacha;
-
-extern crate test;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
@@ -40,162 +34,9 @@ fn main() {
     let key_bytes = &mut [0u8; 32];
     f.read_exact(key_bytes).expect("couldn't read key file");
 
-    let mut key: [u32; 8] = [0u32; 8];
+    let bytes = do_chacha(input.as_slice(), key_bytes, args.nonce, 0, args.rounds);
-    for i in 0..8 {
-        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];
-        x[..c.len()].copy_from_slice(c);
-
-        let cipher = chacha_block(&key, i as u64, args.nonce, args.rounds);
-        let data = u32x16::from(many_u8_to_few_u32(x));
-        let new = cipher ^ data;
-        bytes.extend_from_slice(&few_u32_to_many_u8(new.as_array())[..c.len()]);
-    }
 
     let stdout = stdout();
     let mut out = stdout.lock();
     out.write_all(bytes.as_slice()).unwrap();
 }
-
-#[inline]
-fn many_u8_to_few_u32(data: [u8; 64]) -> [u32; 16] {
-    let mut out = [0u32; 16];
-
-    for (i, c) in data.chunks_exact(4).enumerate() {
-        out[i] = u32::from_ne_bytes(c.try_into().unwrap());
-    }
-
-    out
-}
-
-#[inline]
-fn few_u32_to_many_u8(data: &[u32; 16]) -> [u8; 64] {
-    let mut out = [0u8; 64];
-
-    for (i, c) in data.iter().enumerate() {
-        out[i * 4..][..4].copy_from_slice(&u32::to_ne_bytes(*c));
-    }
-
-    out
-}
-
-#[inline]
-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,
-    ])
-}
-
-fn chacha_rounds(input: u32x16, rounds: usize) -> u32x16 {
-    let mut x: u32x16 = input;
-
-    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
-}
-
-#[inline]
-fn chacha_block(key: &[u32; 8], counter: u64, nonce: u64, rounds: usize) -> u32x16 {
-    let block = make_chacha_block(key, counter, nonce);
-
-    chacha_rounds(block, rounds)
-}
-
-#[inline]
-fn chacha_quarter_round(x: &mut u32x16, a: usize, b: usize, c: usize, d: usize) {
-    x[a] = x[a].wrapping_add(x[b]);
-    x[d] ^= x[a];
-    x[d] = x[d].rotate_left(16);
-
-    x[c] = x[c].wrapping_add(x[d]);
-    x[b] ^= x[c];
-    x[b] = x[b].rotate_left(12);
-
-    x[a] = x[a].wrapping_add(x[b]);
-    x[d] ^= x[a];
-    x[d] = x[d].rotate_left(8);
-
-    x[c] = x[c].wrapping_add(x[d]);
-    x[b] ^= x[c];
-    x[b] = x[b].rotate_left(7);
-}
-
-#[cfg(test)]
-mod tests {
-    use test::{Bencher, black_box};
-    use crate::{chacha_quarter_round, chacha_rounds, make_chacha_block};
-    use std::simd::u32x16;
-
-    #[test]
-    fn test_chacha_quarter_round() {
-        let mut data = [0u32; 16];
-        data[0] = 0x11111111;
-        data[1] = 0x01020304;
-        data[2] = 0x9b8d6f43;
-        data[3] = 0x01234567;
-
-        let mut arr = u32x16::from(data);
-
-        chacha_quarter_round(&mut arr, 0, 1, 2, 3);
-
-        assert_eq!(arr[0], 0xea2a92f4);
-        assert_eq!(arr[1], 0xcb1cf8ce);
-        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));
-        })
-    }
-}