katra_cache/

store.rs

//! The content-addressed persistent store.

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use katra_core::{KatraError, Result, fnv1a_parts};
use serde::{Deserialize, Serialize};

/// The layer of a cache entry (Katra3D §18).
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct CacheKey {
    /// Layer (1 DXIL, 2 IR, 3 SPIR-V, 4 pipeline).
    pub layer: u8,
    /// Content hash of the primary input at this layer.
    pub content_hash: u64,
    /// Secondary identity (e.g. PSO descriptor hash, target feature set).
    pub meta_hash: u64,
}

impl CacheKey {
    /// Build a key from parts.
    pub fn new(layer: u8, content_hash: u64, meta_hash: u64) -> Self {
        CacheKey { layer, content_hash, meta_hash }
    }

    /// A stable filename fragment for this key.
    pub fn file_name(&self) -> String {
        format!("{:02x}_{:016x}_{:016x}.bin", self.layer, self.content_hash, self.meta_hash)
    }
}

/// A cached artifact.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CacheEntry {
    /// The key.
    pub key: CacheKey,
    /// Artifact size in bytes.
    pub size: u64,
    /// Creation wall time (ns).
    pub created_wall_ns: u64,
    /// Hit count.
    pub hits: u64,
    /// File path of the artifact.
    pub path: PathBuf,
}

/// Cache metrics.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct CacheMetrics {
    /// Lookups performed.
    pub lookups: u64,
    /// Lookup hits.
    pub hits: u64,
    /// Artifacts inserted.
    pub inserts: u64,
    /// LRU evictions.
    pub evictions: u64,
    /// Entries invalidated.
    pub invalidations: u64,
    /// Bytes stored.
    pub bytes_stored: u64,
}

/// The persistent layered cache.
pub struct CacheStore {
    root: PathBuf,
    max_bytes: u64,
    index: HashMap<CacheKey, CacheEntry>,
    lru: Vec<CacheKey>,
    metrics: CacheMetrics,
}

impl CacheStore {
    /// Open (or create) a cache rooted at `root` with a byte budget.
    /// A corrupt index is tolerated: it is rebuilt empty.
    pub fn open(root: &Path, max_bytes: u64) -> Result<Self> {
        std::fs::create_dir_all(root).map_err(KatraError::from)?;
        let index_path = root.join("index.bin");
        let index: HashMap<CacheKey, CacheEntry> = match std::fs::read(&index_path) {
            Ok(bytes) => bincode::deserialize(&bytes).unwrap_or_default(),
            Err(_) => HashMap::new(),
        };
        let mut store = CacheStore {
            root: root.to_path_buf(),
            max_bytes,
            index,
            lru: Vec::new(),
            metrics: CacheMetrics::default(),
        };
        // Rebuild the LRU order (oldest first by creation time).
        let mut entries: Vec<&CacheEntry> = store.index.values().collect();
        entries.sort_by_key(|e| e.created_wall_ns);
        store.lru = entries.iter().map(|e| e.key).collect();
        store.metrics.bytes_stored = entries.iter().map(|e| e.size).sum();
        Ok(store)
    }

    /// Look up an artifact; loads it from disk and bumps LRU + hits.
    pub fn lookup(&mut self, key: &CacheKey) -> Option<Vec<u8>> {
        self.metrics.lookups += 1;
        let entry = self.index.get(key)?.clone();
        let data = std::fs::read(&entry.path).ok()?;
        self.metrics.hits += 1;
        self.index.get_mut(key).expect("entry present").hits += 1;
        // Move to the back of LRU.
        self.lru.retain(|k| k != key);
        self.lru.push(*key);
        Some(data)
    }

    /// Whether an artifact is present (without loading).
    pub fn contains(&self, key: &CacheKey) -> bool {
        self.index.contains_key(key)
    }

    /// Insert an artifact (atomic temp + rename).
    pub fn insert(&mut self, key: &CacheKey, data: &[u8]) -> Result<()> {
        if let Some(old) = self.index.get(key) {
            // Replace: drop the old file.
            let _ = std::fs::remove_file(&old.path);
            self.metrics.bytes_stored = self.metrics.bytes_stored.saturating_sub(old.size);
            self.lru.retain(|k| k != key);
        }
        let layer_dir = self.root.join(format!("L{}", key.layer));
        std::fs::create_dir_all(&layer_dir).map_err(KatraError::from)?;
        let final_path = layer_dir.join(key.file_name());
        let tmp_path = layer_dir.join(format!("{}.tmp", key.file_name()));
        std::fs::write(&tmp_path, data).map_err(KatraError::from)?;
        std::fs::rename(&tmp_path, &final_path).map_err(KatraError::from)?;

        let entry = CacheEntry {
            key: *key,
            size: data.len() as u64,
            created_wall_ns: katra_core::wall_now_ns(),
            hits: 0,
            path: final_path,
        };
        self.index.insert(*key, entry.clone());
        self.lru.push(*key);
        self.metrics.inserts += 1;
        self.metrics.bytes_stored += entry.size;
        self.evict_lru();
        self.flush_index()?;
        Ok(())
    }

