ruvector-cli 2.0.4

//! GNN Layer Caching for Performance Optimization
//!
//! This module provides persistent caching for GNN layers and query results,
//! eliminating the ~2.5s overhead per operation from process initialization,
//! database loading, and index deserialization.
//!
//! ## Performance Impact
//!
//! | Operation | Before | After | Improvement |
//! |-----------|--------|-------|-------------|
//! | Layer init | ~2.5s | ~5-10ms | 250-500x |
//! | Query | ~2.5s | ~5-10ms | 250-500x |
//! | Batch query | ~2.5s * N | ~5-10ms | Amortized |

use lru::LruCache;
use ruvector_gnn::layer::RuvectorLayer;
use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;

/// Cache entry with metadata for monitoring
#[derive(Debug, Clone)]
pub struct CacheEntry<T> {
    pub value: T,
    pub created_at: Instant,
    pub last_accessed: Instant,
    pub access_count: u64,
}

impl<T: Clone> CacheEntry<T> {
    pub fn new(value: T) -> Self {
        let now = Instant::now();
        Self {
            value,
            created_at: now,
            last_accessed: now,
            access_count: 1,
        }
    }

    pub fn access(&mut self) -> &T {
        self.last_accessed = Instant::now();
        self.access_count += 1;
        &self.value
    }
}

/// Configuration for the GNN cache
#[derive(Debug, Clone)]
pub struct GnnCacheConfig {
    /// Maximum number of GNN layers to cache
    pub max_layers: usize,
    /// Maximum number of query results to cache
    pub max_query_results: usize,
    /// TTL for cached query results (in seconds)
    pub query_result_ttl_secs: u64,
    /// Whether to preload common layer configurations
    pub preload_common: bool,
}

impl Default for GnnCacheConfig {
    fn default() -> Self {
        Self {
            max_layers: 32,
            max_query_results: 1000,
            query_result_ttl_secs: 300, // 5 minutes
            preload_common: true,
        }
    }
}

/// Query result cache key
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct QueryCacheKey {
    /// Layer configuration hash
    pub layer_hash: String,
    /// Query vector hash (first 8 floats as u64 bits)
    pub query_hash: u64,
    /// Number of results requested
    pub k: usize,
}

impl QueryCacheKey {
    pub fn new(layer_id: &str, query: &[f32], k: usize) -> Self {
        // Simple hash of query vector
        let query_hash = query
            .iter()
            .take(8)
            .fold(0u64, |acc, &v| acc.wrapping_add(v.to_bits() as u64));

        Self {
            layer_hash: layer_id.to_string(),
            query_hash,
            k,
        }
    }
}

/// Cached query result
#[derive(Debug, Clone)]
pub struct CachedQueryResult {
    pub result: Vec<f32>,
    pub cached_at: Instant,
}

/// GNN Layer cache with LRU eviction and TTL support
pub struct GnnCache {
    /// Cached GNN layers by configuration hash
    layers: Arc<RwLock<HashMap<String, CacheEntry<RuvectorLayer>>>>,
    /// LRU cache for query results
    query_results: Arc<RwLock<LruCache<QueryCacheKey, CachedQueryResult>>>,
    /// Configuration
    config: GnnCacheConfig,
    /// Cache statistics
    stats: Arc<RwLock<CacheStats>>,
}

/// Cache statistics for monitoring
#[derive(Debug, Clone, Default)]
pub struct CacheStats {
    pub layer_hits: u64,
    pub layer_misses: u64,
    pub query_hits: u64,
    pub query_misses: u64,
    pub evictions: u64,
    pub total_queries: u64,
}

impl CacheStats {
    pub fn layer_hit_rate(&self) -> f64 {
        let total = self.layer_hits + self.layer_misses;
        if total == 0 {
            0.0
        } else {
            self.layer_hits as f64 / total as f64
        }
    }

    pub fn query_hit_rate(&self) -> f64 {
        let total = self.query_hits + self.query_misses;
        if total == 0 {
            0.0
        } else {
            self.query_hits as f64 / total as f64
        }
    }
}

