ipfrs_semantic/

cache.rs

//! Advanced caching for vector embeddings
//!
//! This module provides high-performance caching strategies for vector embeddings
//! including hot embedding cache, adaptive caching, and cache-aligned storage.

use lru::LruCache;
use parking_lot::RwLock;
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Cache-aligned vector storage
///
/// Vectors are aligned to cache line boundaries (64 bytes) to reduce
/// cache misses and improve SIMD performance.
#[repr(align(64))]
#[derive(Debug, Clone)]
pub struct AlignedVector {
    data: Vec<f32>,
}

impl AlignedVector {
    /// Create a new cache-aligned vector
    pub fn new(data: Vec<f32>) -> Self {
        Self { data }
    }

    /// Create a zeroed cache-aligned vector
    pub fn zeros(len: usize) -> Self {
        Self {
            data: vec![0.0; len],
        }
    }

    /// Get a reference to the underlying data
    pub fn as_slice(&self) -> &[f32] {
        &self.data
    }

    /// Get a mutable reference to the underlying data
    pub fn as_mut_slice(&mut self) -> &mut [f32] {
        &mut self.data
    }

    /// Get the length of the vector
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Check if the vector is empty
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Convert into the underlying Vec
    pub fn into_vec(self) -> Vec<f32> {
        self.data
    }
}

impl From<Vec<f32>> for AlignedVector {
    fn from(data: Vec<f32>) -> Self {
        Self::new(data)
    }
}

impl AsRef<[f32]> for AlignedVector {
    fn as_ref(&self) -> &[f32] {
        &self.data
    }
}

impl AsMut<[f32]> for AlignedVector {
    fn as_mut(&mut self) -> &mut [f32] {
        &mut self.data
    }
}

/// Access statistics for a cached item
#[derive(Debug, Clone)]
struct AccessStats {
    /// Number of times accessed
    access_count: u64,
    /// Last access time
    last_access: Instant,
    /// First access time
    first_access: Instant,
    /// Total time in cache (for adaptive sizing)
    time_in_cache: Duration,
}

impl AccessStats {
    fn new() -> Self {
        let now = Instant::now();
        Self {
            access_count: 1,
            last_access: now,
            first_access: now,
            time_in_cache: Duration::from_secs(0),
        }
    }

    fn record_access(&mut self) {
        self.access_count += 1;
        self.last_access = Instant::now();
        self.time_in_cache = self.last_access.duration_since(self.first_access);
    }

    fn access_frequency(&self) -> f64 {
        if self.time_in_cache.as_secs_f64() > 0.0 {
            self.access_count as f64 / self.time_in_cache.as_secs_f64()
        } else {
            self.access_count as f64
        }
    }
}

/// Cached embedding entry
#[derive(Debug, Clone)]
struct CachedEmbedding {
    vector: AlignedVector,
    stats: AccessStats,
}

/// Hot embedding cache with LRU eviction and adaptive sizing
///
/// Caches frequently accessed embeddings in memory with cache-aligned
/// storage for optimal SIMD performance.
pub struct HotEmbeddingCache {
    /// LRU cache for embeddings
    cache: Arc<RwLock<LruCache<String, CachedEmbedding>>>,
    /// Access statistics (hits/misses)
    hits: Arc<AtomicU64>,
    misses: Arc<AtomicU64>,
    /// Total cache capacity
    capacity: usize,
    /// Prefetch queue for predicted accesses
    prefetch_queue: Arc<RwLock<Vec<String>>>,
}

