use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
use rand::{
    RngExt, SeedableRng,
    distr::{Distribution, StandardUniform},
};
use rand_xoshiro::Xoshiro128PlusPlus;
use small_is_prime::IsPrime;

fn bench_is_prime<T>(c: &mut Criterion)
where
    T: IsPrime + Copy,
    StandardUniform: Distribution<T>,
{
    let mut rng = Xoshiro128PlusPlus::seed_from_u64(42);
    let mut acc = false;
    for filter in [None, Some(false), Some(true)] {
        let filter_label = match filter {
            None => "any",
            Some(true) => "prime",
            Some(false) => "composite",
        };
        c.bench_function(
            &format!("{filter_label} {}", std::any::type_name::<T>()),
            |b| {
                b.iter_batched(
                    || loop {
                        let x = rng.random::<T>();
                        if filter.map(|f| x.is_prime() == f).unwrap_or(true) {
                            break x;
                        }
                    },
                    |x| {
                        acc ^= x.is_prime();
                    },
                    BatchSize::SmallInput,
                );
            },
        );
    }
}

criterion_group!(
    benches,
    bench_is_prime::<u64>,
    bench_is_prime::<u32>,
    bench_is_prime::<u16>,
    bench_is_prime::<u8>
);
criterion_main!(benches);
