use eyre::{bail, ensure, Context, OptionExt, Result};
use lru::LruCache;
use std::{hash::Hash, num::NonZeroUsize};
enum Value<V> {
Detached,
Pinned(V),
Regular(V),
}
impl<V> Value<V> {
fn get(&self) -> Result<&V> {
Ok(match self {
Value::Detached => bail!("Value is detached"),
Value::Pinned(v) => v,
Value::Regular(v) => v,
})
}
fn get_mut(&mut self) -> Result<&mut V> {
Ok(match self {
Value::Detached => bail!("Value is detached"),
Value::Pinned(v) => v,
Value::Regular(v) => v,
})
}
fn out(self, pin_count: &mut usize) -> Result<V> {
Ok(match self {
Value::Detached => bail!("Value is detached"),
Value::Pinned(v) => {
*pin_count -= 1;
v
}
Value::Regular(v) => v,
})
}
}
pub struct Cache<K: Clone + PartialEq + Eq + Hash, V> {
entries: LruCache<K, Value<V>>,
detach_count: usize,
pin_count: usize,
}
impl<K: Clone + PartialEq + Eq + Hash, V> Cache<K, V> {
pub fn attach(&mut self, key: &K, val: V) -> Result<()> {
let (key, d_val) = self
.entries
.pop_entry(key)
.ok_or_eyre("cannot attach: requested entry is not cached")?;
ensure!(
matches!(d_val, Value::Detached),
"only detached entries can be attached"
);
self.entries.push(key, Value::Regular(val));
self.detach_count -= 1;
Ok(())
}
pub fn detach(&mut self, key: &K) -> Result<V> {
let (key, val) = self
.entries
.pop_entry(key)
.ok_or_eyre("cannot detach: requested entry is not cached")?;
self.entries.push(key, Value::Detached);
self.detach_count += 1;
val.out(&mut self.pin_count)
}
pub fn empty(&mut self) -> Result<Vec<(K, V)>> {
let mut entries = Vec::new();
while let Some((key, val)) = self.entries.pop_lru() {
entries.push((key, val.out(&mut self.pin_count)?))
}
Ok(entries)
}
pub fn get(&mut self, key: &K) -> Result<&V> {
self.entries
.get(key)
.ok_or_eyre("cannot get: requested entry is not cached")?
.get()
}
pub fn get_mut(&mut self, key: &K) -> Result<&mut V> {
self.entries
.get_mut(key)
.ok_or_eyre("cannot get_mut: requested entry is not cached")?
.get_mut()
}
pub fn is_cached(&self, key: &K) -> bool {
self.entries.contains(key)
}
pub fn is_pinned(&self, key: &K) -> bool {
matches!(self.entries.peek(key), Some(Value::Pinned(_)))
}
#[allow(dead_code)]
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn margin(&self) -> usize {
self.entries.cap().get() - self.detach_count - self.pin_count
}
pub fn new(cap: NonZeroUsize) -> Self {
Self {
entries: LruCache::new(cap),
detach_count: 0,
pin_count: 0,
}
}
pub fn pin(&mut self, key: &K) -> Result<()> {
let (key, val) = self
.entries
.pop_entry(key)
.ok_or_eyre("cannot pin: requested entry is not cached")?;
self.entries
.push(key, Value::Pinned(val.out(&mut self.pin_count)?));
self.pin_count += 1;
Ok(())
}
pub fn pop(&mut self, key: &K) -> Result<V> {
self.entries
.pop(key)
.ok_or_eyre("cannot pop: requested entry is not cached")?
.out(&mut self.pin_count)
}
#[allow(dead_code)]
pub fn pop_lru(&mut self) -> Option<(K, V)> {
self.refresh().expect("failed to refresh the cache");
self.entries
.pop_lru()
.map(|(k, v)| (k, v.out(&mut 0).unwrap()))
}
pub fn push(&mut self, key: K, val: V) -> Result<Option<(K, V)>> {
self.refresh()
.wrap_err("failed to refresh the cache when pushing new entry")?;
Ok(
if let Some((k, v)) = self.entries.push(key, Value::Regular(val)) {
Some((k, v.out(&mut 0)?))
} else {
None
},
)
}
fn refresh(&mut self) -> Result<()> {
if self.entries.len() < self.entries.cap().get() {
return Ok(());
}
ensure!(
self.margin() > 0,
"cache overflow -- capacity: {}, detached: {}, pinned: {}",
self.entries.cap().get(),
self.detach_count,
self.pin_count
);
loop {
let key = match self.entries.peek_lru() {
Some((key, val)) if !matches!(val, Value::Regular(_)) => key.clone(),
_ => break,
};
self.entries.promote(&key);
}
Ok(())
}
pub fn unpin(&mut self, key: &K) -> Result<()> {
let (key, val) = self
.entries
.pop_entry(key)
.ok_or_eyre("cannot unpin: requested entry is not cached")?;
self.entries
.push(key, Value::Regular(val.out(&mut self.pin_count)?));
Ok(())
}
}