use std::collections::{BTreeMap, HashMap, VecDeque};
use crate::bloom_filter::{BloomFilter, CountingBloomFilter};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EvictionPolicy {
Lru,
Lfu,
Fifo,
Random,
TinyLfu,
ArcCache,
}
#[derive(Debug, Clone)]
pub struct FrequencyCounter {
counters: HashMap<u64, u64>,
window_size: usize,
total_increments: usize,
}
impl FrequencyCounter {
pub fn new(window_size: usize) -> Self {
Self {
counters: HashMap::new(),
window_size: window_size.max(1),
total_increments: 0,
}
}
pub fn increment(&mut self, key: u64) {
*self.counters.entry(key).or_insert(0) += 1;
self.total_increments += 1;
if self.total_increments >= self.window_size {
self.decay_all();
}
}
pub fn frequency(&self, key: u64) -> u64 {
self.counters.get(&key).copied().unwrap_or(0)
}
pub fn decay_all(&mut self) {
self.counters.retain(|_, count| {
*count >>= 1;
*count > 0
});
self.total_increments = 0;
}
pub fn tracked_keys(&self) -> usize {
self.counters.len()
}
pub fn clear(&mut self) {
self.counters.clear();
self.total_increments = 0;
}
}
#[derive(Debug, Clone)]
pub struct LfuEvictionTracker {
freq_buckets: BTreeMap<u64, VecDeque<u64>>,
key_freq: HashMap<u64, u64>,
}
impl LfuEvictionTracker {
pub fn new() -> Self {
Self {
freq_buckets: BTreeMap::new(),
key_freq: HashMap::new(),
}
}
pub fn insert(&mut self, key: u64) {
if self.key_freq.contains_key(&key) {
return;
}
self.key_freq.insert(key, 1);
self.freq_buckets.entry(1).or_default().push_back(key);
}
pub fn promote(&mut self, key: u64) {
let old_freq = match self.key_freq.get_mut(&key) {
Some(f) => *f,
None => return,
};
let new_freq = old_freq.saturating_add(1);
if let Some(bucket) = self.freq_buckets.get_mut(&old_freq) {
bucket.retain(|&k| k != key);
if bucket.is_empty() {
self.freq_buckets.remove(&old_freq);
}
}
self.key_freq.insert(key, new_freq);
self.freq_buckets
.entry(new_freq)
.or_default()
.push_back(key);
}
pub fn evict(&mut self) -> Option<u64> {
let (&min_freq, _) = self.freq_buckets.iter().next()?;
let victim = self
.freq_buckets
.get_mut(&min_freq)
.and_then(|q| q.pop_front())?;
if let Some(bucket) = self.freq_buckets.get(&min_freq) {
if bucket.is_empty() {
self.freq_buckets.remove(&min_freq);
}
}
self.key_freq.remove(&victim);
Some(victim)
}
pub fn remove(&mut self, key: u64) -> bool {
if let Some(freq) = self.key_freq.remove(&key) {
if let Some(bucket) = self.freq_buckets.get_mut(&freq) {
bucket.retain(|&k| k != key);
if bucket.is_empty() {
self.freq_buckets.remove(&freq);
}
}
true
} else {
false
}
}
pub fn frequency(&self, key: u64) -> u64 {
self.key_freq.get(&key).copied().unwrap_or(0)
}
pub fn len(&self) -> usize {
self.key_freq.len()
}
pub fn is_empty(&self) -> bool {
self.key_freq.is_empty()
}
}
impl Default for LfuEvictionTracker {
fn default() -> Self {
Self::new()
}
}
pub struct TinyLfuAdmission {
bloom: BloomFilter,
window_counter: CountingBloomFilter,
freq_counter: FrequencyCounter,
}
impl TinyLfuAdmission {
pub fn new(expected_items: usize) -> Self {
let window_size = expected_items * 4; Self {
bloom: BloomFilter::new(expected_items.max(1), 0.01),
window_counter: CountingBloomFilter::new(expected_items.max(1), 0.01),
freq_counter: FrequencyCounter::new(window_size.max(1)),
}
}
pub fn record_access(&mut self, candidate_key: u64) {
let key_bytes = candidate_key.to_le_bytes();
if !self.bloom.contains(&key_bytes) {
self.bloom.insert(&key_bytes);
} else {
self.window_counter.insert(&key_bytes);
}
self.freq_counter.increment(candidate_key);
}
pub fn should_admit(&mut self, candidate_key: u64, evicted_freq: u64) -> bool {
self.record_access(candidate_key);
let candidate_freq = self.freq_counter.frequency(candidate_key);
candidate_freq > evicted_freq
}
pub fn estimated_frequency(&self, key: u64) -> u64 {
self.freq_counter.frequency(key)
}
pub fn decay(&mut self) {
self.freq_counter.decay_all();
}
}
#[derive(Debug, Clone)]
pub struct ArcTracker {
pub t1_size: usize,
pub t2_size: usize,
pub b1_size: usize,
pub b2_size: usize,
pub p: usize,
capacity: usize,
}
impl ArcTracker {
pub fn new(capacity: usize) -> Self {
Self {
t1_size: 0,
t2_size: 0,
b1_size: 0,
b2_size: 0,
p: 0,
capacity,
}
}
pub fn adapt_on_hit_b1(&mut self) {
let delta = self.b2_size.checked_div(self.b1_size).unwrap_or(0).max(1);
self.p = self.p.saturating_add(delta).min(self.capacity);
}
pub fn adapt_on_hit_b2(&mut self) {
let delta = self.b1_size.checked_div(self.b2_size).unwrap_or(0).max(1);
self.p = self.p.saturating_sub(delta);
}
pub fn on_admit_t1(&mut self) {
self.t1_size += 1;
}
pub fn on_promote_t1_to_t2(&mut self) {
if self.t1_size > 0 {
self.t1_size -= 1;
}
self.t2_size += 1;
}
pub fn on_evict_t1(&mut self) {
if self.t1_size > 0 {
self.t1_size -= 1;
}
self.b1_size += 1;
}
pub fn on_evict_t2(&mut self) {
if self.t2_size > 0 {
self.t2_size -= 1;
}
self.b2_size += 1;
}
pub fn on_remove_b1_ghost(&mut self) {
if self.b1_size > 0 {
self.b1_size -= 1;
}
}
pub fn on_remove_b2_ghost(&mut self) {
if self.b2_size > 0 {
self.b2_size -= 1;
}
}
pub fn live_size(&self) -> usize {
self.t1_size + self.t2_size
}
pub fn ghost_size(&self) -> usize {
self.b1_size + self.b2_size
}
pub fn is_full(&self) -> bool {
self.live_size() >= self.capacity
}
pub fn should_evict_t1(&self) -> bool {
(self.t1_size > 0) && (self.t1_size > self.p || self.t2_size == 0)
}
pub fn capacity(&self) -> usize {
self.capacity
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_frequency_counter_increment() {
let mut fc = FrequencyCounter::new(100);
fc.increment(42);
fc.increment(42);
fc.increment(42);
assert_eq!(fc.frequency(42), 3);
}
#[test]
fn test_frequency_counter_absent_key() {
let fc = FrequencyCounter::new(50);
assert_eq!(fc.frequency(999), 0);
}
#[test]
fn test_frequency_counter_decay() {
let mut fc = FrequencyCounter::new(100);
fc.increment(1);
fc.increment(1);
fc.increment(1);
fc.increment(1); fc.decay_all();
assert_eq!(fc.frequency(1), 2); }
#[test]
fn test_frequency_counter_decay_removes_zero() {
let mut fc = FrequencyCounter::new(100);
fc.increment(77); fc.decay_all(); assert_eq!(fc.frequency(77), 0);
assert_eq!(fc.tracked_keys(), 0);
}
#[test]
fn test_frequency_counter_auto_decay_on_window_fill() {
let window = 8usize;
let mut fc = FrequencyCounter::new(window);
for _ in 0..window {
fc.increment(1);
}
assert_eq!(fc.frequency(1), 4);
}
#[test]
fn test_frequency_counter_clear() {
let mut fc = FrequencyCounter::new(50);
fc.increment(10);
fc.increment(20);
fc.clear();
assert_eq!(fc.tracked_keys(), 0);
assert_eq!(fc.frequency(10), 0);
}
#[test]
fn test_lfu_insert_and_frequency() {
let mut tracker = LfuEvictionTracker::new();
tracker.insert(100);
assert_eq!(tracker.frequency(100), 1);
}
#[test]
fn test_lfu_promote() {
let mut tracker = LfuEvictionTracker::new();
tracker.insert(1);
tracker.insert(2);
tracker.promote(1);
let victim = tracker.evict();
assert_eq!(victim, Some(2));
}
#[test]
fn test_lfu_evict_lowest_frequency() {
let mut tracker = LfuEvictionTracker::new();
tracker.insert(10);
tracker.insert(20);
tracker.insert(30);
tracker.promote(10);
tracker.promote(20);
let victim = tracker.evict();
assert_eq!(victim, Some(30), "key 30 has lowest frequency");
}
#[test]
fn test_lfu_evict_fifo_within_bucket() {
let mut tracker = LfuEvictionTracker::new();
tracker.insert(1);
tracker.insert(2);
tracker.insert(3);
assert_eq!(tracker.evict(), Some(1));
assert_eq!(tracker.evict(), Some(2));
assert_eq!(tracker.evict(), Some(3));
}
#[test]
fn test_lfu_evict_empty() {
let mut tracker = LfuEvictionTracker::new();
assert!(tracker.evict().is_none());
}
#[test]
fn test_lfu_remove() {
let mut tracker = LfuEvictionTracker::new();
tracker.insert(55);
assert!(tracker.remove(55));
assert_eq!(tracker.frequency(55), 0);
assert!(tracker.is_empty());
}
#[test]
fn test_lfu_remove_absent() {
let mut tracker = LfuEvictionTracker::new();
assert!(!tracker.remove(999));
}
#[test]
fn test_lfu_len_and_is_empty() {
let mut tracker = LfuEvictionTracker::new();
assert!(tracker.is_empty());
tracker.insert(1);
tracker.insert(2);
assert_eq!(tracker.len(), 2);
tracker.evict();
assert_eq!(tracker.len(), 1);
}
#[test]
fn test_lfu_insert_duplicate_no_reset() {
let mut tracker = LfuEvictionTracker::new();
tracker.insert(7);
tracker.promote(7);
tracker.promote(7);
tracker.insert(7);
assert_eq!(tracker.frequency(7), 3);
}
#[test]
fn test_tinylfu_should_admit_popular() {
let mut gate = TinyLfuAdmission::new(100);
let key = 42u64;
for _ in 0..20 {
gate.record_access(key);
}
assert!(gate.should_admit(key, 1));
}
#[test]
fn test_tinylfu_should_not_admit_cold() {
let mut gate = TinyLfuAdmission::new(100);
let candidate = 99u64;
let freq = gate.estimated_frequency(candidate);
assert_eq!(freq, 0);
assert!(!gate.should_admit(candidate, 10));
}
#[test]
fn test_tinylfu_decay() {
let mut gate = TinyLfuAdmission::new(100);
for _ in 0..8 {
gate.record_access(1);
}
let before = gate.estimated_frequency(1);
gate.decay();
let after = gate.estimated_frequency(1);
assert!(after <= before, "decay should not increase frequency");
}
#[test]
fn test_arc_initial_state() {
let arc = ArcTracker::new(100);
assert_eq!(arc.t1_size, 0);
assert_eq!(arc.t2_size, 0);
assert_eq!(arc.b1_size, 0);
assert_eq!(arc.b2_size, 0);
assert_eq!(arc.p, 0);
assert_eq!(arc.capacity(), 100);
}
#[test]
fn test_arc_admit_t1() {
let mut arc = ArcTracker::new(10);
arc.on_admit_t1();
arc.on_admit_t1();
assert_eq!(arc.t1_size, 2);
assert_eq!(arc.live_size(), 2);
}
#[test]
fn test_arc_promote_t1_to_t2() {
let mut arc = ArcTracker::new(10);
arc.on_admit_t1();
arc.on_promote_t1_to_t2();
assert_eq!(arc.t1_size, 0);
assert_eq!(arc.t2_size, 1);
}
#[test]
fn test_arc_evict_t1_ghost() {
let mut arc = ArcTracker::new(10);
arc.on_admit_t1();
arc.on_evict_t1();
assert_eq!(arc.t1_size, 0);
assert_eq!(arc.b1_size, 1);
assert_eq!(arc.ghost_size(), 1);
}
#[test]
fn test_arc_evict_t2_ghost() {
let mut arc = ArcTracker::new(10);
arc.on_admit_t1();
arc.on_promote_t1_to_t2();
arc.on_evict_t2();
assert_eq!(arc.t2_size, 0);
assert_eq!(arc.b2_size, 1);
}
#[test]
fn test_arc_adapt_on_hit_b1_increases_p() {
let mut arc = ArcTracker::new(100);
arc.b1_size = 5;
arc.b2_size = 10;
let p_before = arc.p;
arc.adapt_on_hit_b1();
assert!(arc.p > p_before, "p should increase on B1 hit");
}
#[test]
fn test_arc_adapt_on_hit_b2_decreases_p() {
let mut arc = ArcTracker::new(100);
arc.p = 50;
arc.b1_size = 10;
arc.b2_size = 5;
arc.adapt_on_hit_b2();
assert!(arc.p < 50, "p should decrease on B2 hit");
}
#[test]
fn test_arc_p_capped_at_capacity() {
let mut arc = ArcTracker::new(10);
arc.p = 9;
arc.b1_size = 1;
arc.b2_size = 100;
arc.adapt_on_hit_b1();
assert!(arc.p <= 10, "p must not exceed capacity");
}
#[test]
fn test_arc_p_floor_at_zero() {
let mut arc = ArcTracker::new(10);
arc.p = 0;
arc.b2_size = 5;
arc.b1_size = 1;
arc.adapt_on_hit_b2();
assert_eq!(arc.p, 0, "p must not go below zero");
}
#[test]
fn test_arc_is_full() {
let mut arc = ArcTracker::new(3);
arc.on_admit_t1();
arc.on_admit_t1();
assert!(!arc.is_full());
arc.on_admit_t1();
assert!(arc.is_full());
}
#[test]
fn test_arc_should_evict_t1_when_over_target() {
let mut arc = ArcTracker::new(10);
arc.t1_size = 5;
arc.t2_size = 3;
arc.p = 2; assert!(arc.should_evict_t1());
}
#[test]
fn test_arc_should_evict_t2_when_t1_at_target() {
let mut arc = ArcTracker::new(10);
arc.t1_size = 2;
arc.t2_size = 4;
arc.p = 3; assert!(!arc.should_evict_t1());
}
#[test]
fn test_arc_remove_b1_ghost() {
let mut arc = ArcTracker::new(10);
arc.b1_size = 3;
arc.on_remove_b1_ghost();
assert_eq!(arc.b1_size, 2);
}
#[test]
fn test_arc_remove_b2_ghost() {
let mut arc = ArcTracker::new(10);
arc.b2_size = 5;
arc.on_remove_b2_ghost();
assert_eq!(arc.b2_size, 4);
}
#[test]
fn test_eviction_policy_equality() {
assert_eq!(EvictionPolicy::Lru, EvictionPolicy::Lru);
assert_ne!(EvictionPolicy::Lru, EvictionPolicy::Lfu);
assert_ne!(EvictionPolicy::TinyLfu, EvictionPolicy::ArcCache);
}
#[test]
fn test_eviction_policy_clone() {
let p = EvictionPolicy::ArcCache;
let q = p.clone();
assert_eq!(p, q);
}
}