extern crate alloc;
use crate::config::LfuCacheConfig;
use crate::entry::CacheEntry;
use crate::list::{List, ListEntry};
use crate::metrics::{CacheMetrics, LfuCacheMetrics};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct LfuMeta {
pub frequency: u64,
}
impl LfuMeta {
#[inline]
pub fn new(frequency: u64) -> Self {
Self { frequency }
}
#[inline]
pub fn increment(&mut self) -> u64 {
self.frequency += 1;
self.frequency
}
}
use alloc::boxed::Box;
use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;
use core::borrow::Borrow;
use core::hash::{BuildHasher, Hash};
use core::num::NonZeroUsize;
#[cfg(feature = "hashbrown")]
use hashbrown::DefaultHashBuilder;
#[cfg(feature = "hashbrown")]
use hashbrown::HashMap;
#[cfg(not(feature = "hashbrown"))]
use std::collections::hash_map::RandomState as DefaultHashBuilder;
#[cfg(not(feature = "hashbrown"))]
use std::collections::HashMap;
pub(crate) struct LfuSegment<K, V, S = DefaultHashBuilder> {
config: LfuCacheConfig,
min_frequency: usize,
map: HashMap<K, *mut ListEntry<CacheEntry<K, V, LfuMeta>>, S>,
frequency_lists: BTreeMap<usize, List<CacheEntry<K, V, LfuMeta>>>,
metrics: LfuCacheMetrics,
current_size: u64,
}
unsafe impl<K: Send, V: Send, S: Send> Send for LfuSegment<K, V, S> {}
unsafe impl<K: Send, V: Send, S: Sync> Sync for LfuSegment<K, V, S> {}
impl<K: Hash + Eq, V: Clone, S: BuildHasher> LfuSegment<K, V, S> {
#[allow(dead_code)] pub(crate) fn init(config: LfuCacheConfig, hasher: S) -> Self {
let map_capacity = config.capacity.get().next_power_of_two();
LfuSegment {
config,
min_frequency: 1,
map: HashMap::with_capacity_and_hasher(map_capacity, hasher),
frequency_lists: BTreeMap::new(),
metrics: LfuCacheMetrics::new(config.max_size),
current_size: 0,
}
}
#[inline]
pub(crate) fn cap(&self) -> NonZeroUsize {
self.config.capacity
}
#[inline]
pub(crate) fn len(&self) -> usize {
self.map.len()
}
#[inline]
pub(crate) fn is_empty(&self) -> bool {
self.map.is_empty()
}
#[inline]
pub(crate) fn current_size(&self) -> u64 {
self.current_size
}
#[inline]
pub(crate) fn max_size(&self) -> u64 {
self.config.max_size
}
#[inline]
pub(crate) fn metrics(&self) -> &LfuCacheMetrics {
&self.metrics
}
unsafe fn update_frequency_by_node(
&mut self,
node: *mut ListEntry<CacheEntry<K, V, LfuMeta>>,
old_frequency: usize,
) -> *mut ListEntry<CacheEntry<K, V, LfuMeta>>
where
K: Clone + Hash + Eq,
{
let new_frequency = old_frequency + 1;
self.metrics
.record_frequency_increment(old_frequency, new_frequency);
let entry = (*node).get_value();
let key_cloned = entry.key.clone();
let node = *self.map.get(&key_cloned).unwrap();
let boxed_entry = self
.frequency_lists
.get_mut(&old_frequency)
.unwrap()
.remove(node)
.unwrap();
if self.frequency_lists.get(&old_frequency).unwrap().is_empty() {
self.frequency_lists.remove(&old_frequency);
if old_frequency == self.min_frequency {
self.min_frequency = new_frequency;
}
}
let entry_ptr = Box::into_raw(boxed_entry);
(*entry_ptr).get_value_mut().metadata.algorithm.frequency = new_frequency as u64;
let capacity = self.config.capacity;
self.frequency_lists
.entry(new_frequency)
.or_insert_with(|| List::new(capacity));
self.frequency_lists
.get_mut(&new_frequency)
.unwrap()
.attach_from_other_list(entry_ptr);
*self.map.get_mut(&key_cloned).unwrap() = entry_ptr;
self.metrics.update_frequency_levels(&self.frequency_lists);
entry_ptr
}
pub(crate) fn get<Q>(&mut self, key: &Q) -> Option<&V>
where
K: Borrow<Q> + Clone,
Q: ?Sized + Hash + Eq,
{
if let Some(&node) = self.map.get(key) {
unsafe {
let entry = (*node).get_value();
let frequency = entry.metadata.algorithm.frequency as usize;
let object_size = entry.metadata.size;
self.metrics.record_frequency_hit(object_size, frequency);
let new_node = self.update_frequency_by_node(node, frequency);
let new_entry = (*new_node).get_value();
Some(&new_entry.value)
}
} else {
None
}
}
pub(crate) fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
where
K: Borrow<Q> + Clone,
Q: ?Sized + Hash + Eq,
{
if let Some(&node) = self.map.get(key) {
unsafe {
let entry = (*node).get_value();
let frequency = entry.metadata.algorithm.frequency as usize;
let object_size = entry.metadata.size;
self.metrics.record_frequency_hit(object_size, frequency);
let new_node = self.update_frequency_by_node(node, frequency);
let new_entry = (*new_node).get_value_mut();
Some(&mut new_entry.value)
}
} else {
None
}
}
pub(crate) fn put(&mut self, key: K, value: V, size: u64) -> Option<Vec<(K, V)>>
where
K: Clone,
{
if let Some(&node) = self.map.get(&key) {
unsafe {
let entry = (*node).get_value();
let frequency = entry.metadata.algorithm.frequency as usize;
let old_size = entry.metadata.size;
let new_entry = CacheEntry::with_algorithm_metadata(
key.clone(),
value,
size,
LfuMeta::new(frequency as u64),
);
let _old_entry = self
.frequency_lists
.get_mut(&frequency)
.unwrap()
.update(node, new_entry, true);
self.current_size = self.current_size.saturating_sub(old_size);
self.current_size += size;
self.metrics.core.record_size_change(old_size, size);
self.metrics.core.bytes_written_to_cache += size;
return None;
}
}
let mut evicted = Vec::new();
while self.len() >= self.config.capacity.get()
|| (self.current_size + size > self.config.max_size && !self.map.is_empty())
{
if let Some(entry) = self.evict() {
self.metrics.core.evictions += 1;
evicted.push(entry);
} else {
break;
}
}
let frequency = 1;
self.min_frequency = 1;
let capacity = self.config.capacity;
self.frequency_lists
.entry(frequency)
.or_insert_with(|| List::new(capacity));
let cache_entry = CacheEntry::with_algorithm_metadata(
key.clone(),
value,
size,
LfuMeta::new(frequency as u64),
);
if let Some(node) = self
.frequency_lists
.get_mut(&frequency)
.unwrap()
.add(cache_entry)
{
self.map.insert(key, node);
self.current_size += size;
}
self.metrics.core.record_insertion(size);
self.metrics.update_frequency_levels(&self.frequency_lists);
if evicted.is_empty() {
None
} else {
Some(evicted)
}
}
pub(crate) fn remove<Q>(&mut self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
let node = self.map.remove(key)?;
unsafe {
let frequency = (*node).get_value().metadata.algorithm.frequency as usize;
let boxed_entry = self.frequency_lists.get_mut(&frequency)?.remove(node)?;
let entry_ptr = Box::into_raw(boxed_entry);
let cache_entry = (*entry_ptr).take_value();
let removed_size = cache_entry.metadata.size;
let _ = Box::from_raw(entry_ptr);
self.current_size = self.current_size.saturating_sub(removed_size);
self.metrics.core.record_removal(removed_size);
if self.frequency_lists.get(&frequency).unwrap().is_empty() {
self.frequency_lists.remove(&frequency);
if frequency == self.min_frequency {
self.min_frequency = self.frequency_lists.keys().copied().next().unwrap_or(1);
}
}
Some(cache_entry.value)
}
}
pub(crate) fn clear(&mut self) {
self.map.clear();
self.frequency_lists.clear();
self.min_frequency = 1;
self.current_size = 0;
}
#[inline]
pub(crate) fn record_miss(&mut self, object_size: u64) {
self.metrics.record_miss(object_size);
}
#[inline]
pub(crate) fn contains<Q>(&self, key: &Q) -> bool
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.map.contains_key(key)
}
pub(crate) fn peek<Q>(&self, key: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
let &node = self.map.get(key)?;
unsafe {
let entry = (*node).get_value();
Some(&entry.value)
}
}
fn evict(&mut self) -> Option<(K, V)> {
if self.is_empty() {
return None;
}
let min_freq_list = self.frequency_lists.get_mut(&self.min_frequency)?;
let old_entry = min_freq_list.remove_last()?;
let is_list_empty = min_freq_list.is_empty();
unsafe {
let entry_ptr = Box::into_raw(old_entry);
let cache_entry = (*entry_ptr).take_value();
let evicted_size = cache_entry.metadata.size;
self.map.remove(&cache_entry.key);
self.current_size = self.current_size.saturating_sub(evicted_size);
self.metrics.core.record_removal(evicted_size);
if is_list_empty {
self.frequency_lists.remove(&self.min_frequency);
self.min_frequency = self.frequency_lists.keys().copied().next().unwrap_or(1);
}
let _ = Box::from_raw(entry_ptr);
Some((cache_entry.key, cache_entry.value))
}
}
}
impl<K, V, S> core::fmt::Debug for LfuSegment<K, V, S> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("LfuSegment")
.field("capacity", &self.config.capacity)
.field("len", &self.map.len())
.field("min_frequency", &self.min_frequency)
.finish()
}
}
#[derive(Debug)]
pub struct LfuCache<K, V, S = DefaultHashBuilder> {
segment: LfuSegment<K, V, S>,
}
impl<K: Hash + Eq, V: Clone, S: BuildHasher> LfuCache<K, V, S> {
#[inline]
pub fn cap(&self) -> NonZeroUsize {
self.segment.cap()
}
#[inline]
pub fn len(&self) -> usize {
self.segment.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.segment.is_empty()
}
#[inline]
pub fn current_size(&self) -> u64 {
self.segment.current_size()
}
#[inline]
pub fn max_size(&self) -> u64 {
self.segment.max_size()
}
#[inline]
pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
where
K: Borrow<Q> + Clone,
Q: ?Sized + Hash + Eq,
{
self.segment.get(key)
}
#[inline]
pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
where
K: Borrow<Q> + Clone,
Q: ?Sized + Hash + Eq,
{
self.segment.get_mut(key)
}
#[inline]
pub fn put(&mut self, key: K, value: V, size: u64) -> Option<Vec<(K, V)>>
where
K: Clone,
{
self.segment.put(key, value, size)
}
#[inline]
pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
V: Clone,
{
self.segment.remove(key)
}
#[inline]
pub fn clear(&mut self) {
self.segment.clear()
}
#[inline]
pub fn record_miss(&mut self, object_size: u64) {
self.segment.record_miss(object_size);
}
#[inline]
pub fn contains<Q>(&self, key: &Q) -> bool
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.segment.contains(key)
}
#[inline]
pub fn peek<Q>(&self, key: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: ?Sized + Hash + Eq,
{
self.segment.peek(key)
}
}
impl<K: Hash + Eq, V> LfuCache<K, V>
where
V: Clone,
{
pub fn init(
config: LfuCacheConfig,
hasher: Option<DefaultHashBuilder>,
) -> LfuCache<K, V, DefaultHashBuilder> {
LfuCache {
segment: LfuSegment::init(config, hasher.unwrap_or_default()),
}
}
}
impl<K: Hash + Eq, V: Clone, S: BuildHasher> CacheMetrics for LfuCache<K, V, S> {
fn metrics(&self) -> BTreeMap<String, f64> {
self.segment.metrics().metrics()
}
fn algorithm_name(&self) -> &'static str {
self.segment.metrics().algorithm_name()
}
}
#[cfg(test)]
mod tests {
extern crate std;
use alloc::format;
use alloc::vec;
use std::println;
use std::string::ToString;
use super::*;
use crate::config::LfuCacheConfig;
use alloc::string::String;
fn make_cache<K: Hash + Eq + Clone, V: Clone>(cap: usize) -> LfuCache<K, V> {
let config = LfuCacheConfig {
capacity: NonZeroUsize::new(cap).unwrap(),
max_size: u64::MAX,
};
LfuCache::init(config, None)
}
#[test]
fn test_lfu_basic() {
let mut cache = make_cache(3);
assert_eq!(cache.put("a", 1, 1), None);
assert_eq!(cache.put("b", 2, 1), None);
assert_eq!(cache.put("c", 3, 1), None);
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"b"), Some(&2));
let evicted = cache.put("d", 4, 1);
assert!(evicted.is_some());
let (evicted_key, evicted_val) = evicted.unwrap()[0];
assert_eq!(evicted_key, "c");
assert_eq!(evicted_val, 3);
assert_eq!(cache.get(&"a"), Some(&1)); assert_eq!(cache.get(&"b"), Some(&2)); assert_eq!(cache.get(&"d"), Some(&4)); assert_eq!(cache.get(&"c"), None); }
#[test]
fn test_lfu_frequency_ordering() {
let mut cache = make_cache(2);
cache.put("a", 1, 1);
cache.put("b", 2, 1);
cache.get(&"a");
cache.get(&"a");
cache.get(&"a");
cache.get(&"b");
let evicted = cache.put("c", 3, 1);
assert_eq!(evicted.unwrap()[0].0, "b");
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"c"), Some(&3));
assert_eq!(cache.get(&"b"), None);
}
#[test]
fn test_lfu_update_existing() {
let mut cache = make_cache(2);
cache.put("a", 1, 1);
cache.get(&"a");
let old_value = cache.put("a", 10, 1);
assert!(old_value.is_none());
cache.put("b", 2, 1);
cache.put("c", 3, 1);
assert_eq!(cache.get(&"a"), Some(&10));
assert_eq!(cache.get(&"c"), Some(&3));
assert_eq!(cache.get(&"b"), None);
}
#[test]
fn test_lfu_remove() {
let mut cache = make_cache(3);
cache.put("a", 1, 1);
cache.put("b", 2, 1);
cache.put("c", 3, 1);
assert_eq!(cache.remove(&"b"), Some(2));
assert_eq!(cache.remove(&"b"), None);
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"c"), Some(&3));
assert_eq!(cache.len(), 2);
}
#[test]
fn test_lfu_clear() {
let mut cache = make_cache(3);
cache.put("a", 1, 1);
cache.put("b", 2, 1);
cache.put("c", 3, 1);
assert_eq!(cache.len(), 3);
cache.clear();
assert_eq!(cache.len(), 0);
assert!(cache.is_empty());
cache.put("d", 4, 1);
assert_eq!(cache.get(&"d"), Some(&4));
}
#[test]
fn test_lfu_get_mut() {
let mut cache = make_cache(2);
cache.put("a", 1, 1);
if let Some(value) = cache.get_mut(&"a") {
*value = 10;
}
assert_eq!(cache.get(&"a"), Some(&10));
}
#[test]
fn test_lfu_complex_values() {
let mut cache = make_cache(2);
#[derive(Debug, Clone, PartialEq)]
struct ComplexValue {
id: usize,
data: String,
}
cache.put(
"a",
ComplexValue {
id: 1,
data: "a-data".to_string(),
},
1,
);
cache.put(
"b",
ComplexValue {
id: 2,
data: "b-data".to_string(),
},
1,
);
if let Some(value) = cache.get_mut(&"a") {
value.id = 100;
value.data = "a-modified".to_string();
}
let a = cache.get(&"a").unwrap();
assert_eq!(a.id, 100);
assert_eq!(a.data, "a-modified");
}
#[test]
fn test_miri_stacked_borrows_fix() {
let mut cache = make_cache(10);
cache.put("a", 1, 1);
cache.put("b", 2, 1);
cache.put("c", 3, 1);
for _ in 0..3 {
assert_eq!(cache.get(&"a"), Some(&1));
assert_eq!(cache.get(&"b"), Some(&2));
assert_eq!(cache.get(&"c"), Some(&3));
}
assert_eq!(cache.len(), 3);
if let Some(val) = cache.get_mut(&"a") {
*val += 10;
}
assert_eq!(cache.get(&"a"), Some(&11));
}
#[test]
fn test_lfu_segment_directly() {
let config = LfuCacheConfig {
capacity: NonZeroUsize::new(3).unwrap(),
max_size: u64::MAX,
};
let mut segment: LfuSegment<&str, i32, DefaultHashBuilder> =
LfuSegment::init(config, DefaultHashBuilder::default());
assert_eq!(segment.len(), 0);
assert!(segment.is_empty());
assert_eq!(segment.cap().get(), 3);
segment.put("a", 1, 1);
segment.put("b", 2, 1);
assert_eq!(segment.len(), 2);
assert_eq!(segment.get(&"a"), Some(&1));
assert_eq!(segment.get(&"a"), Some(&1));
assert_eq!(segment.get(&"b"), Some(&2));
}
#[test]
fn test_lfu_concurrent_access() {
extern crate std;
use std::sync::{Arc, Mutex};
use std::thread;
use std::vec::Vec;
let cache = Arc::new(Mutex::new(make_cache::<String, i32>(100)));
let num_threads = 4;
let ops_per_thread = 100;
let mut handles: Vec<std::thread::JoinHandle<()>> = Vec::new();
for t in 0..num_threads {
let cache = Arc::clone(&cache);
handles.push(thread::spawn(move || {
for i in 0..ops_per_thread {
let key = std::format!("key_{}_{}", t, i);
let mut guard = cache.lock().unwrap();
guard.put(key.clone(), i, 1);
if i % 3 == 0 {
let _ = guard.get(&key);
let _ = guard.get(&key);
}
}
}));
}
for handle in handles {
handle.join().unwrap();
}
let mut guard = cache.lock().unwrap();
assert!(guard.len() <= 100);
guard.clear(); }
#[test]
fn test_lfu_size_aware_tracking() {
let mut cache = make_cache(10);
assert_eq!(cache.current_size(), 0);
assert_eq!(cache.max_size(), u64::MAX);
cache.put("a", 1, 100);
cache.put("b", 2, 200);
cache.put("c", 3, 150);
assert_eq!(cache.current_size(), 450);
assert_eq!(cache.len(), 3);
cache.clear();
assert_eq!(cache.current_size(), 0);
}
#[test]
fn test_lfu_init_constructor() {
let config = LfuCacheConfig {
capacity: NonZeroUsize::new(1000).unwrap(),
max_size: 1024 * 1024,
};
let cache: LfuCache<String, i32> = LfuCache::init(config, None);
assert_eq!(cache.current_size(), 0);
assert_eq!(cache.max_size(), 1024 * 1024);
}
#[test]
fn test_lfu_with_limits_constructor() {
let config = LfuCacheConfig {
capacity: NonZeroUsize::new(100).unwrap(),
max_size: 1024 * 1024,
};
let cache: LfuCache<String, String> = LfuCache::init(config, None);
assert_eq!(cache.current_size(), 0);
assert_eq!(cache.max_size(), 1024 * 1024);
assert_eq!(cache.cap().get(), 100);
}
#[test]
fn test_lfu_frequency_list_accumulation() {
use std::time::Instant;
let mut cache = make_cache(500);
let num_ops = if cfg!(miri) { 5_000 } else { 50_000 };
let start = Instant::now();
for i in 0..num_ops {
let key = if i % 5 == 0 {
format!("popular_{}", i % 100) } else {
format!("long_tail_{}", i) };
cache.put(key.clone(), i, 1);
if i % 10 == 0 {
for j in 0..10 {
cache.get(&format!("popular_{}", j));
}
}
}
let elapsed = start.elapsed();
let num_freq_lists = cache.segment.frequency_lists.len();
let empty_lists = cache
.segment
.frequency_lists
.values()
.filter(|l| l.is_empty())
.count();
println!(
"{} ops in {:?} ({:.0} ops/sec)",
num_ops,
elapsed,
num_ops as f64 / elapsed.as_secs_f64()
);
println!(
"Frequency lists: {} total, {} empty",
num_freq_lists, empty_lists
);
let mut freq_counts: std::collections::BTreeMap<usize, usize> =
std::collections::BTreeMap::new();
for (freq, list) in &cache.segment.frequency_lists {
if !list.is_empty() {
*freq_counts.entry(*freq).or_insert(0) += list.len();
}
}
println!("Frequency distribution (non-empty): {:?}", freq_counts);
assert!(
empty_lists == 0,
"LFU should not accumulate empty frequency lists. Found {} empty lists out of {} total",
empty_lists,
num_freq_lists
);
}
#[test]
fn test_lfu_contains_non_promoting() {
let mut cache = make_cache(2);
cache.put("a", 1, 1);
cache.put("b", 2, 1);
assert!(cache.contains(&"a"));
assert!(cache.contains(&"b"));
assert!(!cache.contains(&"c"));
cache.get(&"b");
assert!(cache.contains(&"a"));
cache.put("c", 3, 1);
assert!(!cache.contains(&"a")); assert!(cache.contains(&"b")); assert!(cache.contains(&"c")); }
#[test]
fn test_put_returns_none_when_no_eviction() {
let mut cache = make_cache(10);
assert!(cache.put("a", 1, 1).is_none());
assert!(cache.put("b", 2, 1).is_none());
}
#[test]
fn test_put_returns_single_eviction() {
let mut cache = make_cache(2);
cache.put("a", 1, 1);
cache.put("b", 2, 1);
let result = cache.put("c", 3, 1);
assert_eq!(result, Some(vec![("a", 1)]));
}
#[test]
fn test_put_replacement_returns_none() {
let mut cache = make_cache(10);
cache.put("a", 1, 1);
let result = cache.put("a", 2, 1);
assert!(result.is_none());
assert_eq!(cache.get(&"a"), Some(&2));
}
#[test]
fn test_put_returns_multiple_evictions_size_based() {
let config = LfuCacheConfig {
capacity: NonZeroUsize::new(10).unwrap(),
max_size: 100,
};
let mut cache = LfuCache::init(config, None);
for i in 0..10 {
cache.put(i, i, 10);
}
let result = cache.put(99, 99, 50);
let evicted = result.unwrap();
assert_eq!(evicted.len(), 5);
}
}