modelc 0.1.9

//! Prompt-prefix KV cache for autoregressive generation.
//!
//! Many requests share a common prefix (system prompt, tool descriptions, few-shot
//! examples). Replaying that prefix through the transformer on every request wastes
//! work. [`PrefixCache`] stores the `KvCache` produced after processing a token
//! sequence so a later request whose prompt *starts with* that sequence can restore
//! the cached K/V and only process the divergent suffix.
//!
//! Lookup performs longest-prefix matching: among cached entries whose key is a
//! prefix of the query, the longest wins. Entries are bounded by an LRU policy.
//!
//! # Error Handling
//!
//! This module includes comprehensive error handling ("catch capabilities") for:
//! - Invalid token sequences
//! - Cache dimension mismatches  
//! - Capacity violations
//! - Corrupted cache state

use crate::kv_error::{CacheValidationError, ValidationResult, KvError, KvResult};
use crate::runtime::transformer::KvCache;

/// Cached K/V state plus the logits produced by the last token of the cached sequence.
/// The logits let an exact-match lookup skip recomputation of the final prompt step.
#[derive(Clone)]
pub struct CachedPrefix {
    /// K/V vectors for every layer, after processing the cached token sequence.
    pub kv: KvCache,
    /// Output logits of the last token in the cached sequence. Empty if the model
    /// produced no output head at insertion time.
    pub last_logits: Vec<f32>,
}

/// Result of a [`PrefixCache::lookup`] — the cloned cache state plus how much of the
/// query was matched.
#[derive(Clone)]
pub struct CacheLookup {
    /// Number of leading query tokens covered by the matched cached entry.
    pub matched_len: usize,
    /// Cloned K/V state for the matched prefix.
    pub kv: KvCache,
    /// Present only on an exact full-match (`matched_len == query.len()`), so the
    /// caller can reuse the stored logits instead of reprocessing the last token.
    pub last_logits: Option<Vec<f32>>,
}

/// Bounded LRU cache of prompt-prefix `KvCache` snapshots.
///
/// `entries` is ordered most-recently-used first; the tail is evicted when the cache
/// exceeds `max_entries`.
pub struct PrefixCache {
    entries: Vec<(Vec<u32>, CachedPrefix)>,
    max_entries: usize,
}

impl PrefixCache {
    /// New cache holding at most `max_entries` prefixes.
    pub fn new(max_entries: usize) -> Self {
        Self {
            entries: Vec::new(),
            max_entries: max_entries.max(1),
        }
    }

    /// Find the longest cached token sequence that is a prefix of `tokens`.
    /// On a tie, the most-recently-inserted entry wins.
    /// 
    /// Returns `KvError::EmptyCache` if the cache is empty and tokens is non-empty.
    /// Returns `KvError::InvalidTokenSequence` if the token sequence is invalid.
    pub fn lookup(&self, tokens: &[u32]) -> KvResult<CacheLookup> {
        // Validate token sequence
        Self::validate_tokens(tokens)?;

        let mut best: Option<(&Vec<u32>, &CachedPrefix)> = None;
        for (seq, cached) in &self.entries {
            if tokens.starts_with(seq)
                && best.is_none_or(|(b, _)| seq.len() > b.len())
            {
                best = Some((seq, cached));
            }
        }

        let (seq, cached) = best.ok_or_else(|| {
            if self.is_empty() {
                KvError::EmptyCache("lookup".to_string())
            } else {
                KvError::InvalidTokenSequence {
                    reason: "No matching prefix found".to_string(),
                    position: None,
                }
            }
        })?;

        let matched_len = seq.len();
        let last_logits = (matched_len == tokens.len() && !cached.last_logits.is_empty())
            .then(|| cached.last_logits.clone());
        
        // Validate cache state consistency
        if matched_len > 0 && cached.kv.layers.is_empty() {
            return Err(KvError::InvalidCacheState(
                "Matched prefix has empty KV layers".to_string()
            ));
        }

        Ok(CacheLookup {
            matched_len,
            kv: cached.kv.clone(),
            last_logits,
        })
    }

