memo-cache 0.12.0

use criterion::{Criterion, criterion_group, criterion_main};
use memo_cache::MemoCache;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use rand_distr::{Distribution, Normal};
use std::{
    collections::{BTreeMap, HashMap, HashSet},
    hint::black_box,
    ops::RangeInclusive,
};

// Pseudo random number generator seed value (used for all benches).
const RNG_SEED_VALUE: u64 = 42;

// Uniform input data distribution settings.
const UNIFORM_RANGE_NARROW: RangeInclusive<i32> = -30..=30;
const UNIFORM_RANGE_WIDE: RangeInclusive<i32> = -100..=100;
const UNIFORM_RANGE_VERY_WIDE: RangeInclusive<i32> = -1000..=1000;

// Normal input data distribution settings.
const NORMAL_VARIANCE_NARROW: f32 = 5.0;
const NORMAL_VARIANCE_WIDE: f32 = 50.0;
const NORMAL_VARIANCE_VERY_WIDE: f32 = 500.0;

// Sequence length used in the benchmarks.
const SEQUENCE_LENGTH: usize = 1000;

// Macro to benchmark MemoCache with a specific size (const generic).
macro_rules! bench_memo_cache_size {
    ($b:expr, $sequence:expr, $size:expr) => {{
        let mut cache = MemoCache::<_, _, $size>::new();
        let mut seq = SequenceIterator::new($sequence.clone());
        $b.iter(|| {
            let input = seq.next();
            black_box(cache.get_or_insert_with(&input, |_| fake_expensive_calculation(input)));
        })
    }};
}

fn fake_expensive_calculation(x: i32) -> i32 {
    let mut x = x;
    for i in 0..10_000 {
        x = x.wrapping_add(i);
        x = x.wrapping_mul(17);
        x = x.rotate_left(5);
    }
    x
}

struct SequenceIterator {
    sequence: Vec<i32>,
    index: usize,
}

impl SequenceIterator {
    fn new(sequence: Vec<i32>) -> Self {
        Self { sequence, index: 0 }
    }

    fn next(&mut self) -> i32 {
        let val = self.sequence[self.index % self.sequence.len()];
        self.index += 1;
        val
    }
}

#[derive(Clone, Copy)]
enum TestDistribution {
    UniformNarrow,
    UniformWide,
    UniformVeryWide,
    NormalNarrow,
    NormalWide,
    NormalVeryWide,
}

impl TestDistribution {
    fn generate_sequence(&self) -> Vec<i32> {
        let mut rng = ChaCha8Rng::seed_from_u64(RNG_SEED_VALUE);

        match self {
            Self::UniformNarrow => (0..SEQUENCE_LENGTH)
                .map(|_| rng.random_range(UNIFORM_RANGE_NARROW))
                .collect(),
            Self::UniformWide => (0..SEQUENCE_LENGTH)
                .map(|_| rng.random_range(UNIFORM_RANGE_WIDE))
                .collect(),
            Self::UniformVeryWide => (0..SEQUENCE_LENGTH)
                .map(|_| rng.random_range(UNIFORM_RANGE_VERY_WIDE))
                .collect(),
            Self::NormalNarrow => {
                let normal = Normal::new(0.0, NORMAL_VARIANCE_NARROW).unwrap();
                (0..SEQUENCE_LENGTH)
                    .map(|_| normal.sample(&mut rng) as i32)
                    .collect()
            }
            Self::NormalWide => {
                let normal = Normal::new(0.0, NORMAL_VARIANCE_WIDE).unwrap();
                (0..SEQUENCE_LENGTH)
                    .map(|_| normal.sample(&mut rng) as i32)
                    .collect()
            }
            Self::NormalVeryWide => {
                let normal = Normal::new(0.0, NORMAL_VARIANCE_VERY_WIDE).unwrap();
                (0..SEQUENCE_LENGTH)
                    .map(|_| normal.sample(&mut rng) as i32)
                    .collect()
            }
        }
    }

    fn name(&self) -> &'static str {
        match self {
            Self::UniformNarrow => "Uniform distribution - narrow",
            Self::UniformWide => "Uniform distribution - wide",
            Self::UniformVeryWide => "Uniform distribution - very wide",
            Self::NormalNarrow => "Normal distribution - narrow",
            Self::NormalWide => "Normal distribution - wide",
            Self::NormalVeryWide => "Normal distribution - very wide",
        }
    }
}

fn analyze_sequence(name: &str, sequence: &[i32], cache_size: usize) {
    let unique_values = sequence.iter().collect::<HashSet<_>>().len();
    let duplicates = sequence.len() - unique_values;
    let hit_rate = (duplicates as f64 / sequence.len() as f64) * 100.0;

    println!("\nAnalysis for {}:", name);
    println!("  - Total operations:   {}", sequence.len());
    println!("  - Unique values:      {unique_values}");
    println!("  - Duplicate accesses: {duplicates}");
    println!("  - Potential hit rate: {hit_rate:.1}%");
    println!(
        "  - Estimated memory use (HashMap):   {} bytes",
        unique_values * (std::mem::size_of::<i32>() + std::mem::size_of::<i32>() + 8) + 48
    );
    println!(
        "  - Estimated memory use (BTreeMap):  {} bytes",
        unique_values * (std::mem::size_of::<i32>() + std::mem::size_of::<i32>() + 12) + 128
    );
    println!(
        "  - Estimated memory use (MemoCache): {} bytes",
        cache_size * (std::mem::size_of::<i32>() + std::mem::size_of::<i32>() + 1) + 4
    );
    println!();

    if unique_values > cache_size {
        println!(
            "  ⚠️ Working set ({}) exceeds MemoCache capacity ({})",
            unique_values, cache_size
        );
    }
}

fn bench_distribution(c: &mut Criterion, distribution: TestDistribution, cache_size: usize) {
    let group_name = format!("{} (size={})", distribution.name(), cache_size);
    let mut g = c.benchmark_group(&group_name);
    let sequence = distribution.generate_sequence();

    analyze_sequence(&group_name, &sequence, cache_size);

    g.bench_function("HashMap", |b| {
        let mut cache = HashMap::new();
        let mut seq = SequenceIterator::new(sequence.clone());

        b.iter(|| {
            let input = seq.next();
            black_box(
                cache
                    .entry(input)
                    .or_insert_with(|| fake_expensive_calculation(input)),
            );
        })
    });

    g.bench_function("BTreeMap", |b| {
        let mut cache = BTreeMap::new();
        let mut seq = SequenceIterator::new(sequence.clone());

        b.iter(|| {
            let input = seq.next();
            black_box(
                cache
                    .entry(input)
                    .or_insert_with(|| fake_expensive_calculation(input)),
            );
        })
    });

    g.bench_function("MemoCache", |b| match cache_size {
        8 => bench_memo_cache_size!(b, sequence, 8),
        16 => bench_memo_cache_size!(b, sequence, 16),
        32 => bench_memo_cache_size!(b, sequence, 32),
        64 => bench_memo_cache_size!(b, sequence, 64),
        128 => bench_memo_cache_size!(b, sequence, 128),
        _ => panic!("Unsupported cache size: {}", cache_size),
    });

    g.finish();
}

fn bench_all_distributions(c: &mut Criterion) {
    for dist in [
        TestDistribution::UniformNarrow,
        TestDistribution::UniformWide,
        TestDistribution::UniformVeryWide,
        TestDistribution::NormalNarrow,
        TestDistribution::NormalWide,
        TestDistribution::NormalVeryWide,
    ] {
        for size in [8, 16, 32, 64, 128] {
            bench_distribution(c, dist, size);
        }
    }
}

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