impl HotEmbeddingCache {
    /// Create a new hot embedding cache
    pub fn new(capacity: usize) -> Self {
        Self {
            cache: Arc::new(RwLock::new(LruCache::new(
                NonZeroUsize::new(capacity).unwrap(),
            ))),
            hits: Arc::new(AtomicU64::new(0)),
            misses: Arc::new(AtomicU64::new(0)),
            capacity,
            prefetch_queue: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Get an embedding from the cache
    pub fn get(&self, key: &str) -> Option<AlignedVector> {
        let mut cache = self.cache.write();
        if let Some(entry) = cache.get_mut(key) {
            entry.stats.record_access();
            self.hits.fetch_add(1, Ordering::Relaxed);
            Some(entry.vector.clone())
        } else {
            self.misses.fetch_add(1, Ordering::Relaxed);
            None
        }
    }

    /// Insert an embedding into the cache
    pub fn insert(&self, key: String, vector: Vec<f32>) {
        let aligned = AlignedVector::new(vector);
        let entry = CachedEmbedding {
            vector: aligned,
            stats: AccessStats::new(),
        };
        self.cache.write().put(key, entry);
    }

    /// Get cache statistics
    pub fn stats(&self) -> HotCacheStats {
        let hits = self.hits.load(Ordering::Relaxed);
        let misses = self.misses.load(Ordering::Relaxed);
        let total = hits + misses;
        let hit_rate = if total > 0 {
            hits as f64 / total as f64
        } else {
            0.0
        };

        let cache = self.cache.read();
        let size = cache.len();

        HotCacheStats {
            hits,
            misses,
            hit_rate,
            size,
            capacity: self.capacity,
        }
    }

    /// Clear the cache
    pub fn clear(&self) {
        self.cache.write().clear();
        self.hits.store(0, Ordering::Relaxed);
        self.misses.store(0, Ordering::Relaxed);
    }

    /// Get the current size of the cache
    pub fn len(&self) -> usize {
        self.cache.read().len()
    }

    /// Check if the cache is empty
    pub fn is_empty(&self) -> bool {
        self.cache.read().is_empty()
    }

    /// Prefetch embeddings (add to prefetch queue)
    pub fn prefetch(&self, keys: Vec<String>) {
        let mut queue = self.prefetch_queue.write();
        queue.extend(keys);
    }

    /// Get hot embeddings (most frequently accessed)
    pub fn get_hot_keys(&self, top_n: usize) -> Vec<String> {
        let cache = self.cache.read();
        let mut entries: Vec<_> = cache
            .iter()
            .map(|(k, v)| (k.clone(), v.stats.access_frequency()))
            .collect();

        entries.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        entries.into_iter().take(top_n).map(|(k, _)| k).collect()
    }
}

/// Statistics for hot embedding cache
#[derive(Debug, Clone)]
pub struct HotCacheStats {
    /// Number of cache hits
    pub hits: u64,
    /// Number of cache misses
    pub misses: u64,
    /// Hit rate (hits / total accesses)
    pub hit_rate: f64,
    /// Current cache size
    pub size: usize,
    /// Maximum cache capacity
    pub capacity: usize,
}

/// Adaptive caching strategy
///
/// Dynamically adjusts cache size and eviction policy based on
/// access patterns and hit rates.
pub struct AdaptiveCacheStrategy {
    /// Current cache size target
    target_size: Arc<RwLock<usize>>,
    /// Minimum cache size
    min_size: usize,
    /// Maximum cache size
    max_size: usize,
    /// Target hit rate (0.0-1.0)
    target_hit_rate: f64,
    /// Adjustment factor for cache sizing
    adjustment_factor: f64,
}

impl AdaptiveCacheStrategy {
    /// Create a new adaptive caching strategy
    pub fn new(min_size: usize, max_size: usize, target_hit_rate: f64) -> Self {
        Self {
            target_size: Arc::new(RwLock::new((min_size + max_size) / 2)),
            min_size,
            max_size,
            target_hit_rate,
            adjustment_factor: 1.1, // 10% adjustment per iteration
        }
    }

    /// Adjust cache size based on current hit rate
    pub fn adjust(&self, current_hit_rate: f64) -> usize {
        let mut target = self.target_size.write();

        if current_hit_rate < self.target_hit_rate {
            // Hit rate too low, increase cache size
            let new_size =
                (*target as f64 * self.adjustment_factor).min(self.max_size as f64) as usize;
            *target = new_size;
        } else if current_hit_rate > self.target_hit_rate + 0.05 {
            // Hit rate high enough, can reduce cache size
            let new_size =
                (*target as f64 / self.adjustment_factor).max(self.min_size as f64) as usize;
            *target = new_size;
        }

        *target
    }

    /// Get the current target cache size
    pub fn target_size(&self) -> usize {
        *self.target_size.read()
    }

    /// Reset to default size
    pub fn reset(&self) {
        *self.target_size.write() = (self.min_size + self.max_size) / 2;
    }
}

/// Cache invalidation policy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InvalidationPolicy {
    /// Time-to-live based invalidation
    TTL(Duration),
    /// Event-driven invalidation (manual)
    Event,
    /// Never invalidate (manual only)
    Never,
}

/// Cache invalidation tracker
pub struct CacheInvalidator {
    /// Invalidation policy
    policy: InvalidationPolicy,
    /// Timestamp of last invalidation
    last_invalidation: Arc<RwLock<Instant>>,
}

impl CacheInvalidator {
    /// Create a new cache invalidator
    pub fn new(policy: InvalidationPolicy) -> Self {
        Self {
            policy,
            last_invalidation: Arc::new(RwLock::new(Instant::now())),
        }
    }

    /// Check if cache should be invalidated
    pub fn should_invalidate(&self) -> bool {
        match self.policy {
            InvalidationPolicy::TTL(ttl) => {
                let elapsed = self.last_invalidation.read().elapsed();
                elapsed >= ttl
            }
            InvalidationPolicy::Event => false, // Manual invalidation only
            InvalidationPolicy::Never => false,
        }
    }

