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

This commit is contained in:
Neemek 2025-09-10 14:02:30 +02:00
parent 943f041759
commit 1942ba7d8f
Signed by: neemek
GPG key ID: 28360A8951CD0E9B
3 changed files with 181 additions and 163 deletions

4
Cargo.lock generated
View file

@ -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"

177
src/lib.rs Normal file
View file

@ -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));
})
}
}

View file

@ -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;
extern crate test;
use chacha::do_chacha;
#[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];
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 bytes = do_chacha(input.as_slice(), key_bytes, args.nonce, 0, args.rounds);
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));
})
}
}