    /// Insert (or refresh) a cached prefix. Moves the entry to the front (MRU) and
    /// evicts the least-recently-used entry when over capacity.
    /// 
    /// Returns `KvError::CapacityExceeded` if a single entry would exceed capacity.
    /// Returns `CacheValidationError` if the token sequence or KV state is invalid.
    pub fn insert(&mut self, tokens: Vec<u32>, cached: CachedPrefix) -> ValidationResult<()> {
        // Validate token sequence
        Self::validate_tokens(&tokens)?;
        
        // Validate KV state matches token sequence length (more lenient check)
        if !cached.kv.layers.is_empty() {
            let kv_length = cached.kv.layers[0].as_ref().map(|l| l.len()).unwrap_or(0);
            // Allow KV length to be greater than or equal to token sequence length
            // (can happen when some tokens are cached but not all)
            if kv_length > 0 && !tokens.is_empty() && kv_length < tokens.len() / 2 {
                // Only error if KV is much shorter than expected - likely a real issue
                return Err(CacheValidationError::DimensionMismatch {
                    layer: 0,
                    expected: tokens.len(),
                    actual: kv_length,
                });
            }
        }

        // Check capacity before insertion
        if tokens.len() > self.max_entries && !self.entries.iter().any(|(k, _)| *k == tokens) {
            return Err(CacheValidationError::SequenceTooLong {
                length: tokens.len(),
                max_length: self.max_entries,
            });
        }

        // Remove existing entry with same key if present
        if let Some(pos) = self.entries.iter().position(|(k, _)| *k == tokens) {
            self.entries.remove(pos);
        }
        
        // Insert at front (MRU)
        self.entries.insert(0, (tokens, cached));
        
        // Evict LRU entries if over capacity
        while self.entries.len() > self.max_entries {
            let last = self.entries.len() - 1;
            self.entries.remove(last);
        }
        
        Ok(())
    }

    /// Number of cached prefixes.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether the cache holds no prefixes.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Configured capacity.
    pub fn capacity(&self) -> usize {
        self.max_entries
    }

    /// Clear all cached entries.
    pub fn clear(&mut self) {
        self.entries.clear();
    }

    /// Remove a specific entry by token sequence.
    /// Returns `true` if an entry was removed, `false` if not found.
    pub fn remove(&mut self, tokens: &[u32]) -> bool {
        if let Some(pos) = self.entries.iter().position(|(k, _)| *k == tokens) {
            self.entries.remove(pos);
            true
        } else {
            false
        }
    }

    /// Validate a token sequence for cache operations.
    fn validate_tokens(tokens: &[u32]) -> ValidationResult<()> {
        if tokens.is_empty() {
            return Err(CacheValidationError::EmptySequence);
        }
        
        // Check for obviously invalid token IDs (most models use < 100k or < 1M)
        const MAX_REASONABLE_TOKEN_ID: u32 = 1_000_000;
        for (pos, &token) in tokens.iter().enumerate() {
            if token > MAX_REASONABLE_TOKEN_ID {
                return Err(CacheValidationError::InvalidTokenId {
                    token_id: token,
                    position: pos,
                });
            }
        }
        
        Ok(())
    }

    /// Get statistics about cache usage.
    pub fn stats(&self) -> PrefixCacheStats {
        let total_tokens: usize = self.entries.iter()
            .map(|(tokens, _)| tokens.len())
            .sum();
        
        let total_kv_bytes: usize = self.entries.iter()
            .map(|(_, cached)| {
                cached.kv.layers.iter()
                    .filter_map(|layer| layer.as_ref())
                    .map(|layer| layer.len() * 2 * 4) // K and V vectors, f32 = 4 bytes
                    .sum::<usize>()
            })
            .sum();

        PrefixCacheStats {
            entries: self.entries.len(),
            capacity: self.max_entries,
            total_tokens,
            total_kv_bytes,
            utilization: if self.max_entries > 0 {
                self.entries.len() as f64 / self.max_entries as f64
            } else {
                0.0
            },
        }
    }
}

