memo-cache 0.13.0

use criterion::{Criterion, Throughput, 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,
    mem::size_of,
    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 (standard deviations).
const NORMAL_STD_DEV_NARROW: f32 = 5.0;
const NORMAL_STD_DEV_WIDE: f32 = 50.0;
const NORMAL_STD_DEV_VERY_WIDE: f32 = 500.0;

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

// Cache sizes used in the benchmarks (must match the dispatch tables below).
const CACHE_SIZES: [usize; 5] = [8, 16, 32, 64, 128];

// Macro to benchmark steady-state MemoCache use with a specific size (const generic).
//
// The cache is prefilled with one pass over the sequence so that every measured iteration
// starts from (and returns the cache to) the same steady state, doing identical work.
macro_rules! bench_memo_cache_size {
    ($g:expr, $sequence:expr, $size:expr) => {{
        let sequence: &[i32] = $sequence;
        $g.bench_function("MemoCache", |b| {
            let mut cache = MemoCache::<_, _, $size>::new();
            for &x in sequence {
                cache.get_or_insert_with(&x, |_| fake_expensive_calculation(x));
            }
            b.iter(|| {
                for &x in sequence {
                    black_box(cache.get_or_insert_with(&x, |_| fake_expensive_calculation(x)));
                }
            })
        });
    }};
}

// Macro to benchmark the MemoCache hit path (lookups only, on a full cache) with a
// specific size (const generic).
macro_rules! bench_memo_cache_hit_path {
    ($g:expr, $keys:expr, $size:expr) => {{
        let keys: &[i32] = $keys;
        $g.bench_function("MemoCache", |b| {
            let mut cache = MemoCache::<i32, i32, $size>::new();
            for &k in keys {
                cache.insert(k, k.wrapping_mul(2));
            }
            b.iter(|| {
                for k in keys {
                    black_box(cache.get(k));
                }
            })
        });
    }};
}

// Macro to benchmark the MemoCache churn (miss/evict) path with a specific size (const
// generic). The working set is twice the cache size and the compute function is trivial,
// so the numbers isolate the cost of key search, eviction-slot search, and replacement.
macro_rules! bench_memo_cache_churn {
    ($g:expr, $keys:expr, $size:expr) => {{
        let keys: &[i32] = $keys;
        $g.bench_function("MemoCache", |b| {
            let mut cache = MemoCache::<i32, i32, $size>::new();
            for &k in keys {
                cache.get_or_insert_with(&k, |&k| k);
            }
            b.iter(|| {
                for &k in keys {
                    black_box(cache.get_or_insert_with(&k, |&k| k));
                }
            })
        });
    }};
}

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
}

#[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);

        let uniform = |rng: &mut ChaCha8Rng, range: RangeInclusive<i32>| {
            (0..SEQUENCE_LENGTH)
                .map(|_| rng.random_range(range.clone()))
                .collect()
        };

        // NOTE: Round instead of truncate; `as i32` truncates toward zero, which would
        //       make the 0 bucket twice as wide as every other one.
        let normal = |rng: &mut ChaCha8Rng, std_dev: f32| {
            let normal = Normal::new(0.0, std_dev).unwrap();
            (0..SEQUENCE_LENGTH)
                .map(|_| normal.sample(rng).round() as i32)
                .collect()
        };

        match self {
            Self::UniformNarrow => uniform(&mut rng, UNIFORM_RANGE_NARROW),
            Self::UniformWide => uniform(&mut rng, UNIFORM_RANGE_WIDE),
            Self::UniformVeryWide => uniform(&mut rng, UNIFORM_RANGE_VERY_WIDE),
            Self::NormalNarrow => normal(&mut rng, NORMAL_STD_DEV_NARROW),
            Self::NormalWide => normal(&mut rng, NORMAL_STD_DEV_WIDE),
            Self::NormalVeryWide => normal(&mut rng, NORMAL_STD_DEV_VERY_WIDE),
        }
    }

    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",
        }
    }
}

/// Get the exact size in bytes of a `MemoCache<i32, i32, SIZE>` for a runtime size value.
fn memo_cache_size_bytes(cache_size: usize) -> usize {
    match cache_size {
        8 => size_of::<MemoCache<i32, i32, 8>>(),
        16 => size_of::<MemoCache<i32, i32, 16>>(),
        32 => size_of::<MemoCache<i32, i32, 32>>(),
        64 => size_of::<MemoCache<i32, i32, 64>>(),
        128 => size_of::<MemoCache<i32, i32, 128>>(),
        _ => panic!("Unsupported cache size: {cache_size}"),
    }
}

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;

    // Measure the real allocation of a HashMap holding the unique values; hashbrown sizes
    // its table by bucket count (usable capacity is ~7/8 of the buckets), so derive the
    // bucket count from the reported capacity.
    let map: HashMap<i32, i32> = sequence.iter().map(|&x| (x, 0)).collect();
    let buckets = (map.capacity() * 8 / 7).next_power_of_two();
    let hashmap_bytes = buckets * (size_of::<(i32, i32)>() + 1) + size_of::<HashMap<i32, i32>>();

    // BTreeMap nodes hold up to 11 entries and are ~70% full on random insertion;
    // ~13 bytes per (i32, i32) entry including node overhead.
    let btreemap_bytes = unique_values * 13 + size_of::<BTreeMap<i32, i32>>();

    println!("\nAnalysis for {name}:");
    println!("  - Total operations:   {}", sequence.len());
    println!("  - Unique values:      {unique_values}");
    println!("  - Duplicate accesses: {duplicates}");
    println!(
        "  - Upper-bound hit rate: {hit_rate:.1}% (attainable by the unbounded maps; for MemoCache this is an upper bound)"
    );
    println!("  - Memory use (HashMap, ≈ from real capacity): {hashmap_bytes} bytes");
    println!("  - Memory use (BTreeMap, rough estimate):      {btreemap_bytes} bytes");
    println!(
        "  - Memory use (MemoCache, exact):              {} bytes",
        memo_cache_size_bytes(cache_size)
    );
    println!();

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

