poimen-memory/crates/mem-cli/src/cache_alignment.rs

/// Phase 6: Cache Alignment & KV Cache Optimization
///
/// Optimize memory system for LLM KV cache efficiency:
/// 1. Cache metrics tracking (hit ratio, evictions)
/// 2. Wiki-link ordering by cache locality
/// 3. Chunk pre-loading for hot paths
/// 4. Monitor KV cache hit ratio during retrieval
///
/// Target:
/// - KV cache hit ratio > 80%
/// - Chunk loading latency < 50ms (cache) vs 200ms (disk)
/// - Reduce context recomputation by 60%

use anyhow::Result;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex};

/// Cache metrics for tracking
#[derive(Debug, Clone, Copy)]
pub struct CacheMetrics {
    pub hits: u64,
    pub misses: u64,
    pub evictions: u64,
    pub avg_load_ms: f32,
}

impl CacheMetrics {
    pub fn new() -> Self {
        Self {
            hits: 0,
            misses: 0,
            evictions: 0,
            avg_load_ms: 0.0,
        }
    }

    pub fn hit_ratio(&self) -> f32 {
        let total = self.hits + self.misses;
        if total == 0 {
            0.0
        } else {
            self.hits as f32 / total as f32
        }
    }

    pub fn total_requests(&self) -> u64 {
        self.hits + self.misses
    }
}

/// Chunk with cache locality info
#[derive(Debug, Clone)]
pub struct CachedChunk {
    pub chunk_id: String,
    pub text: String,
    pub score: f32,
    pub cache_distance: u32,      // Hops from root in wiki-graph
    pub access_count: u64,
    pub last_accessed_slot: u32,  // Transformer position in context
}

/// LRU Cache for chunks
pub struct LruChunkCache {
    capacity: usize,
    cache: Arc<Mutex<HashMap<String, (String, u64)>>>, // id -> (text, access_time)
    access_queue: Arc<Mutex<VecDeque<String>>>,
    metrics: Arc<Mutex<CacheMetrics>>,
}

impl LruChunkCache {
    pub fn new(capacity: usize) -> Self {
        Self {
            capacity,
            cache: Arc::new(Mutex::new(HashMap::new())),
            access_queue: Arc::new(Mutex::new(VecDeque::new())),
            metrics: Arc::new(Mutex::new(CacheMetrics::new())),
        }
    }

    /// Get chunk from cache
    pub fn get(&self, chunk_id: &str) -> Option<String> {
        let mut cache = self.cache.lock().unwrap();
        let mut metrics = self.metrics.lock().unwrap();

        if let Some((text, _)) = cache.get_mut(chunk_id) {
            metrics.hits += 1;
            let result = text.clone();

            // Update access tracking (move to end of queue)
            let mut queue = self.access_queue.lock().unwrap();
            queue.retain(|id| id != chunk_id);
            queue.push_back(chunk_id.to_string());

            Some(result)
        } else {
            metrics.misses += 1;
            None
        }
    }

    /// Put chunk in cache with LRU eviction
    pub fn put(&self, chunk_id: &str, text: &str) -> Option<String> {
        let mut cache = self.cache.lock().unwrap();
        let mut queue = self.access_queue.lock().unwrap();
        let mut metrics = self.metrics.lock().unwrap();

        // If cache is full, evict LRU item
        let evicted = if cache.len() >= self.capacity {
            if let Some(lru_id) = queue.pop_front() {
                metrics.evictions += 1;
                cache.remove(&lru_id)
            } else {
                None
            }
        } else {
            None
        };

        // Insert new chunk
        cache.insert(chunk_id.to_string(), (text.to_string(), 0u64));
        queue.push_back(chunk_id.to_string());

        evicted.map(|(text, _)| text)
    }

    pub fn metrics(&self) -> CacheMetrics {
        *self.metrics.lock().unwrap()
    }

    pub fn clear(&self) {
        self.cache.lock().unwrap().clear();
        self.access_queue.lock().unwrap().clear();
    }
}

/// Cache Locality Analyzer
pub struct CacheLocalityAnalyzer;

impl CacheLocalityAnalyzer {
    /// Calculate cache distance (hops from root)
    pub fn calculate_distance(
        chunk_id: &str,
        root_id: &str,
        graph: &HashMap<String, Vec<String>>,
    ) -> u32 {
        if chunk_id == root_id {
            return 0;
        }

        let mut visited = std::collections::HashSet::new();
        let mut queue = VecDeque::new();

        queue.push_back((root_id.to_string(), 0u32));
        visited.insert(root_id.to_string());

        while let Some((current, distance)) = queue.pop_front() {
            if current == chunk_id {
                return distance;
            }

            if let Some(neighbors) = graph.get(&current) {
                for neighbor in neighbors {
                    if !visited.contains(neighbor) {
                        visited.insert(neighbor.clone());
                        queue.push_back((neighbor.clone(), distance + 1));
                    }
                }
            }
        }

        u32::MAX // Unreachable
    }

