vyre-wgpu 0.1.0

use crate::runtime::cache::lru::IntrusiveLru;
use rustc_hash::FxHashMap;

/// Maximum number of compiled compute pipelines retained by the runtime cache.
pub const MAX_PIPELINE_CACHE_ENTRIES: usize = 1024;

/// Build a deterministic cache key from shader source and entry point.
#[inline]
pub fn cache_key(wgsl_source: &str, entry_point: &str) -> String {
    let mut key = String::with_capacity(wgsl_source.len() + entry_point.len() + 1);
    key.push_str(entry_point);
    key.push('\0');
    key.push_str(wgsl_source);
    key
}

pub(crate) struct PipelineCache {
    pub(crate) map: FxHashMap<String, wgpu::ComputePipeline>,
    tokens_by_key: FxHashMap<String, u64>,
    keys_by_token: FxHashMap<u64, String>,
    lru: IntrusiveLru<u64, ()>,
    next_token: u64,
}

impl PipelineCache {
    #[inline]
    pub(crate) fn new() -> Self {
        Self {
            map: FxHashMap::default(),
            tokens_by_key: FxHashMap::default(),
            keys_by_token: FxHashMap::default(),
            lru: IntrusiveLru::with_capacity(MAX_PIPELINE_CACHE_ENTRIES),
            next_token: 1,
        }
    }

    #[inline]
    pub(crate) fn get(&mut self, key: &str) -> Option<wgpu::ComputePipeline> {
        let pipeline = self.map.get(key).cloned()?;
        if let Some(token) = self.tokens_by_key.get(key) {
            self.lru.touch(*token);
        }
        Some(pipeline)
    }

    #[inline]
    pub(crate) fn insert(&mut self, key: String, pipeline: wgpu::ComputePipeline) {
        if let Some(token) = self.tokens_by_key.get(&key) {
            self.lru.touch(*token);
            self.map.insert(key, pipeline);
            return;
        }
        self.evict_if_full();
        let token = self.allocate_token();
        self.tokens_by_key.insert(key.clone(), token);
        self.keys_by_token.insert(token, key.clone());
        self.lru.ensure(token);
        self.map.insert(key, pipeline);
    }

    fn evict_if_full(&mut self) {
        if self.map.len() < MAX_PIPELINE_CACHE_ENTRIES {
            return;
        }
        let Some(token) = self.lru.iter_coldest().next().map(|(token, _)| *token) else {
            return;
        };
        self.lru.remove(&token);
        if let Some(key) = self.keys_by_token.remove(&token) {
            self.tokens_by_key.remove(&key);
            self.map.remove(&key);
        }
    }

    fn allocate_token(&mut self) -> u64 {
        while self.keys_by_token.contains_key(&self.next_token) {
            self.next_token = self.next_token.wrapping_add(1).max(1);
        }
        let token = self.next_token;
        self.next_token = self.next_token.wrapping_add(1).max(1);
        token
    }
}