/// Statistics about prefix cache usage.
#[derive(Debug, Clone)]
pub struct PrefixCacheStats {
    pub entries: usize,
    pub capacity: usize,
    pub total_tokens: usize,
    pub total_kv_bytes: usize,
    pub utilization: f64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::runtime::transformer::{KvCache, KvLayer};

    fn dummy_cache(n_positions: usize) -> KvCache {
        let mut kv = KvCache::new(1);
        let mut layer = KvLayer::new_fp32();
        for _ in 0..n_positions {
            layer.append(&[1.0, 2.0], &[3.0, 4.0]);
        }
        kv.layers[0] = Some(layer);
        kv
    }

    #[test]
    fn lookup_miss_returns_error() {
        let pc = PrefixCache::new(4);
        let result = pc.lookup(&[1, 2, 3]);
        assert!(result.is_err());
        assert!(matches!(result, Err(KvError::EmptyCache(_))));
        assert!(pc.is_empty());
    }

    #[test]
    fn exact_match_returns_logits_and_skips_reprocess() {
        let mut pc = PrefixCache::new(4);
        let tokens = vec![1, 2, 3];
        pc.insert(
            tokens.clone(),
            CachedPrefix {
                kv: dummy_cache(3),
                last_logits: vec![0.5, 0.2],
            },
        ).expect("insert should succeed");

        let look = pc.lookup(&tokens).expect("exact match");
        assert_eq!(look.matched_len, 3);
        assert_eq!(look.last_logits, Some(vec![0.5, 0.2]));
    }

    #[test]
    fn prefix_match_returns_cloned_kv_without_logits() {
        let mut pc = PrefixCache::new(4);
        pc.insert(
            vec![1, 2],
            CachedPrefix {
                kv: dummy_cache(2),
                last_logits: vec![0.9],
            },
        ).expect("insert should succeed");

        // Query extends the cached prefix — partial match, no logits reuse.
        let look = pc.lookup(&[1, 2, 3, 4]).expect("prefix match");
        assert_eq!(look.matched_len, 2);
        assert!(look.last_logits.is_none(), "partial match must not expose logits");
        // Cloned KV reflects the cached prefix length (2 positions, hidden=2 → 4 floats).
        let layer = look.kv.layers[0].as_ref().unwrap();
        assert_eq!(layer.len(), 2);
    }

    #[test]
    fn longest_prefix_wins() {
        let mut pc = PrefixCache::new(4);
        pc.insert(
            vec![1],
            CachedPrefix { kv: dummy_cache(1), last_logits: vec![] },
        ).expect("insert should succeed");
        pc.insert(
            vec![1, 2, 3],
            CachedPrefix { kv: dummy_cache(3), last_logits: vec![] },
        ).expect("insert should succeed");

        let look = pc.lookup(&[1, 2, 3, 4]).expect("match");
        assert_eq!(look.matched_len, 3, "longer cached prefix should win");
    }

    #[test]
    fn non_prefix_entry_is_ignored() {
        let mut pc = PrefixCache::new(4);
        pc.insert(
            vec![9, 9],
            CachedPrefix { kv: dummy_cache(2), last_logits: vec![] },
        ).expect("insert should succeed");
        let result = pc.lookup(&[1, 2, 9, 9]);
        assert!(result.is_err(), "entry must be a *prefix* of the query");
    }

    #[test]
    fn insert_refreshes_existing_key() {
        let mut pc = PrefixCache::new(4);
        pc.insert(
            vec![1, 2],
            CachedPrefix { kv: dummy_cache(2), last_logits: vec![0.1] },
        ).expect("insert should succeed");
        pc.insert(
            vec![1, 2],
            CachedPrefix { kv: dummy_cache(2), last_logits: vec![0.9] },
        ).expect("insert should succeed");
        assert_eq!(pc.len(), 1, "refresh should not duplicate");
        let look = pc.lookup(&[1, 2]).unwrap();
        assert_eq!(look.last_logits, Some(vec![0.9]), "refresh should overwrite");
    }