    /// Order chunks by cache locality (closest first)
    pub fn order_by_locality(
        chunks: Vec<CachedChunk>,
        root_id: &str,
        graph: &HashMap<String, Vec<String>>,
    ) -> Vec<CachedChunk> {
        let mut ordered = chunks;

        ordered.sort_by_key(|c| {
            Self::calculate_distance(&c.chunk_id, root_id, graph)
        });

        ordered
    }
}

/// KV Cache Alignment Optimizer
pub struct KvCacheAligner {
    context_window: usize,           // Max tokens per context
    chunk_avg_tokens: usize,         // Average tokens per chunk
    cache: Arc<LruChunkCache>,
}

impl KvCacheAligner {
    pub fn new(context_window: usize, chunk_avg_tokens: usize, cache_size: usize) -> Self {
        Self {
            context_window,
            chunk_avg_tokens,
            cache: Arc::new(LruChunkCache::new(cache_size)),
        }
    }

    /// Predict if chunk will fit in context window
    pub fn will_fit(&self, chunk_count: usize) -> bool {
        chunk_count * self.chunk_avg_tokens <= self.context_window
    }

    /// Calculate tokens used by chunks
    pub fn calculate_tokens(&self, chunks: &[CachedChunk]) -> usize {
        chunks.len() * self.chunk_avg_tokens
    }

    /// Assign slot positions in context (for cache locality)
    pub fn assign_slots(&self, chunks: &[CachedChunk]) -> Vec<(String, u32)> {
        chunks
            .iter()
            .enumerate()
            .map(|(i, chunk)| (chunk.chunk_id.clone(), i as u32))
            .collect()
    }

    /// Pre-load hot chunks into cache
    pub fn preload_hot_chunks(&self, hot_chunks: Vec<(&str, &str)>) -> Result<()> {
        for (chunk_id, text) in hot_chunks {
            self.cache.put(chunk_id, text);
        }
        Ok(())
    }

    pub fn get_metrics(&self) -> CacheMetrics {
        self.cache.metrics()
    }

    pub fn get_from_cache(&self, chunk_id: &str) -> Option<String> {
        self.cache.get(chunk_id)
    }
}


/// Retrieval Timing Profiler
pub struct RetrievalProfiler {
    timings: Arc<Mutex<Vec<(String, u64)>>>, // (stage_name, duration_ms)
}

impl RetrievalProfiler {
    pub fn new() -> Self {
        Self {
            timings: Arc::new(Mutex::new(Vec::new())),
        }
    }

    pub fn record(&self, stage: &str, duration_ms: u64) {
        self.timings
            .lock()
            .unwrap()
            .push((stage.to_string(), duration_ms));
    }

    pub fn summary(&self) -> Vec<(String, u64)> {
        self.timings.lock().unwrap().clone()
    }

    pub fn total_time(&self) -> u64 {
        self.timings
            .lock()
            .unwrap()
            .iter()
            .map(|(_, duration)| duration)
            .sum()
    }

    pub fn stage_time(&self, stage: &str) -> Option<u64> {
        self.timings
            .lock()
            .unwrap()
            .iter()
            .find(|(s, _)| s == stage)
            .map(|(_, d)| *d)
    }