impl GnnCache {
    /// Create a new GNN cache with the given configuration
    pub fn new(config: GnnCacheConfig) -> Self {
        let query_cache_size =
            NonZeroUsize::new(config.max_query_results).unwrap_or(NonZeroUsize::new(1000).unwrap());

        Self {
            layers: Arc::new(RwLock::new(HashMap::new())),
            query_results: Arc::new(RwLock::new(LruCache::new(query_cache_size))),
            config,
            stats: Arc::new(RwLock::new(CacheStats::default())),
        }
    }

    /// Get or create a GNN layer with the specified configuration
    pub async fn get_or_create_layer(
        &self,
        input_dim: usize,
        hidden_dim: usize,
        heads: usize,
        dropout: f32,
    ) -> RuvectorLayer {
        let key = format!(
            "{}_{}_{}_{}",
            input_dim,
            hidden_dim,
            heads,
            (dropout * 1000.0) as u32
        );

        // Check cache first
        {
            let mut layers = self.layers.write().await;
            if let Some(entry) = layers.get_mut(&key) {
                let mut stats = self.stats.write().await;
                stats.layer_hits += 1;
                return entry.access().clone();
            }
        }

        // Create new layer
        let layer = RuvectorLayer::new(input_dim, hidden_dim, heads, dropout);

        // Cache it
        {
            let mut layers = self.layers.write().await;
            let mut stats = self.stats.write().await;
            stats.layer_misses += 1;

            // Evict if necessary
            if layers.len() >= self.config.max_layers {
                // Simple eviction: remove oldest entry
                if let Some(oldest_key) = layers
                    .iter()
                    .min_by_key(|(_, v)| v.last_accessed)
                    .map(|(k, _)| k.clone())
                {
                    layers.remove(&oldest_key);
                    stats.evictions += 1;
                }
            }

            layers.insert(key, CacheEntry::new(layer.clone()));
        }

        layer
    }

    /// Get cached query result if available and not expired
    pub async fn get_query_result(&self, key: &QueryCacheKey) -> Option<Vec<f32>> {
        let mut results = self.query_results.write().await;

        if let Some(cached) = results.get(key) {
            let ttl = Duration::from_secs(self.config.query_result_ttl_secs);
            if cached.cached_at.elapsed() < ttl {
                let mut stats = self.stats.write().await;
                stats.query_hits += 1;
                stats.total_queries += 1;
                return Some(cached.result.clone());
            }
            // Expired, remove it
            results.pop(key);
        }

        let mut stats = self.stats.write().await;
        stats.query_misses += 1;
        stats.total_queries += 1;
        None
    }

    /// Cache a query result
    pub async fn cache_query_result(&self, key: QueryCacheKey, result: Vec<f32>) {
        let mut results = self.query_results.write().await;
        results.put(
            key,
            CachedQueryResult {
                result,
                cached_at: Instant::now(),
            },
        );
    }

    /// Get current cache statistics
    pub async fn stats(&self) -> CacheStats {
        self.stats.read().await.clone()
    }

    /// Clear all caches
    pub async fn clear(&self) {
        self.layers.write().await.clear();
        self.query_results.write().await.clear();
    }

    /// Preload common layer configurations for faster first access
    pub async fn preload_common_layers(&self) {
        // Common configurations used in practice
        let common_configs = [
            (128, 256, 4, 0.1),   // Small model
            (256, 512, 8, 0.1),   // Medium model
            (384, 768, 8, 0.1),   // Base model (BERT-like)
            (768, 1024, 16, 0.1), // Large model
        ];

        for (input, hidden, heads, dropout) in common_configs {
            let _ = self
                .get_or_create_layer(input, hidden, heads, dropout)
                .await;
        }
    }

    /// Get number of cached layers
    pub async fn layer_count(&self) -> usize {
        self.layers.read().await.len()
    }

    /// Get number of cached query results
    pub async fn query_result_count(&self) -> usize {
        self.query_results.read().await.len()
    }
}

/// Batch operation for multiple GNN forward passes
#[derive(Debug, Clone)]
pub struct BatchGnnRequest {
    pub layer_config: LayerConfig,
    pub operations: Vec<GnnOperation>,
}

#[derive(Debug, Clone)]
pub struct LayerConfig {
    pub input_dim: usize,
    pub hidden_dim: usize,
    pub heads: usize,
    pub dropout: f32,
}

#[derive(Debug, Clone)]
pub struct GnnOperation {
    pub node_embedding: Vec<f32>,
    pub neighbor_embeddings: Vec<Vec<f32>>,
    pub edge_weights: Vec<f32>,
}