    /// Invalidate entries at `layer >= from_layer` (e.g. driver update →
    /// `invalidate_from(4)` keeps layers 1–3).
    pub fn invalidate_from(&mut self, from_layer: u8) -> usize {
        let doomed: Vec<CacheKey> =
            self.index.keys().filter(|k| k.layer >= from_layer).copied().collect();
        let mut n = 0;
        for key in doomed {
            if let Some(e) = self.index.remove(&key) {
                let _ = std::fs::remove_file(&e.path);
                self.metrics.bytes_stored = self.metrics.bytes_stored.saturating_sub(e.size);
                self.metrics.invalidations += 1;
                n += 1;
            }
            self.lru.retain(|k| *k != key);
        }
        let _ = self.flush_index();
        n
    }

    /// Evict LRU entries until under budget.
    fn evict_lru(&mut self) -> usize {
        let mut evicted = 0;
        while self.metrics.bytes_stored > self.max_bytes && !self.lru.is_empty() {
            // Evict the least-recently-used entry.
            let key = self.lru.remove(0);
            if let Some(e) = self.index.remove(&key) {
                let _ = std::fs::remove_file(&e.path);
                self.metrics.bytes_stored = self.metrics.bytes_stored.saturating_sub(e.size);
                self.metrics.evictions += 1;
                evicted += 1;
            }
        }
        evicted
    }

    /// Persist the index.
    pub fn flush_index(&mut self) -> Result<()> {
        let bytes = bincode::serialize(&self.index)
            .map_err(|e| KatraError::Protocol(format!("cache index: {e}")))?;
        let final_path = self.root.join("index.bin");
        let tmp_path = self.root.join("index.bin.tmp");
        std::fs::write(&tmp_path, &bytes).map_err(KatraError::from)?;
        std::fs::rename(&tmp_path, &final_path).map_err(KatraError::from)?;
        Ok(())
    }

    /// Metrics snapshot.
    pub fn metrics(&self) -> &CacheMetrics {
        &self.metrics
    }

    /// Current byte usage.
    pub fn bytes_stored(&self) -> u64 {
        self.metrics.bytes_stored
    }

    /// Entry count.
    pub fn entry_count(&self) -> usize {
        self.index.len()
    }

    /// Stable identity hash of the whole index (for audits).
    pub fn index_fingerprint(&self) -> u64 {
        let mut keys: Vec<u64> = self
            .index
            .keys()
            .map(|k| fnv1a_parts(&[k.layer as u64, k.content_hash, k.meta_hash]))
            .collect();
        keys.sort();
        fnv1a_parts(&keys)
    }
}

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

    fn temp_cache(name: &str) -> PathBuf {
        let p = std::env::temp_dir().join(format!("katra-cache-{name}-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&p);
        p
    }

    #[test]
    fn roundtrip_persists() {
        let root = temp_cache("roundtrip");
        let k = CacheKey::new(3, 0xabc, 0xdef);
        {
            let mut c = CacheStore::open(&root, 1 << 20).unwrap();
            c.insert(&k, b"spirv-bytes").unwrap();
            assert_eq!(c.lookup(&k).unwrap(), b"spirv-bytes");
        }
        // Reopen: still there.
        let mut c = CacheStore::open(&root, 1 << 20).unwrap();
        assert!(c.contains(&k));
        assert_eq!(c.lookup(&k).unwrap(), b"spirv-bytes");
        let _ = std::fs::remove_dir_all(&root);
    }

    #[test]
    fn invalidate_cascade_keeps_upper_layers() {
        let root = temp_cache("cascade");
        let mut c = CacheStore::open(&root, 1 << 20).unwrap();
        for layer in 1..=4u8 {
            c.insert(&CacheKey::new(layer, layer as u64, 0), b"x").unwrap();
        }
        let n = c.invalidate_from(4);
        assert_eq!(n, 1); // only layer 4 removed
        assert!(c.contains(&CacheKey::new(1, 1, 0)));
        assert!(c.contains(&CacheKey::new(2, 2, 0)));
        assert!(c.contains(&CacheKey::new(3, 3, 0)));
        assert!(!c.contains(&CacheKey::new(4, 4, 0)));
        let _ = std::fs::remove_dir_all(&root);
    }

    #[test]
    fn lru_eviction_respects_budget() {
        let root = temp_cache("lru");
        let mut c = CacheStore::open(&root, 32).unwrap(); // tiny budget
        for i in 0..10u64 {
            c.insert(&CacheKey::new(1, i, 0), b"12345678901234567890").unwrap(); // 20 bytes each
        }
        assert!(c.bytes_stored() <= 32);
        assert!(c.entry_count() <= 2);
        let _ = std::fs::remove_dir_all(&root);
    }

    #[test]
    fn corrupt_index_tolerated() {
        let root = temp_cache("corrupt");
        let index_path = root.join("index.bin");
        std::fs::create_dir_all(&root).unwrap();
        std::fs::write(&index_path, b"garbage").unwrap();
        let c = CacheStore::open(&root, 1 << 20).unwrap();
        assert_eq!(c.entry_count(), 0);
        let _ = std::fs::remove_dir_all(&root);
    }
}