use std::any::Any;
use std::fmt;
use std::num::NonZeroUsize;
use lru::LruCache;
use tenferro_tensor::CacheStats;
pub const DEFAULT_EXTENSION_CACHE_CAPACITY: usize = 256;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ExtensionCacheKey {
pub family_id: &'static str,
pub cache_name: &'static str,
pub discriminator: u64,
}
impl ExtensionCacheKey {
pub const fn new(
family_id: &'static str,
cache_name: &'static str,
discriminator: u64,
) -> Self {
Self {
family_id,
cache_name,
discriminator,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ExtensionCacheSelector {
All,
Family { family_id: &'static str },
Cache {
family_id: &'static str,
cache_name: &'static str,
},
}
impl ExtensionCacheSelector {
pub fn matches(&self, key: &ExtensionCacheKey) -> bool {
match *self {
Self::All => true,
Self::Family { family_id } => key.family_id == family_id,
Self::Cache {
family_id,
cache_name,
} => key.family_id == family_id && key.cache_name == cache_name,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ExtensionCacheLimits {
max_entries: NonZeroUsize,
}
impl ExtensionCacheLimits {
pub const fn new(max_entries: NonZeroUsize) -> Self {
Self { max_entries }
}
pub const fn max_entries(self) -> NonZeroUsize {
self.max_entries
}
}
impl Default for ExtensionCacheLimits {
fn default() -> Self {
Self {
max_entries: NonZeroUsize::new(DEFAULT_EXTENSION_CACHE_CAPACITY)
.unwrap_or(NonZeroUsize::MIN),
}
}
}
trait ExtensionCacheValue: Send + Sync {
fn as_any(&self) -> &(dyn Any + Send + Sync);
fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync);
fn retained_bytes(&self) -> usize;
}
struct FixedRetainedBytes<T> {
value: T,
retained_bytes: usize,
}
impl<T> ExtensionCacheValue for FixedRetainedBytes<T>
where
T: Any + Send + Sync + 'static,
{
fn as_any(&self) -> &(dyn Any + Send + Sync) {
&self.value
}
fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
&mut self.value
}
fn retained_bytes(&self) -> usize {
self.retained_bytes
}
}
struct DynamicRetainedBytes<T, F> {
value: T,
retained_bytes: F,
}
impl<T, F> ExtensionCacheValue for DynamicRetainedBytes<T, F>
where
T: Any + Send + Sync + 'static,
F: Fn(&T) -> usize + Send + Sync + 'static,
{
fn as_any(&self) -> &(dyn Any + Send + Sync) {
&self.value
}
fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
&mut self.value
}
fn retained_bytes(&self) -> usize {
(self.retained_bytes)(&self.value)
}
}
struct ExtensionCacheEntry {
value: Box<dyn ExtensionCacheValue>,
}
pub struct ExtensionCacheStore {
limits: ExtensionCacheLimits,
entries: LruCache<ExtensionCacheKey, ExtensionCacheEntry>,
}
impl fmt::Debug for ExtensionCacheStore {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ExtensionCacheStore")
.field("limits", &self.limits)
.field("stats", &self.stats(ExtensionCacheSelector::All))
.finish_non_exhaustive()
}
}
impl ExtensionCacheStore {
pub fn new() -> Self {
Self::with_limits(ExtensionCacheLimits::default())
}
pub fn with_limits(limits: ExtensionCacheLimits) -> Self {
Self {
entries: LruCache::new(limits.max_entries()),
limits,
}
}
pub const fn limits(&self) -> ExtensionCacheLimits {
self.limits
}
pub fn set_limits(&mut self, limits: ExtensionCacheLimits) {
self.entries.resize(limits.max_entries());
self.limits = limits;
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn put<T>(&mut self, key: ExtensionCacheKey, value: T, retained_bytes: usize)
where
T: Any + Send + Sync + 'static,
{
self.entries.put(
key,
ExtensionCacheEntry {
value: Box::new(FixedRetainedBytes {
value,
retained_bytes,
}),
},
);
}
pub fn put_with_retained_bytes<T, F>(
&mut self,
key: ExtensionCacheKey,
value: T,
retained_bytes: F,
) where
T: Any + Send + Sync + 'static,
F: Fn(&T) -> usize + Send + Sync + 'static,
{
self.entries.put(
key,
ExtensionCacheEntry {
value: Box::new(DynamicRetainedBytes {
value,
retained_bytes,
}),
},
);
}
pub fn get<T>(&mut self, key: &ExtensionCacheKey) -> Option<&T>
where
T: Any + Send + Sync + 'static,
{
self.entries
.get(key)
.and_then(|entry| entry.value.as_any().downcast_ref::<T>())
}
pub fn get_mut<T>(&mut self, key: &ExtensionCacheKey) -> Option<&mut T>
where
T: Any + Send + Sync + 'static,
{
self.entries
.get_mut(key)
.and_then(|entry| entry.value.as_any_mut().downcast_mut::<T>())
}
pub fn clear_selected(&mut self, selector: ExtensionCacheSelector) {
if selector == ExtensionCacheSelector::All {
self.entries.clear();
return;
}
let keys: Vec<_> = self
.entries
.iter()
.map(|(key, _)| *key)
.filter(|key| selector.matches(key))
.collect();
for key in keys {
self.entries.pop(&key);
}
}
pub fn clear(&mut self) {
self.entries.clear();
}
pub fn stats(&self, selector: ExtensionCacheSelector) -> CacheStats {
CacheStats {
entries: self
.entries
.iter()
.filter(|(key, _)| selector.matches(key))
.count(),
retained_bytes: self
.entries
.iter()
.filter(|(key, _)| selector.matches(key))
.map(|(_, entry)| entry.value.retained_bytes())
.fold(0usize, usize::saturating_add),
}
}
}
impl Default for ExtensionCacheStore {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;