use bytemuck::*;
use criterion::{criterion_group, criterion_main, Criterion, Fun};
/**
 * This module benchmarks the performance of iterating over chunks of data vs. reinterpreting an
 * array as chunked data.
 */
use rand::distributions::{Distribution, Standard};
use rand::prelude::*;

static SEED: [u8; 32] = [3; 32];
static BUF_SIZE: usize = 750_000;

#[inline]
fn make_random_vec<T>() -> Vec<T>
where
    Standard: Distribution<T>,
{
    let n = BUF_SIZE;

    let mut rng: StdRng = SeedableRng::from_seed(SEED);
    let mut vec = Vec::new();
    for _ in 0..n {
        vec.push(rng.gen::<T>());
    }
    vec
}

fn compute(x: f64, y: f64, z: f64) -> [f64; 3] {
    [
        x * 3.6321 + 42314.0 * y + z * 2.1,
        y * 3.6321 + 42314.0 * z + x * 2.1,
        z * 3.6321 + 42314.0 * x + y * 2.1,
    ]
}

// Vanilla iteration over a single vector
fn chunks(v: &mut Vec<f64>) {
    for a in v.chunks_exact_mut(3) {
        let res = compute(a[0], a[1], a[2]);
        a[0] = res[0];
        a[1] = res[1];
        a[2] = res[2];
    }
}

fn reinterpret(v: &mut Vec<f64>) {
    let s: &mut [[f64; 3]] = cast_slice_mut(v.as_mut_slice());
    for a in s.iter_mut() {
        *a = compute(a[0], a[1], a[2]);
    }
}

fn chunks_iter(c: &mut Criterion) {
    let chunks = Fun::new("Chunks", move |b, _| {
        let mut v = make_random_vec::<f64>();
        b.iter(|| {
            chunks(&mut v);
        })
    });

    let reinterpret = Fun::new("Reinterpret", move |b, _| {
        let mut v = make_random_vec::<f64>();
        b.iter(|| {
            reinterpret(&mut v);
        })
    });

    let fns = vec![chunks, reinterpret];
    c.bench_functions("chunks_iter", fns, ());
}

criterion_group!(benches, chunks_iter);
criterion_main!(benches);