/// Benchmark steady-state memoization: every iteration performs one full pass over the
/// input sequence on a prefilled cache, so per-iteration work is identical and
/// deterministic. Misses pay `fake_expensive_calculation`, so these numbers are dominated
/// by miss rate × recompute cost; see the hit-path benchmarks for cache overhead itself.
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.throughput(Throughput::Elements(SEQUENCE_LENGTH as u64));

    g.bench_function("HashMap", |b| {
        let mut cache = HashMap::new();
        for &x in &sequence {
            cache
                .entry(x)
                .or_insert_with(|| fake_expensive_calculation(x));
        }
        b.iter(|| {
            for &x in &sequence {
                black_box(
                    cache
                        .entry(x)
                        .or_insert_with(|| fake_expensive_calculation(x)),
                );
            }
        });
    });

    g.bench_function("BTreeMap", |b| {
        let mut cache = BTreeMap::new();
        for &x in &sequence {
            cache
                .entry(x)
                .or_insert_with(|| fake_expensive_calculation(x));
        }
        b.iter(|| {
            for &x in &sequence {
                black_box(
                    cache
                        .entry(x)
                        .or_insert_with(|| fake_expensive_calculation(x)),
                );
            }
        });
    });

    match cache_size {
        8 => bench_memo_cache_size!(g, &sequence, 8),
        16 => bench_memo_cache_size!(g, &sequence, 16),
        32 => bench_memo_cache_size!(g, &sequence, 32),
        64 => bench_memo_cache_size!(g, &sequence, 64),
        128 => bench_memo_cache_size!(g, &sequence, 128),
        _ => panic!("Unsupported cache size: {cache_size}"),
    }

    g.finish();
}

/// Benchmark the hit path only: lookups on a full cache, no expensive recompute. This
/// isolates the cost of the cache itself (linear scan vs. hashing vs. tree search).
fn bench_hit_path(c: &mut Criterion, cache_size: usize) {
    let group_name = format!("Hit-path lookup (size={cache_size})");
    let mut g = c.benchmark_group(&group_name);

    g.throughput(Throughput::Elements(cache_size as u64));

    #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
    let keys = (0..cache_size as i32).collect::<Vec<_>>();

    g.bench_function("HashMap", |b| {
        let cache: HashMap<i32, i32> = keys.iter().map(|&k| (k, k.wrapping_mul(2))).collect();
        b.iter(|| {
            for k in &keys {
                black_box(cache.get(k));
            }
        });
    });

    g.bench_function("BTreeMap", |b| {
        let cache: BTreeMap<i32, i32> = keys.iter().map(|&k| (k, k.wrapping_mul(2))).collect();
        b.iter(|| {
            for k in &keys {
                black_box(cache.get(k));
            }
        });
    });

    match cache_size {
        8 => bench_memo_cache_hit_path!(g, &keys, 8),
        16 => bench_memo_cache_hit_path!(g, &keys, 16),
        32 => bench_memo_cache_hit_path!(g, &keys, 32),
        64 => bench_memo_cache_hit_path!(g, &keys, 64),
        128 => bench_memo_cache_hit_path!(g, &keys, 128),
        _ => panic!("Unsupported cache size: {cache_size}"),
    }

    g.finish();
}

/// Benchmark the churn path: a cycling working set of twice the cache size with a trivial
/// compute function, so (almost) every access misses and evicts. This isolates the cache's
/// own miss/insert overhead, which the steady-state benchmarks drown in recompute cost.
fn bench_churn_path(c: &mut Criterion, cache_size: usize) {
    let group_name = format!("Churn path (size={cache_size})");
    let mut g = c.benchmark_group(&group_name);

    #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
    let keys = (0..2 * cache_size as i32).collect::<Vec<_>>();

    g.throughput(Throughput::Elements(keys.len() as u64));

    match cache_size {
        8 => bench_memo_cache_churn!(g, &keys, 8),
        16 => bench_memo_cache_churn!(g, &keys, 16),
        32 => bench_memo_cache_churn!(g, &keys, 32),
        64 => bench_memo_cache_churn!(g, &keys, 64),
        128 => bench_memo_cache_churn!(g, &keys, 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 CACHE_SIZES {
            bench_distribution(c, dist, size);
        }
    }
}

fn bench_all_hit_paths(c: &mut Criterion) {
    for size in CACHE_SIZES {
        bench_hit_path(c, size);
    }
}

fn bench_all_churn_paths(c: &mut Criterion) {
    for size in CACHE_SIZES {
        bench_churn_path(c, size);
    }
}

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