use std::num::NonZeroUsize;
use std::sync::Mutex;
use lru::LruCache;
use xxhash_rust::xxh3::Xxh3;
use crate::eval::{CanvasInfo, TileId};
use crate::value::PortValue;
pub type Hash128 = u128;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CacheKey(pub Hash128);
impl CacheKey {
pub fn build(
canvas: CanvasInfo,
tile: Option<TileId>,
params_hash: Hash128,
inputs: &[Hash128],
) -> Self {
let mut h = Xxh3::new();
h.update(&canvas.tile_size.to_le_bytes());
h.update(&canvas.pad.to_le_bytes());
if let Some(t) = tile {
h.update(&[t.z]);
h.update(&t.x.to_le_bytes());
h.update(&t.y.to_le_bytes());
}
h.update(¶ms_hash.to_le_bytes());
for i in inputs {
h.update(&i.to_le_bytes());
}
CacheKey(h.digest128())
}
}
pub const DEFAULT_CAPACITY: usize = 4096;
pub const DEFAULT_BYTE_BUDGET: usize = 8 * 1024 * 1024;
pub struct Cache {
inner: Mutex<Inner>,
byte_budget: usize,
}
struct Inner {
lru: LruCache<CacheKey, PortValue>,
bytes: usize,
}
impl Default for Cache {
fn default() -> Self {
Self::new()
}
}
impl Cache {
pub fn new() -> Self {
Self::with_limits(DEFAULT_CAPACITY, DEFAULT_BYTE_BUDGET)
}
pub fn with_capacity(cap: usize) -> Self {
Self::with_limits(cap, DEFAULT_BYTE_BUDGET)
}
pub fn with_limits(cap: usize, byte_budget: usize) -> Self {
let cap = NonZeroUsize::new(cap.max(1)).expect("cap.max(1) is non-zero");
Self {
inner: Mutex::new(Inner {
lru: LruCache::new(cap),
bytes: 0,
}),
byte_budget,
}
}
pub fn get(&self, key: CacheKey) -> Option<PortValue> {
self.lock().lru.get(&key).cloned()
}
pub fn insert(&self, key: CacheKey, value: PortValue) {
let bytes = value.approx_bytes();
let mut inner = self.lock();
if let Some(old) = inner.lru.put(key, value) {
inner.bytes = inner.bytes.saturating_sub(old.approx_bytes());
}
inner.bytes += bytes;
while inner.bytes > self.byte_budget && inner.lru.peek_lru().is_some_and(|(k, _)| *k != key)
{
let Some((_, evicted)) = inner.lru.pop_lru() else {
break;
};
inner.bytes = inner.bytes.saturating_sub(evicted.approx_bytes());
}
}
pub fn len(&self) -> usize {
self.lock().lru.len()
}
pub fn is_empty(&self) -> bool {
self.lock().lru.is_empty()
}
pub fn clear(&self) {
let mut inner = self.lock();
inner.lru.clear();
inner.bytes = 0;
}
pub fn capacity(&self) -> usize {
self.lock().lru.cap().get()
}
pub fn byte_budget(&self) -> usize {
self.byte_budget
}
pub fn bytes(&self) -> usize {
self.lock().bytes
}
fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
self.inner.lock().unwrap_or_else(|e| e.into_inner())
}
}