use crate::types::{RowId, Value};
use lru::LruCache;
use parking_lot::RwLock;
use std::hash::{Hash, Hasher};
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
#[derive(Debug, Clone)]
enum FastKey {
Integer(i64),
Float(u64), Bool(bool),
Text(Arc<str>),
Timestamp(u64),
Null,
Complex(u8), }
impl Hash for FastKey {
fn hash<H: Hasher>(&self, state: &mut H) {
std::mem::discriminant(self).hash(state);
match self {
FastKey::Integer(i) => i.hash(state),
FastKey::Float(bits) => bits.hash(state),
FastKey::Bool(b) => b.hash(state),
FastKey::Text(s) => s.hash(state),
FastKey::Timestamp(ts) => ts.hash(state),
FastKey::Null | FastKey::Complex(_) => {}
}
}
}
impl PartialEq for FastKey {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(FastKey::Integer(a), FastKey::Integer(b)) => a == b,
(FastKey::Float(a), FastKey::Float(b)) => a == b,
(FastKey::Bool(a), FastKey::Bool(b)) => a == b,
(FastKey::Text(a), FastKey::Text(b)) => a == b, (FastKey::Timestamp(a), FastKey::Timestamp(b)) => a == b,
(FastKey::Null, FastKey::Null) => true,
(FastKey::Complex(a), FastKey::Complex(b)) => a == b,
_ => false,
}
}
}
impl Eq for FastKey {}
impl FastKey {
fn from_value(value: &Value) -> Self {
match value {
Value::Integer(i) => FastKey::Integer(*i),
Value::Float(f) => FastKey::Float(f.to_bits()),
Value::Bool(b) => FastKey::Bool(*b),
Value::Text(s) => FastKey::Text(Arc::clone(&s.0)), Value::Timestamp(ts) => FastKey::Timestamp(ts.as_micros_u64()),
Value::Null => FastKey::Null,
Value::Vector(_) => FastKey::Complex(7),
Value::Tensor(_) => FastKey::Complex(8),
Value::Spatial(_) => FastKey::Complex(9),
Value::TextDoc(_) => FastKey::Complex(10),
}
}
}
pub struct CachedIndex {
cache: RwLock<LruCache<FastKey, Arc<Vec<RowId>>>>,
hit_count: AtomicU64,
miss_count: AtomicU64,
}
impl CachedIndex {
pub fn new(capacity: usize) -> Self {
let capacity = NonZeroUsize::new(capacity).unwrap_or(NonZeroUsize::new(1000).unwrap());
Self {
cache: RwLock::new(LruCache::new(capacity)),
hit_count: AtomicU64::new(0),
miss_count: AtomicU64::new(0),
}
}
pub fn get(&self, key: &Value) -> Option<Arc<Vec<RowId>>> {
let fk = FastKey::from_value(key);
let mut cache = self.cache.write(); if let Some(ids) = cache.get(&fk) {
self.hit_count.fetch_add(1, Ordering::Relaxed);
return Some(Arc::clone(ids));
}
None
}
pub fn put(&self, key: Value, ids: Vec<RowId>) {
let fk = FastKey::from_value(&key);
let mut cache = self.cache.write();
cache.put(fk, Arc::new(ids));
}
pub fn record_miss(&self) {
self.miss_count.fetch_add(1, Ordering::Relaxed);
}
pub fn hit_rate(&self) -> f64 {
let hits = self.hit_count.load(Ordering::Relaxed) as f64;
let misses = self.miss_count.load(Ordering::Relaxed) as f64;
if hits + misses == 0.0 {
0.0
} else {
hits / (hits + misses)
}
}
pub fn stats(&self) -> CacheStats {
let cache = self.cache.read();
CacheStats {
capacity: cache.cap().get(),
size: cache.len(),
hits: self.hit_count.load(Ordering::Relaxed),
misses: self.miss_count.load(Ordering::Relaxed),
hit_rate: self.hit_rate(),
}
}
pub fn clear(&self) {
let mut cache = self.cache.write();
cache.clear();
self.hit_count.store(0, Ordering::Relaxed);
self.miss_count.store(0, Ordering::Relaxed);
}
pub fn invalidate(&self, key: &Value) {
let fk = FastKey::from_value(key);
let mut cache = self.cache.write();
cache.pop(&fk);
}
pub fn try_invalidate(&self, key: &Value) {
let fk = FastKey::from_value(key);
if let Some(mut cache) = self.cache.try_write() {
cache.pop(&fk);
}
}
pub fn invalidate_range(&self, start: &Value, end: &Value) {
let _ = (start, end);
let mut cache = self.cache.write();
cache.clear();
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub capacity: usize,
pub size: usize,
pub hits: u64,
pub misses: u64,
pub hit_rate: f64,
}
impl std::fmt::Display for CacheStats {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(
f,
"Cache: {}/{} entries, {:.1}% hit rate ({} hits, {} misses)",
self.size,
self.capacity,
self.hit_rate * 100.0,
self.hits,
self.misses
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_basic() {
let cache = CachedIndex::new(100);
assert_eq!(cache.stats().size, 0);
assert_eq!(cache.hit_rate(), 0.0);
cache.put(Value::Integer(1), vec![100, 200]);
cache.put(Value::Integer(2), vec![300]);
assert_eq!(*cache.get(&Value::Integer(1)).unwrap(), vec![100, 200]);
assert_eq!(*cache.get(&Value::Integer(2)).unwrap(), vec![300]);
assert_eq!(cache.get(&Value::Integer(3)), None);
cache.record_miss();
let stats = cache.stats();
assert_eq!(stats.size, 2);
assert_eq!(stats.hits, 2);
assert_eq!(stats.misses, 1);
assert!((stats.hit_rate - 0.666).abs() < 0.01);
}
#[test]
fn test_cache_lru_eviction() {
let cache = CachedIndex::new(2);
cache.put(Value::Integer(1), vec![100]);
cache.put(Value::Integer(2), vec![200]);
cache.get(&Value::Integer(1));
cache.put(Value::Integer(3), vec![300]);
assert_eq!(*cache.get(&Value::Integer(1)).unwrap(), vec![100]);
assert_eq!(*cache.get(&Value::Integer(3)).unwrap(), vec![300]);
assert_eq!(cache.get(&Value::Integer(2)), None);
}
#[test]
fn test_cache_invalidate() {
let cache = CachedIndex::new(100);
cache.put(Value::Integer(1), vec![100]);
assert_eq!(*cache.get(&Value::Integer(1)).unwrap(), vec![100]);
cache.invalidate(&Value::Integer(1));
assert_eq!(cache.get(&Value::Integer(1)), None);
}
#[test]
fn test_cache_clear() {
let cache = CachedIndex::new(100);
cache.put(Value::Integer(1), vec![100]);
cache.put(Value::Integer(2), vec![200]);
cache.get(&Value::Integer(1));
assert_eq!(cache.stats().size, 2);
assert_eq!(cache.stats().hits, 1);
cache.clear();
assert_eq!(cache.stats().size, 0);
assert_eq!(cache.stats().hits, 0);
assert_eq!(cache.stats().misses, 0);
}
#[test]
fn test_cache_text_key() {
let cache = CachedIndex::new(100);
cache.put(Value::text("hello".to_string()), vec![1, 2]);
assert_eq!(
*cache.get(&Value::text("hello".to_string())).unwrap(),
vec![1, 2]
);
assert_eq!(cache.get(&Value::text("world".to_string())), None);
}
#[test]
fn test_cache_text_no_collision() {
let cache = CachedIndex::new(100);
cache.put(Value::text("alpha".to_string()), vec![1]);
cache.put(Value::text("beta".to_string()), vec![2]);
assert_eq!(
*cache.get(&Value::text("alpha".to_string())).unwrap(),
vec![1]
);
assert_eq!(
*cache.get(&Value::text("beta".to_string())).unwrap(),
vec![2]
);
assert_eq!(cache.get(&Value::text("gamma".to_string())), None);
}
#[test]
fn test_cache_float_key() {
let cache = CachedIndex::new(100);
cache.put(Value::Float(3.14), vec![42]);
assert_eq!(*cache.get(&Value::Float(3.14)).unwrap(), vec![42]);
assert_eq!(cache.get(&Value::Float(2.71)), None);
}
}