use crate::types::RowId;
use parking_lot::RwLock;
use std::num::NonZeroUsize;
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum PkKey {
Int(i64),
Float(u64),
Text(Box<str>),
Bool(bool),
Null,
}
impl PkKey {
pub fn from_value(value: &crate::types::Value) -> Self {
match value {
crate::types::Value::Integer(i) => PkKey::Int(*i),
crate::types::Value::Float(f) => PkKey::Float(f.to_bits()),
crate::types::Value::Text(s) => PkKey::Text(s.as_str().to_string().into_boxed_str()),
crate::types::Value::Bool(b) => PkKey::Bool(*b),
crate::types::Value::Null => PkKey::Null,
_ => PkKey::Null,
}
}
pub fn from_hash_key(s: &str) -> Self {
if let Some(rest) = s.strip_prefix("i:") {
if let Ok(i) = rest.parse::<i64>() {
return PkKey::Int(i);
}
} else if let Some(rest) = s.strip_prefix("f:") {
if let Ok(bits) = rest.parse::<u64>() {
return PkKey::Float(bits);
}
} else if let Some(rest) = s.strip_prefix("t:") {
return PkKey::Text(rest.into());
} else if let Some(rest) = s.strip_prefix("b:") {
if let Ok(b) = rest.parse::<bool>() {
return PkKey::Bool(b);
}
}
PkKey::Null
}
}
pub struct PkLookupCache {
cache: RwLock<lru::LruCache<PkKey, RowId>>,
}
impl PkLookupCache {
pub fn new(capacity: usize) -> Self {
Self {
cache: RwLock::new(lru::LruCache::new(
NonZeroUsize::new(capacity.max(1)).unwrap(),
)),
}
}
pub fn insert(&self, key: PkKey, row_id: RowId) {
let mut cache = self.cache.write();
cache.put(key, row_id);
}
pub fn get(&self, key: &str) -> Option<RowId> {
let pk_key = PkKey::from_hash_key(key);
self.get_pk(&pk_key)
}
pub fn get_pk(&self, key: &PkKey) -> Option<RowId> {
{
let cache = self.cache.read();
if let Some(&row_id) = cache.peek(key) {
return Some(row_id);
}
}
None
}
pub fn remove(&self, key: &str) {
let mut cache = self.cache.write();
let pk_key = PkKey::from_hash_key(key);
cache.pop(&pk_key);
}
pub fn remove_pk(&self, key: &PkKey) {
let mut cache = self.cache.write();
cache.pop(key);
}
pub fn insert_if_absent(&self, key: PkKey, row_id: RowId) -> Result<(), RowId> {
let mut cache = self.cache.write();
if let Some(&existing) = cache.get(&key) {
return Err(existing);
}
cache.put(key, row_id);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pk_key_from_integer() {
let key = PkKey::from_value(&crate::types::Value::Integer(42));
assert_eq!(key, PkKey::Int(42));
}
#[test]
fn test_pk_key_from_text() {
let key = PkKey::from_value(&crate::types::Value::Text("hello".into()));
match key {
PkKey::Text(s) => assert_eq!(&*s, "hello"),
_ => panic!("expected Text"),
}
}
#[test]
fn test_pk_key_from_null() {
let key = PkKey::from_value(&crate::types::Value::Null);
assert_eq!(key, PkKey::Null);
}
#[test]
fn test_insert_and_get() {
let cache = PkLookupCache::new(100);
cache.insert(PkKey::Int(1), 100);
cache.insert(PkKey::Int(2), 200);
assert_eq!(cache.get_pk(&PkKey::Int(1)), Some(100));
assert_eq!(cache.get_pk(&PkKey::Int(2)), Some(200));
assert_eq!(cache.get_pk(&PkKey::Int(99)), None);
}
#[test]
fn test_remove() {
let cache = PkLookupCache::new(100);
cache.insert(PkKey::Int(1), 100);
cache.remove_pk(&PkKey::Int(1));
assert_eq!(cache.get_pk(&PkKey::Int(1)), None);
}
#[test]
fn test_overwrite() {
let cache = PkLookupCache::new(100);
cache.insert(PkKey::Int(1), 100);
cache.insert(PkKey::Int(1), 999);
assert_eq!(cache.get_pk(&PkKey::Int(1)), Some(999));
}
#[test]
fn test_lru_eviction() {
let cache = PkLookupCache::new(10);
for i in 0..20i64 {
cache.insert(PkKey::Int(i), i as RowId);
}
assert_eq!(cache.get_pk(&PkKey::Int(0)), None);
assert!(cache.get_pk(&PkKey::Int(19)).is_some());
}
}