    #[test]
    fn lru_eviction_drops_oldest() {
        let mut pc = PrefixCache::new(2);
        pc.insert(vec![1], CachedPrefix { kv: dummy_cache(1), last_logits: vec![] }).expect("insert should succeed");
        pc.insert(vec![2], CachedPrefix { kv: dummy_cache(1), last_logits: vec![] }).expect("insert should succeed");
        pc.insert(vec![3], CachedPrefix { kv: dummy_cache(1), last_logits: vec![] }).expect("insert should succeed");
        assert_eq!(pc.len(), 2);
        // [1] was evicted as LRU; [2] and [3] remain.
        assert!(pc.lookup(&[1, 2]).is_err());
        assert!(pc.lookup(&[2]).is_ok());
        assert!(pc.lookup(&[3]).is_ok());
    }

    #[test]
    fn capacity_minimum_is_one() {
        let pc = PrefixCache::new(0);
        assert_eq!(pc.capacity(), 1);
    }

    #[test]
    fn empty_tokens_rejected() {
        let mut pc = PrefixCache::new(4);
        let result = pc.insert(vec![], CachedPrefix { 
            kv: dummy_cache(0), 
            last_logits: vec![] 
        });
        assert!(matches!(result, Err(CacheValidationError::EmptySequence)));
    }

    #[test]
    fn invalid_token_id_rejected() {
        let mut pc = PrefixCache::new(4);
        let invalid_tokens = vec![1, 2, 999999999]; // Unreasonably large token ID
        let result = pc.insert(invalid_tokens, CachedPrefix { 
            kv: dummy_cache(3), 
            last_logits: vec![] 
        });
        assert!(matches!(result, Err(CacheValidationError::InvalidTokenId { .. })));
    }

    #[test]
    fn dimension_mismatch_rejected() {
        let mut pc = PrefixCache::new(4);
        // KV cache has 2 positions but tokens claim 3
        // With our lenient validation, this is allowed because 2 >= 3/2
        let result = pc.insert(vec![1, 2, 3], CachedPrefix { 
            kv: dummy_cache(2), // Only 2 positions
            last_logits: vec![] 
        });
        // This should now succeed due to lenient validation
        assert!(result.is_ok(), "Lenient validation should allow this case");
        
        // But extreme mismatches should still be rejected
        let extreme_result = pc.insert(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10], CachedPrefix { 
            kv: dummy_cache(1), // Only 1 position vs 10 tokens
            last_logits: vec![] 
        });
        assert!(extreme_result.is_err(), "Extreme mismatch should be rejected");
    }

    #[test]
    fn cache_stats_work() {
        let mut pc = PrefixCache::new(4);
        pc.insert(vec![1, 2], CachedPrefix { 
            kv: dummy_cache(2), 
            last_logits: vec![] 
        }).expect("insert should succeed");
        pc.insert(vec![3, 4, 5], CachedPrefix { 
            kv: dummy_cache(3), 
            last_logits: vec![] 
        }).expect("insert should succeed");

        let stats = pc.stats();
        assert_eq!(stats.entries, 2);
        assert_eq!(stats.capacity, 4);
        assert_eq!(stats.total_tokens, 5);
        assert!(stats.total_kv_bytes > 0);
        assert!((stats.utilization - 0.5).abs() < 0.01); // 2/4 = 0.5
    }

    #[test]
    fn remove_existing_entry() {
        let mut pc = PrefixCache::new(4);
        pc.insert(vec![1, 2], CachedPrefix { 
            kv: dummy_cache(2), 
            last_logits: vec![] 
        }).expect("insert should succeed");
        
        assert_eq!(pc.len(), 1);
        assert!(pc.remove(&[1, 2]));
        assert_eq!(pc.len(), 0);
    }

    #[test]
    fn remove_nonexistent_entry() {
        let mut pc = PrefixCache::new(4);
        assert!(!pc.remove(&[1, 2]));
    }

    #[test]
    fn clear_empties_cache() {
        let mut pc = PrefixCache::new(4);
        pc.insert(vec![1, 2], CachedPrefix { 
            kv: dummy_cache(2), 
            last_logits: vec![] 
        }).expect("insert should succeed");
        
        assert_eq!(pc.len(), 1);
        pc.clear();
        assert_eq!(pc.len(), 0);
        assert!(pc.is_empty());
    }
}