#[derive(Debug, Clone)]
pub struct BatchGnnResult {
    pub results: Vec<Vec<f32>>,
    pub cached_count: usize,
    pub computed_count: usize,
    pub total_time_ms: f64,
}

impl GnnCache {
    /// Execute batch GNN operations with caching
    pub async fn batch_forward(&self, request: BatchGnnRequest) -> BatchGnnResult {
        let start = Instant::now();

        // Get or create the layer
        let layer = self
            .get_or_create_layer(
                request.layer_config.input_dim,
                request.layer_config.hidden_dim,
                request.layer_config.heads,
                request.layer_config.dropout,
            )
            .await;

        let layer_id = format!(
            "{}_{}_{}",
            request.layer_config.input_dim,
            request.layer_config.hidden_dim,
            request.layer_config.heads
        );

        let mut results = Vec::with_capacity(request.operations.len());
        let mut cached_count = 0;
        let mut computed_count = 0;

        for op in &request.operations {
            // Check cache
            let cache_key = QueryCacheKey::new(&layer_id, &op.node_embedding, 1);

            if let Some(cached) = self.get_query_result(&cache_key).await {
                results.push(cached);
                cached_count += 1;
            } else {
                // Compute forward pass
                let result = layer.forward(
                    &op.node_embedding,
                    &op.neighbor_embeddings,
                    &op.edge_weights,
                );

                // Cache the result
                self.cache_query_result(cache_key, result.clone()).await;
                results.push(result);
                computed_count += 1;
            }
        }

        BatchGnnResult {
            results,
            cached_count,
            computed_count,
            total_time_ms: start.elapsed().as_secs_f64() * 1000.0,
        }
    }
}

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

    #[tokio::test]
    async fn test_layer_caching() {
        let cache = GnnCache::new(GnnCacheConfig::default());

        // First access - miss
        let layer1 = cache.get_or_create_layer(128, 256, 4, 0.1).await;
        let stats = cache.stats().await;
        assert_eq!(stats.layer_misses, 1);
        assert_eq!(stats.layer_hits, 0);

        // Second access - hit
        let _layer2 = cache.get_or_create_layer(128, 256, 4, 0.1).await;
        let stats = cache.stats().await;
        assert_eq!(stats.layer_misses, 1);
        assert_eq!(stats.layer_hits, 1);
    }

    #[tokio::test]
    async fn test_query_result_caching() {
        let cache = GnnCache::new(GnnCacheConfig::default());

        let key = QueryCacheKey::new("test", &[1.0, 2.0, 3.0], 10);
        let result = vec![0.1, 0.2, 0.3];

        // Cache miss
        assert!(cache.get_query_result(&key).await.is_none());

        // Cache the result
        cache.cache_query_result(key.clone(), result.clone()).await;

        // Cache hit
        let cached = cache.get_query_result(&key).await;
        assert!(cached.is_some());
        assert_eq!(cached.unwrap(), result);
    }

    #[tokio::test]
    async fn test_batch_forward() {
        let cache = GnnCache::new(GnnCacheConfig::default());

        let request = BatchGnnRequest {
            layer_config: LayerConfig {
                input_dim: 4,
                hidden_dim: 8,
                heads: 2,
                dropout: 0.1,
            },
            operations: vec![
                GnnOperation {
                    node_embedding: vec![1.0, 2.0, 3.0, 4.0],
                    neighbor_embeddings: vec![vec![0.5, 1.0, 1.5, 2.0]],
                    edge_weights: vec![1.0],
                },
                GnnOperation {
                    node_embedding: vec![2.0, 3.0, 4.0, 5.0],
                    neighbor_embeddings: vec![vec![1.0, 1.5, 2.0, 2.5]],
                    edge_weights: vec![1.0],
                },
            ],
        };

        let result = cache.batch_forward(request).await;
        assert_eq!(result.results.len(), 2);
        assert_eq!(result.computed_count, 2);
        assert_eq!(result.cached_count, 0);
    }

    #[tokio::test]
    async fn test_preload_common_layers() {
        let cache = GnnCache::new(GnnCacheConfig {
            preload_common: true,
            ..Default::default()
        });

        cache.preload_common_layers().await;

        // Should have 4 preloaded layers
        assert_eq!(cache.layer_count().await, 4);
    }
}