use core::hash::Hash;
use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::sync::{Mutex, MutexGuard};
use crate::cache::Cache;
use crate::error::CacheError;
pub struct LfuCache<K, V> {
capacity: NonZeroUsize,
inner: Mutex<Inner<K, V>>,
}
struct Entry<V> {
value: V,
count: u64,
last_access: u64,
}
struct Inner<K, V> {
map: HashMap<K, Entry<V>>,
clock: u64,
}
impl<K, V> LfuCache<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
pub fn new(capacity: usize) -> Result<Self, CacheError> {
let cap = NonZeroUsize::new(capacity).ok_or(CacheError::InvalidCapacity)?;
Ok(Self::with_capacity(cap))
}
pub fn with_capacity(capacity: NonZeroUsize) -> Self {
let cap = capacity.get();
Self {
capacity,
inner: Mutex::new(Inner {
map: HashMap::with_capacity(cap),
clock: 0,
}),
}
}
fn lock_inner(&self) -> MutexGuard<'_, Inner<K, V>> {
match self.inner.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
}
}
}
impl<K, V> Cache<K, V> for LfuCache<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
fn get(&self, key: &K) -> Option<V> {
let mut inner = self.lock_inner();
inner.clock = inner.clock.wrapping_add(1);
let now = inner.clock;
let entry = inner.map.get_mut(key)?;
entry.count = entry.count.saturating_add(1);
entry.last_access = now;
Some(entry.value.clone())
}
fn insert(&self, key: K, value: V) -> Option<V> {
let mut inner = self.lock_inner();
inner.clock = inner.clock.wrapping_add(1);
let now = inner.clock;
if let Some(existing) = inner.map.get_mut(&key) {
let old = core::mem::replace(&mut existing.value, value);
existing.count = existing.count.saturating_add(1);
existing.last_access = now;
return Some(old);
}
if inner.map.len() >= self.capacity.get() {
if let Some(victim) = find_victim(&inner.map) {
let _ = inner.map.remove(&victim);
}
}
let _ = inner.map.insert(
key,
Entry {
value,
count: 1,
last_access: now,
},
);
None
}
fn remove(&self, key: &K) -> Option<V> {
let mut inner = self.lock_inner();
inner.map.remove(key).map(|e| e.value)
}
fn contains_key(&self, key: &K) -> bool {
self.lock_inner().map.contains_key(key)
}
fn len(&self) -> usize {
self.lock_inner().map.len()
}
fn clear(&self) {
let mut inner = self.lock_inner();
inner.map.clear();
inner.clock = 0;
}
fn capacity(&self) -> usize {
self.capacity.get()
}
}
fn find_victim<K, V>(map: &HashMap<K, Entry<V>>) -> Option<K>
where
K: Clone,
{
map.iter()
.min_by(|(_, a), (_, b)| {
a.count
.cmp(&b.count)
.then(a.last_access.cmp(&b.last_access))
})
.map(|(k, _)| k.clone())
}