use std::{
borrow::Borrow,
hash::{BuildHasher, Hash},
num::NonZeroUsize,
};
use hashbrown::HashTable;
use crate::collections::map::RandomState;
struct CacheSlot<K, V> {
key: K,
value: V,
newer: Option<u32>,
older: Option<u32>,
}
pub struct BoundedLruCache<K, V> {
index: HashTable<u32>,
hasher: RandomState,
slots: Vec<Option<CacheSlot<K, V>>>,
free: Vec<u32>,
newest: Option<u32>,
oldest: Option<u32>,
cap: NonZeroUsize,
}
fn occupied<K, V>(slots: &[Option<CacheSlot<K, V>>], slot: u32) -> &CacheSlot<K, V> {
slots[slot as usize]
.as_ref()
.expect("an indexed or linked cache slot is always occupied")
}
impl<K, V> BoundedLruCache<K, V>
where
K: Eq + Hash,
{
pub fn new(cap: NonZeroUsize) -> Self {
Self {
index: HashTable::new(),
hasher: RandomState::default(),
slots: Vec::new(),
free: Vec::new(),
newest: None,
oldest: None,
cap,
}
}
pub fn with_capacity_at_least_one(cap: usize) -> Self {
let cap = NonZeroUsize::new(cap).unwrap_or(NonZeroUsize::MIN);
Self::new(cap)
}
pub fn len(&self) -> usize {
self.index.len()
}
pub fn is_empty(&self) -> bool {
self.index.is_empty()
}
pub fn cap(&self) -> NonZeroUsize {
self.cap
}
pub fn clear(&mut self) {
*self = Self::new(self.cap);
}
pub fn contains<Q>(&self, key: &Q) -> bool
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.find(key).is_some()
}
pub fn get<Q>(&mut self, key: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let slot = self.find(key)?;
self.promote(slot);
Some(&self.slot(slot).value)
}
pub fn peek<Q>(&self, key: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let slot = self.find(key)?;
Some(&self.slot(slot).value)
}
pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let slot = self.find(key)?;
self.promote(slot);
Some(&mut self.slot_mut(slot).value)
}
pub fn push(&mut self, key: K, value: V) -> Option<(K, V)> {
if let Some(slot) = self.find(&key) {
self.promote(slot);
let old_value = std::mem::replace(&mut self.slot_mut(slot).value, value);
return Some((key, old_value));
}
let evicted = if self.index.len() == self.cap.get() {
self.pop_lru()
} else {
None
};
let hash = self.hasher.hash_one(&key);
let slot = self.claim_slot(key, value);
let (slots, hasher) = (&self.slots, &self.hasher);
self.index.insert_unique(hash, slot, |&slot| {
hasher.hash_one(&occupied(slots, slot).key)
});
self.link_newest(slot);
evicted
}
pub fn put(&mut self, key: K, value: V) -> Option<V> {
self.push(key, value).map(|(_, value)| value)
}
pub fn peek_lru(&self) -> Option<(&K, &V)> {
let entry = self.slot(self.oldest?);
Some((&entry.key, &entry.value))
}
pub fn pop_lru(&mut self) -> Option<(K, V)> {
let slot = self.oldest?;
let hash = self.hasher.hash_one(&self.slot(slot).key);
if let Ok(entry) = self.index.find_entry(hash, |&indexed| indexed == slot) {
entry.remove();
}
self.unlink(slot);
let entry = self.release_slot(slot);
Some((entry.key, entry.value))
}
pub fn pop(&mut self, key: &K) -> Option<V> {
let hash = self.hasher.hash_one(key);
let slots = &self.slots;
let (slot, _) = self
.index
.find_entry(hash, |&slot| occupied(slots, slot).key == *key)
.ok()?
.remove();
self.unlink(slot);
Some(self.release_slot(slot).value)
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
let mut next = self.newest;
std::iter::from_fn(move || {
let entry = self.slot(next?);
next = entry.older;
Some((&entry.key, &entry.value))
})
}
fn find<Q>(&self, key: &Q) -> Option<u32>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
let slots = &self.slots;
self.index
.find(self.hasher.hash_one(key), |&slot| {
occupied(slots, slot).key.borrow() == key
})
.copied()
}
fn slot(&self, slot: u32) -> &CacheSlot<K, V> {
occupied(&self.slots, slot)
}
fn slot_mut(&mut self, slot: u32) -> &mut CacheSlot<K, V> {
self.slots[slot as usize]
.as_mut()
.expect("a linked cache slot is always occupied")
}
fn promote(&mut self, slot: u32) {
if self.newest == Some(slot) {
return;
}
self.unlink(slot);
self.link_newest(slot);
}
fn link_newest(&mut self, slot: u32) {
let previous_newest = self.newest;
{
let entry = self.slot_mut(slot);
entry.newer = None;
entry.older = previous_newest;
}
if let Some(previous) = previous_newest {
self.slot_mut(previous).newer = Some(slot);
}
self.newest = Some(slot);
if self.oldest.is_none() {
self.oldest = Some(slot);
}
}
fn unlink(&mut self, slot: u32) {
let (newer, older) = {
let entry = self.slot_mut(slot);
(entry.newer.take(), entry.older.take())
};
match newer {
Some(newer) => self.slot_mut(newer).older = older,
None => self.newest = older,
}
match older {
Some(older) => self.slot_mut(older).newer = newer,
None => self.oldest = newer,
}
}
fn claim_slot(&mut self, key: K, value: V) -> u32 {
let entry = CacheSlot {
key,
value,
newer: None,
older: None,
};
match self.free.pop() {
Some(slot) => {
self.slots[slot as usize] = Some(entry);
slot
}
None => {
let slot =
u32::try_from(self.slots.len()).expect("a cache holds fewer than 2^32 entries");
self.slots.push(Some(entry));
slot
}
}
}
fn release_slot(&mut self, slot: u32) -> CacheSlot<K, V> {
let entry = self.slots[slot as usize]
.take()
.expect("a slot being released is always occupied");
self.free.push(slot);
entry
}
}
#[cfg(test)]
#[path = "tests/bounded_lru_tests.rs"]
mod tests;