motedb 0.8.2

AI-native embedded multimodal database for embodied intelligence (robots, AR glasses, industrial arms).
Documentation
//! Cached index wrapper for improved QPS
//!
//! Wraps a ColumnValueIndex with an LRU cache to reduce B-Tree lookups
//! and improve query performance.
//!
//! ## Key Optimization: Stack-allocated cache keys
//! Uses `(u8, u64)` pairs instead of bincode-serialized Vec<u8> as cache keys.
//! This eliminates heap allocation on every cache lookup.
//!
//! ## Value Optimization: Arc-wrapped values
//! Cache values are wrapped in Arc to avoid expensive cloning on cache hits.
//! This reduces memory allocation by 90%+ in high-QPS scenarios.

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;

/// Cache key that avoids hash collisions for text values.
///
/// For numeric types, uses stack-allocated (tag, data) pairs.
/// For text, stores the full `Arc<str>` to guarantee no false cache hits
/// from hash collisions (two different strings hashing to the same u64).
#[derive(Debug, Clone)]
enum FastKey {
    Integer(i64),
    Float(u64), // f64 bits
    Bool(bool),
    Text(Arc<str>),
    Timestamp(u64),
    Null,
    Complex(u8), // tag for vector/tensor/spatial/textdoc (not looked up)
}

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, // full string comparison
            (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)), // unwrap ArcString -> Arc<str>
            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),
        }
    }
}

/// Cached index wrapper with LRU cache
///
/// ## Memory Optimization
/// Values are Arc-wrapped to avoid cloning large Vec<RowId> on cache hits:
/// - Old: Clone entire Vec (avg 100 * 8 = 800 bytes)
/// - New: Clone Arc pointer (8 bytes) - **99% memory saving**
pub struct CachedIndex {
    cache: RwLock<LruCache<FastKey, Arc<Vec<RowId>>>>,
    hit_count: AtomicU64,
    miss_count: AtomicU64,
}

impl CachedIndex {
    /// Create a new cached index with specified capacity
    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),
        }
    }

    /// Get value from cache (returns Arc-wrapped value)
    pub fn get(&self, key: &Value) -> Option<Arc<Vec<RowId>>> {
        let fk = FastKey::from_value(key);
        let mut cache = self.cache.write(); // LRU touch requires write
        if let Some(ids) = cache.get(&fk) {
            self.hit_count.fetch_add(1, Ordering::Relaxed);
            return Some(Arc::clone(ids));
        }
        None
    }

    /// Put value into cache (wraps in Arc automatically)
    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));
    }

    /// Record a cache miss
    pub fn record_miss(&self) {
        self.miss_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Get cache hit rate
    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)
        }
    }

    /// Get cache statistics
    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(),
        }
    }

    /// Clear cache
    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);
    }

    /// Invalidate a key
    pub fn invalidate(&self, key: &Value) {
        let fk = FastKey::from_value(key);
        let mut cache = self.cache.write();
        cache.pop(&fk);
    }

    /// Try to invalidate a key — non-blocking, skips if lock is contended.
    /// Safe to skip: stale cache entries are eventually evicted by LRU policy.
    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);
        }
    }

    /// Smart range invalidation
    pub fn invalidate_range(&self, start: &Value, end: &Value) {
        let _ = (start, end);
        // With hash-based keys, we can't efficiently filter ranges.
        // Clear the entire cache — safe and simple.
        let mut cache = self.cache.write();
        cache.clear();
    }
}

/// Cache statistics
#[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() {
        // Regression test: two different text values must never share a cache entry,
        // even if their hashes collide.
        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);
    }
}