use std::{
collections::{HashMap, VecDeque},
hash::Hash,
};
pub(in crate::db::session) struct BoundedCache<K, V> {
entries: HashMap<K, BoundedCacheEntry<V>>,
insertion_order: VecDeque<K>,
max_entries: usize,
max_retained_weight: usize,
retained_weight: usize,
}
struct BoundedCacheEntry<V> {
value: V,
weight: usize,
}
pub(in crate::db::session) struct BoundedCacheInsertOutcome<V> {
pub(in crate::db::session) replaced: Option<V>,
pub(in crate::db::session) evicted: usize,
pub(in crate::db::session) rejected_oversize: bool,
}
impl<K, V> BoundedCache<K, V>
where
K: Clone + Eq + Hash,
{
pub(in crate::db::session) fn new(max_entries: usize) -> Self {
Self::new_weighted(max_entries, usize::MAX)
}
pub(in crate::db::session) fn new_weighted(
max_entries: usize,
max_retained_weight: usize,
) -> Self {
Self {
entries: HashMap::new(),
insertion_order: VecDeque::new(),
max_entries,
max_retained_weight,
retained_weight: 0,
}
}
pub(in crate::db::session) fn get(&self, key: &K) -> Option<&V> {
self.entries.get(key).map(|entry| &entry.value)
}
pub(in crate::db::session) fn get_mut(&mut self, key: &K) -> Option<&mut V> {
self.entries.get_mut(key).map(|entry| &mut entry.value)
}
pub(in crate::db::session) fn insert(&mut self, key: K, value: V) -> Option<V> {
self.insert_weighted(key, value, 0).replaced
}
pub(in crate::db::session) fn insert_weighted(
&mut self,
key: K,
value: V,
weight: usize,
) -> BoundedCacheInsertOutcome<V> {
if self.max_entries == 0 || weight > self.max_retained_weight {
return BoundedCacheInsertOutcome {
replaced: None,
evicted: 0,
rejected_oversize: true,
};
}
if let Some(replaced) = self.entries.remove(&key) {
self.retained_weight = self.retained_weight.saturating_sub(replaced.weight);
self.insertion_order.retain(|existing| existing != &key);
let evicted = self.evict_until_new_key_fits(weight);
self.insertion_order.push_back(key.clone());
self.retained_weight = self.retained_weight.saturating_add(weight);
self.entries
.insert(key, BoundedCacheEntry { value, weight });
return BoundedCacheInsertOutcome {
replaced: Some(replaced.value),
evicted,
rejected_oversize: false,
};
}
let evicted = self.evict_until_new_key_fits(weight);
self.insertion_order.push_back(key.clone());
self.retained_weight = self.retained_weight.saturating_add(weight);
self.entries
.insert(key, BoundedCacheEntry { value, weight });
BoundedCacheInsertOutcome {
replaced: None,
evicted,
rejected_oversize: false,
}
}
pub(in crate::db::session) fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub(in crate::db::session) fn len(&self) -> usize {
self.entries.len()
}
pub(in crate::db::session) fn keys(&self) -> impl Iterator<Item = &K> {
self.entries.keys()
}
#[cfg(test)]
pub(in crate::db::session) fn values(&self) -> impl Iterator<Item = &V> {
self.entries.values().map(|entry| &entry.value)
}
#[cfg(test)]
pub(in crate::db::session) const fn retained_weight(&self) -> usize {
self.retained_weight
}
fn evict_until_new_key_fits(&mut self, new_weight: usize) -> usize {
let mut evicted = 0usize;
while self.entries.len() >= self.max_entries
|| self.retained_weight.saturating_add(new_weight) > self.max_retained_weight
{
if !self.evict_oldest() {
break;
}
evicted = evicted.saturating_add(1);
}
evicted
}
fn evict_oldest(&mut self) -> bool {
let Some(oldest) = self.insertion_order.pop_front() else {
self.entries.clear();
self.retained_weight = 0;
return false;
};
if let Some(entry) = self.entries.remove(&oldest) {
self.retained_weight = self.retained_weight.saturating_sub(entry.weight);
}
true
}
}
impl<K, V> Default for BoundedCache<K, V>
where
K: Clone + Eq + Hash,
{
fn default() -> Self {
Self::new(1024)
}
}
#[cfg(test)]
mod tests {
use super::BoundedCache;
#[test]
fn bounded_cache_evicts_oldest_inserted_key() {
let mut cache = BoundedCache::new(2);
cache.insert("a", 1);
cache.insert("b", 2);
cache.insert("c", 3);
assert!(cache.get(&"a").is_none());
assert_eq!(cache.get(&"b"), Some(&2));
assert_eq!(cache.get(&"c"), Some(&3));
assert_eq!(cache.len(), 2);
}
#[test]
fn weighted_cache_evicts_until_new_entry_fits() {
let mut cache = BoundedCache::new_weighted(4, 10);
cache.insert_weighted("a", 1, 4);
cache.insert_weighted("b", 2, 4);
let outcome = cache.insert_weighted("c", 3, 4);
assert_eq!(outcome.evicted, 1);
assert!(cache.get(&"a").is_none());
assert_eq!(cache.retained_weight(), 8);
}
#[test]
fn weighted_cache_rejects_single_oversize_entry() {
let mut cache = BoundedCache::new_weighted(4, 10);
let outcome = cache.insert_weighted("a", 1, 11);
assert!(outcome.rejected_oversize);
assert!(cache.is_empty());
assert_eq!(cache.retained_weight(), 0);
}
#[test]
fn weighted_cache_replacement_cannot_evict_itself() {
let mut cache = BoundedCache::new_weighted(4, 10);
cache.insert_weighted("a", 1, 4);
cache.insert_weighted("b", 2, 4);
let outcome = cache.insert_weighted("a", 3, 8);
assert_eq!(outcome.replaced, Some(1));
assert_eq!(outcome.evicted, 1);
assert_eq!(cache.get(&"a"), Some(&3));
assert!(cache.get(&"b").is_none());
assert_eq!(cache.retained_weight(), 8);
}
}