    pub fn clear(&self) {
        self.timings.lock().unwrap().clear();
    }
}

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

    #[test]
    fn test_cache_metrics_hit_ratio() {
        let mut metrics = CacheMetrics::new();
        metrics.hits = 80;
        metrics.misses = 20;
        assert_eq!(metrics.hit_ratio(), 0.8);
    }

    #[test]
    fn test_lru_cache_get_hit() {
        let cache = LruChunkCache::new(10);
        cache.put("chunk1", "content1");

        let result = cache.get("chunk1");
        assert_eq!(result, Some("content1".to_string()));
        assert_eq!(cache.metrics().hits, 1);
    }

    #[test]
    fn test_lru_cache_get_miss() {
        let cache = LruChunkCache::new(10);
        let result = cache.get("nonexistent");
        assert_eq!(result, None);
        assert_eq!(cache.metrics().misses, 1);
    }

    #[test]
    fn test_lru_cache_eviction() {
        let cache = LruChunkCache::new(2);
        cache.put("chunk1", "content1");
        cache.put("chunk2", "content2");
        cache.put("chunk3", "content3"); // Should evict chunk1

        assert_eq!(cache.get("chunk1"), None); // Evicted
        assert_eq!(cache.get("chunk2"), Some("content2".to_string()));
        assert_eq!(cache.get("chunk3"), Some("content3".to_string()));
        assert_eq!(cache.metrics().evictions, 1);
    }

    #[test]
    fn test_cache_locality_distance_root() {
        let mut graph = HashMap::new();
        let distance = CacheLocalityAnalyzer::calculate_distance("root", "root", &graph);
        assert_eq!(distance, 0);
    }

    #[test]
    fn test_cache_locality_distance_neighbors() {
        let mut graph = HashMap::new();
        graph.insert("root".to_string(), vec!["child1".to_string()]);
        graph.insert("child1".to_string(), vec!["child2".to_string()]);

        let dist_child1 = CacheLocalityAnalyzer::calculate_distance("child1", "root", &graph);
        let dist_child2 = CacheLocalityAnalyzer::calculate_distance("child2", "root", &graph);

        assert_eq!(dist_child1, 1);
        assert_eq!(dist_child2, 2);
    }

    #[test]
    fn test_cache_locality_ordering() {
        let mut graph = HashMap::new();
        graph.insert("root".to_string(), vec!["near".to_string(), "far".to_string()]);
        graph.insert("far".to_string(), vec!["farther".to_string()]);

        let chunks = vec![
            CachedChunk {
                chunk_id: "farther".to_string(),
                text: "".to_string(),
                score: 0.9,
                cache_distance: u32::MAX,
                access_count: 1,
                last_accessed_slot: 0,
            },
            CachedChunk {
                chunk_id: "near".to_string(),
                text: "".to_string(),
                score: 0.8,
                cache_distance: 1,
                access_count: 1,
                last_accessed_slot: 0,
            },
        ];

        let ordered = CacheLocalityAnalyzer::order_by_locality(chunks, "root", &graph);
        assert_eq!(ordered[0].chunk_id, "near"); // Closest first
    }

    #[test]
    fn test_kv_cache_will_fit() {
        let aligner = KvCacheAligner::new(1000, 100, 10);
        assert!(aligner.will_fit(5)); // 500 tokens < 1000
        assert!(!aligner.will_fit(15)); // 1500 tokens > 1000
    }

    #[test]
    fn test_kv_cache_assign_slots() {
        let aligner = KvCacheAligner::new(1000, 100, 10);
        let chunks = vec![
            CachedChunk {
                chunk_id: "chunk1".to_string(),
                text: "".to_string(),
                score: 0.9,
                cache_distance: 0,
                access_count: 1,
                last_accessed_slot: 0,
            },
            CachedChunk {
                chunk_id: "chunk2".to_string(),
                text: "".to_string(),
                score: 0.8,
                cache_distance: 1,
                access_count: 1,
                last_accessed_slot: 0,
            },
        ];

        let slots = aligner.assign_slots(&chunks);
        assert_eq!(slots[0], ("chunk1".to_string(), 0));
        assert_eq!(slots[1], ("chunk2".to_string(), 1));
    }

    #[test]
    fn test_preload_hot_chunks() {
        let aligner = KvCacheAligner::new(1000, 100, 10);
        let hot_chunks = vec![("chunk1", "content1"), ("chunk2", "content2")];

        aligner.preload_hot_chunks(hot_chunks).unwrap();

        // Verify that preloaded chunks are in cache by retrieving them
        // This will increment metrics
        assert_eq!(aligner.get_from_cache("chunk1"), Some("content1".to_string()));
        assert_eq!(aligner.get_from_cache("chunk2"), Some("content2".to_string()));

        let metrics = aligner.get_metrics();
        assert!(metrics.total_requests() >= 2);
        assert!(metrics.hits >= 2);
    }

    #[test]
    fn test_retrieval_profiler_record() {
        let profiler = RetrievalProfiler::new();
        profiler.record("wiki_scope", 10);
        profiler.record("tfidf_filter", 50);
        profiler.record("semantic_rerank", 100);

        let summary = profiler.summary();
        assert_eq!(summary.len(), 3);
        assert_eq!(profiler.total_time(), 160);
    }

    #[test]
    fn test_retrieval_profiler_stage_time() {
        let profiler = RetrievalProfiler::new();
        profiler.record("wiki_scope", 10);
        profiler.record("semantic_rerank", 100);

        assert_eq!(profiler.stage_time("wiki_scope"), Some(10));
        assert_eq!(profiler.stage_time("nonexistent"), None);
    }

    #[test]
    fn test_cache_metrics_total_requests() {
        let mut metrics = CacheMetrics::new();
        metrics.hits = 60;
        metrics.misses = 40;
        assert_eq!(metrics.total_requests(), 100);
    }
}