use crate::store::hashmap::HashMapStore;
use crate::store::traits::{StoreCore, StoreMut};
use crate::traits::{Cache, EvictingCache, FrequencyTracking, VictimInspectable};
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::hash::Hash;
use std::sync::Arc;
#[cfg(feature = "metrics")]
use crate::metrics::metrics_impl::LfuMetrics;
#[cfg(feature = "metrics")]
use crate::metrics::snapshot::LfuMetricsSnapshot;
#[cfg(feature = "metrics")]
use crate::metrics::traits::{
CoreMetricsRecorder, LfuMetricsReadRecorder, LfuMetricsRecorder, MetricsSnapshotProvider,
};
#[derive(Debug)]
pub struct HeapLfuCache<K, V> {
store: HashMapStore<K, Arc<V>>,
frequencies: HashMap<K, u64>,
freq_heap: BinaryHeap<Reverse<(u64, K)>>,
#[cfg(feature = "metrics")]
metrics: LfuMetrics,
}
impl<K, V> Clone for HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
V: Clone,
{
fn clone(&self) -> Self {
Self {
store: self.store.clone(),
frequencies: self.frequencies.clone(),
freq_heap: self.freq_heap.clone(),
#[cfg(feature = "metrics")]
metrics: self.metrics.clone(),
}
}
}
impl<K, V> HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
const MAX_HEAP_FACTOR: usize = 4;
#[must_use]
pub fn new(capacity: usize) -> Self {
HeapLfuCache {
store: HashMapStore::new(capacity),
frequencies: HashMap::with_capacity(capacity),
freq_heap: BinaryHeap::with_capacity(capacity),
#[cfg(feature = "metrics")]
metrics: LfuMetrics::default(),
}
}
#[must_use]
pub fn capacity(&self) -> usize {
self.store.capacity()
}
#[must_use]
pub fn len(&self) -> usize {
self.store.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.store.is_empty()
}
#[must_use]
pub fn contains(&self, key: &K) -> bool {
self.store.contains(key)
}
#[must_use]
pub fn frequency(&self, key: &K) -> Option<u64> {
#[cfg(feature = "metrics")]
(&self.metrics).record_frequency_call();
let result = self.frequencies.get(key).copied();
#[cfg(feature = "metrics")]
if result.is_some() {
(&self.metrics).record_frequency_found();
}
result
}
pub fn clear(&mut self) {
#[cfg(feature = "metrics")]
self.metrics.record_clear();
self.store.clear();
self.frequencies.clear();
self.freq_heap.clear();
}
fn add_to_heap(&mut self, key: &K, frequency: u64) {
self.freq_heap.push(Reverse((frequency, key.clone())));
self.maybe_rebuild_heap();
}
fn maybe_rebuild_heap(&mut self) {
let live_entries = self.store.len().max(1);
let max_heap_len = live_entries.saturating_mul(Self::MAX_HEAP_FACTOR);
if self.freq_heap.len() <= max_heap_len {
return;
}
self.freq_heap.clear();
self.freq_heap.reserve(self.frequencies.len());
for (key, freq) in &self.frequencies {
self.freq_heap.push(Reverse((*freq, key.clone())));
}
}
fn pop_lfu_internal(&mut self) -> Option<(K, u64)> {
let mut stale_pops = 0usize;
while let Some(Reverse((heap_freq, key))) = self.freq_heap.peek() {
if let Some(¤t_freq) = self.frequencies.get(key) {
if *heap_freq == current_freq {
let Reverse((freq, key)) = self.freq_heap.pop().unwrap();
return Some((key, freq));
}
}
self.freq_heap.pop();
stale_pops += 1;
if stale_pops >= self.store.len().max(1) {
self.maybe_rebuild_heap();
stale_pops = 0;
}
}
None
}
fn ensure_capacity(&mut self) -> Option<(K, Arc<V>)> {
if self.store.len() >= self.store.capacity() {
self.pop_lfu()
} else {
None
}
}
pub fn pop_lfu(&mut self) -> Option<(K, Arc<V>)> {
#[cfg(feature = "metrics")]
self.metrics.record_pop_lfu_call();
let (lfu_key, _freq) = self.pop_lfu_internal()?;
let value = self.store.remove(&lfu_key)?;
self.frequencies.remove(&lfu_key);
self.store.record_eviction();
#[cfg(feature = "metrics")]
self.metrics.record_pop_lfu_found();
Some((lfu_key, value))
}
pub fn peek_lfu(&self) -> Option<(&K, &Arc<V>)> {
#[cfg(feature = "metrics")]
(&self.metrics).record_peek_lfu_call();
if self.frequencies.is_empty() {
return None;
}
let min_freq = *self.frequencies.values().min()?;
for (key, &freq) in &self.frequencies {
if freq == min_freq {
let result = self.store.peek(key).map(|value| (key, value));
#[cfg(feature = "metrics")]
if result.is_some() {
(&self.metrics).record_peek_lfu_found();
}
return result;
}
}
None
}
pub fn increment_frequency(&mut self, key: &K) -> Option<u64> {
if let Some(freq) = self.frequencies.get_mut(key) {
*freq += 1;
let new_freq = *freq;
self.add_to_heap(key, new_freq);
Some(new_freq)
} else {
None
}
}
pub fn reset_frequency(&mut self, key: &K) -> Option<u64> {
if let Some(freq) = self.frequencies.get_mut(key) {
let old_freq = *freq;
*freq = 1;
self.add_to_heap(key, 1);
Some(old_freq)
} else {
None
}
}
}
impl<K, V> Cache<K, Arc<V>> for HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
fn contains(&self, key: &K) -> bool {
self.store.contains(key)
}
fn len(&self) -> usize {
self.store.len()
}
fn capacity(&self) -> usize {
self.store.capacity()
}
fn peek(&self, key: &K) -> Option<&Arc<V>> {
self.store.peek(key)
}
fn insert(&mut self, key: K, value: Arc<V>) -> Option<Arc<V>> {
#[cfg(feature = "metrics")]
self.metrics.record_insert_call();
if self.store.contains(&key) {
#[cfg(feature = "metrics")]
self.metrics.record_insert_update();
return self.store.try_insert(key, value).ok().flatten();
}
#[cfg(feature = "metrics")]
if self.store.len() >= self.store.capacity() {
self.metrics.record_evict_call();
}
let _evicted = self.ensure_capacity();
#[cfg(feature = "metrics")]
if _evicted.is_some() {
self.metrics.record_evicted_entry();
}
#[cfg(feature = "metrics")]
self.metrics.record_insert_new();
if self.store.try_insert(key.clone(), value).is_err() {
return None;
}
self.frequencies.insert(key.clone(), 1);
self.add_to_heap(&key, 1);
None
}
fn get(&mut self, key: &K) -> Option<&Arc<V>> {
if self.store.contains(key) {
#[cfg(feature = "metrics")]
self.metrics.record_get_hit();
let new_freq = self.frequencies.get_mut(key).map(|f| {
*f += 1;
*f
})?;
self.add_to_heap(key, new_freq);
self.store.get(key)
} else {
#[cfg(feature = "metrics")]
self.metrics.record_get_miss();
None
}
}
fn remove(&mut self, key: &K) -> Option<Arc<V>> {
let value = self.store.remove(key);
let had_frequency = self.frequencies.remove(key).is_some();
if value.is_some() || had_frequency {
self.maybe_rebuild_heap();
}
value
}
fn clear(&mut self) {
HeapLfuCache::clear(self);
}
}
impl<K, V> EvictingCache<K, Arc<V>> for HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
fn evict_one(&mut self) -> Option<(K, Arc<V>)> {
self.pop_lfu()
}
}
impl<K, V> VictimInspectable<K, Arc<V>> for HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
fn peek_victim(&self) -> Option<(&K, &Arc<V>)> {
self.peek_lfu()
}
}
impl<K, V> FrequencyTracking<K, Arc<V>> for HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
fn frequency(&self, key: &K) -> Option<u64> {
HeapLfuCache::frequency(self, key)
}
}
#[cfg(feature = "metrics")]
impl<K, V> HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
#[must_use]
pub fn metrics_snapshot(&self) -> LfuMetricsSnapshot {
LfuMetricsSnapshot {
get_calls: self.metrics.get_calls,
get_hits: self.metrics.get_hits,
get_misses: self.metrics.get_misses,
insert_calls: self.metrics.insert_calls,
insert_updates: self.metrics.insert_updates,
insert_new: self.metrics.insert_new,
evict_calls: self.metrics.evict_calls,
evicted_entries: self.metrics.evicted_entries,
pop_lfu_calls: self.metrics.pop_lfu_calls,
pop_lfu_found: self.metrics.pop_lfu_found,
peek_lfu_calls: self.metrics.peek_lfu_calls.get(),
peek_lfu_found: self.metrics.peek_lfu_found.get(),
frequency_calls: self.metrics.frequency_calls.get(),
frequency_found: self.metrics.frequency_found.get(),
reset_frequency_calls: self.metrics.reset_frequency_calls,
reset_frequency_found: self.metrics.reset_frequency_found,
increment_frequency_calls: self.metrics.increment_frequency_calls,
increment_frequency_found: self.metrics.increment_frequency_found,
cache_len: self.store.len(),
capacity: self.store.capacity(),
}
}
}
#[cfg(feature = "metrics")]
impl<K, V> MetricsSnapshotProvider<LfuMetricsSnapshot> for HeapLfuCache<K, V>
where
K: Eq + Hash + Clone + Ord,
{
fn snapshot(&self) -> LfuMetricsSnapshot {
self.metrics_snapshot()
}
}
#[cfg(test)]
mod heap_lfu_tests {
use super::*;
use crate::policy::lfu::LfuCache;
#[test]
fn test_heap_lfu_basic_operations() {
let mut cache: HeapLfuCache<String, i32> = HeapLfuCache::new(3);
assert_eq!(cache.insert("key1".to_string(), Arc::new(100)), None);
assert_eq!(cache.insert("key2".to_string(), Arc::new(200)), None);
assert_eq!(cache.insert("key3".to_string(), Arc::new(300)), None);
assert_eq!(cache.len(), 3);
assert_eq!(cache.capacity(), 3);
assert_eq!(cache.get(&"key1".to_string()).map(Arc::as_ref), Some(&100));
assert_eq!(cache.frequency(&"key1".to_string()), Some(2));
assert_eq!(cache.get(&"key2".to_string()).map(Arc::as_ref), Some(&200));
assert_eq!(cache.get(&"key2".to_string()).map(Arc::as_ref), Some(&200)); assert_eq!(cache.frequency(&"key2".to_string()), Some(3));
assert!(cache.contains(&"key1".to_string()));
assert!(!cache.contains(&"nonexistent".to_string()));
}
#[test]
fn test_heap_lfu_eviction_order() {
let mut cache: HeapLfuCache<String, i32> = HeapLfuCache::new(3);
cache.insert("key1".to_string(), Arc::new(100));
cache.insert("key2".to_string(), Arc::new(200));
cache.insert("key3".to_string(), Arc::new(300));
cache.get(&"key2".to_string()); cache.get(&"key2".to_string()); cache.get(&"key3".to_string());
assert_eq!(cache.frequency(&"key1".to_string()), Some(1)); assert_eq!(cache.frequency(&"key2".to_string()), Some(3)); assert_eq!(cache.frequency(&"key3".to_string()), Some(2));
cache.insert("key4".to_string(), Arc::new(400));
assert!(!cache.contains(&"key1".to_string()));
assert_eq!(cache.get(&"key1".to_string()).map(Arc::as_ref), None);
assert!(cache.contains(&"key2".to_string()));
assert!(cache.contains(&"key3".to_string()));
assert!(cache.contains(&"key4".to_string()));
assert_eq!(cache.len(), 3);
}
#[test]
fn test_heap_lfu_pop_lfu() {
let mut cache: HeapLfuCache<String, i32> = HeapLfuCache::new(3);
cache.insert("low".to_string(), Arc::new(1));
cache.insert("med".to_string(), Arc::new(2));
cache.insert("high".to_string(), Arc::new(3));
cache.get(&"med".to_string()); cache.get(&"high".to_string()); cache.get(&"high".to_string());
assert_eq!(cache.frequency(&"low".to_string()), Some(1));
assert_eq!(cache.frequency(&"med".to_string()), Some(2));
assert_eq!(cache.frequency(&"high".to_string()), Some(3));
let (key, value) = cache.pop_lfu().unwrap();
assert_eq!(key, "low".to_string());
assert_eq!(*value, 1);
assert_eq!(cache.len(), 2);
let (key, value) = cache.pop_lfu().unwrap();
assert_eq!(key, "med".to_string());
assert_eq!(*value, 2);
assert_eq!(cache.len(), 1);
let (key, value) = cache.pop_lfu().unwrap();
assert_eq!(key, "high".to_string());
assert_eq!(*value, 3);
assert_eq!(cache.len(), 0);
assert_eq!(cache.pop_lfu(), None);
}
#[test]
fn test_heap_lfu_stale_entry_handling() {
let mut cache: HeapLfuCache<i32, i32> = HeapLfuCache::new(3);
cache.insert(1, Arc::new(10));
cache.insert(2, Arc::new(20));
cache.insert(3, Arc::new(30));
cache.get(&1); cache.get(&1); cache.get(&2);
cache.remove(&1);
cache.insert(4, Arc::new(40));
assert!(!cache.contains(&1));
assert!(cache.contains(&2));
assert!(cache.contains(&3));
assert!(cache.contains(&4));
assert_eq!(cache.len(), 3);
let (key, _) = cache.pop_lfu().unwrap();
assert!(key == 3 || key == 4); }
#[test]
fn test_remove_clears_stale_frequency_entries() {
let mut cache: HeapLfuCache<String, i32> = HeapLfuCache::new(2);
cache.insert("key1".to_string(), Arc::new(10));
cache.insert("key2".to_string(), Arc::new(20));
for _ in 0..10 {
cache.increment_frequency(&"key1".to_string());
}
let _ = cache.store.remove(&"key1".to_string());
assert!(cache.frequency(&"key1".to_string()).is_some());
cache.remove(&"key1".to_string());
assert!(cache.frequency(&"key1".to_string()).is_none());
let has_key1_in_heap = cache
.freq_heap
.iter()
.any(|Reverse((_, key))| key == "key1");
assert!(!has_key1_in_heap);
}
#[test]
fn test_pop_lfu_internal_rebuilds_after_stale_pops() {
let mut cache: HeapLfuCache<String, i32> = HeapLfuCache::new(2);
cache.insert("key1".to_string(), Arc::new(10));
cache.insert("key2".to_string(), Arc::new(20));
for _ in 0..10 {
cache.increment_frequency(&"key1".to_string());
}
for _ in 0..4 {
cache.increment_frequency(&"key2".to_string());
}
let heap_len_before = cache.freq_heap.len();
let (lfu_key, lfu_freq) = cache.pop_lfu_internal().unwrap();
assert_eq!(lfu_key, "key2".to_string());
assert_eq!(lfu_freq, 5);
assert!(cache.freq_heap.len() < heap_len_before);
assert_eq!(cache.freq_heap.len(), 1);
}
#[test]
#[cfg_attr(miri, ignore)]
fn test_heap_lfu_performance_comparison() {
use std::time::Instant;
let cache_size = 100;
let mut std_cache = LfuCache::new(cache_size);
for i in 0..cache_size {
std_cache.insert(i, Arc::new(i * 10));
}
let start = Instant::now();
for _ in 0..10 {
if let Some((key, value)) = std_cache.pop_lfu() {
std_cache.insert(key + cache_size, value); }
}
let std_duration = start.elapsed();
let mut heap_cache: HeapLfuCache<usize, usize> = HeapLfuCache::new(cache_size);
for i in 0..cache_size {
heap_cache.insert(i, Arc::new(i * 10));
}
let start = Instant::now();
for _ in 0..10 {
if let Some((key, value)) = heap_cache.pop_lfu() {
heap_cache.insert(key + cache_size, value); }
}
let heap_duration = start.elapsed();
println!("Performance Comparison:");
println!(" Standard LFU (O(n)): {:?}", std_duration);
println!(" Heap LFU (O(log n)): {:?}", heap_duration);
assert_eq!(std_cache.len(), cache_size);
assert_eq!(heap_cache.len(), cache_size);
}
}