    /// Mark cache as invalidated
    pub fn invalidate(&self) {
        *self.last_invalidation.write() = Instant::now();
    }

    /// Get time since last invalidation
    pub fn time_since_invalidation(&self) -> Duration {
        self.last_invalidation.read().elapsed()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_aligned_vector_creation() {
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let aligned = AlignedVector::new(data.clone());

        assert_eq!(aligned.len(), 4);
        assert_eq!(aligned.as_slice(), &data[..]);
    }

    #[test]
    fn test_aligned_vector_alignment() {
        let aligned = AlignedVector::zeros(100);

        // Check that the AlignedVector struct itself is aligned to 64 bytes
        // Note: The inner Vec's data is heap-allocated and uses standard allocator alignment
        assert_eq!(
            std::mem::align_of::<AlignedVector>(),
            64,
            "AlignedVector struct should be aligned to 64 bytes"
        );

        // Verify the Vec data pointer has at least the natural alignment for f32
        let ptr = aligned.as_slice().as_ptr() as usize;
        assert_eq!(
            ptr % std::mem::align_of::<f32>(),
            0,
            "Data pointer should be properly aligned for f32"
        );
    }

    #[test]
    fn test_hot_cache_basic() {
        let cache = HotEmbeddingCache::new(10);

        // Insert some vectors
        cache.insert("key1".to_string(), vec![1.0, 2.0, 3.0]);
        cache.insert("key2".to_string(), vec![4.0, 5.0, 6.0]);

        // Test retrieval
        let vec1 = cache.get("key1").unwrap();
        assert_eq!(vec1.as_slice(), &[1.0, 2.0, 3.0]);

        let vec2 = cache.get("key2").unwrap();
        assert_eq!(vec2.as_slice(), &[4.0, 5.0, 6.0]);

        // Test miss
        assert!(cache.get("key3").is_none());
    }

    #[test]
    fn test_hot_cache_stats() {
        let cache = HotEmbeddingCache::new(10);

        cache.insert("key1".to_string(), vec![1.0, 2.0, 3.0]);

        // One hit
        cache.get("key1");
        // One miss
        cache.get("key2");

        let stats = cache.stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
        assert_eq!(stats.hit_rate, 0.5);
    }

    #[test]
    fn test_hot_cache_lru() {
        let cache = HotEmbeddingCache::new(2);

        cache.insert("key1".to_string(), vec![1.0]);
        cache.insert("key2".to_string(), vec![2.0]);
        cache.insert("key3".to_string(), vec![3.0]); // Should evict key1

        assert!(cache.get("key1").is_none());
        assert!(cache.get("key2").is_some());
        assert!(cache.get("key3").is_some());
    }

    #[test]
    fn test_adaptive_strategy() {
        let strategy = AdaptiveCacheStrategy::new(100, 1000, 0.8);

        let initial_size = strategy.target_size();
        assert_eq!(initial_size, 550); // (100 + 1000) / 2

        // Low hit rate should increase size
        let new_size = strategy.adjust(0.5);
        assert!(new_size > initial_size);

        // High hit rate should decrease size
        strategy.reset();
        let new_size = strategy.adjust(0.95);
        assert!(new_size < initial_size);
    }

    #[test]
    fn test_cache_invalidator_ttl() {
        let invalidator = CacheInvalidator::new(InvalidationPolicy::TTL(Duration::from_millis(10)));

        assert!(!invalidator.should_invalidate());

        std::thread::sleep(Duration::from_millis(15));

        assert!(invalidator.should_invalidate());
    }

    #[test]
    fn test_cache_invalidator_never() {
        let invalidator = CacheInvalidator::new(InvalidationPolicy::Never);

        std::thread::sleep(Duration::from_millis(10));

        assert!(!invalidator.should_invalidate());
    }

    #[test]
    fn test_hot_keys_tracking() {
        let cache = HotEmbeddingCache::new(10);

        cache.insert("key1".to_string(), vec![1.0]);
        cache.insert("key2".to_string(), vec![2.0]);
        cache.insert("key3".to_string(), vec![3.0]);

        // Add a small delay to allow time_in_cache to accumulate
        std::thread::sleep(Duration::from_millis(1));

        // Access key1 multiple times
        for _ in 0..5 {
            cache.get("key1");
        }

        // Access key2 a few times
        for _ in 0..2 {
            cache.get("key2");
        }

        let hot_keys = cache.get_hot_keys(2);
        assert_eq!(hot_keys.len(), 2);
        // key1 should be in the hot keys (either first or second)
        assert!(hot_keys.contains(&"key1".to_string()));
        assert!(hot_keys.contains(&"key2".to_string()));
    }
}