Skip to main content

shape_vm/
blob_cache_v2.rs

1//! Blob-level cache for content-addressed function blobs.
2//! Caches individual FunctionBlobs by content hash on disk (~/.shape/cache/blobs/).
3
4use crate::bytecode::{FunctionBlob, FunctionHash};
5use std::collections::HashMap;
6use std::fs;
7use std::path::PathBuf;
8
9/// Cache for individual function blobs, keyed by content hash.
10pub struct BlobCache {
11    /// In-memory cache layer.
12    memory: HashMap<FunctionHash, FunctionBlob>,
13    /// Optional disk cache directory.
14    disk_root: Option<PathBuf>,
15    /// Cache statistics.
16    stats: CacheStats,
17}
18
19#[derive(Debug, Default, Clone)]
20pub struct CacheStats {
21    pub memory_hits: u64,
22    pub disk_hits: u64,
23    pub misses: u64,
24    pub insertions: u64,
25    pub evictions: u64,
26}
27
28impl BlobCache {
29    /// Create a memory-only cache.
30    pub fn memory_only() -> Self {
31        Self {
32            memory: HashMap::new(),
33            disk_root: None,
34            stats: CacheStats::default(),
35        }
36    }
37
38    /// Create a cache with disk persistence.
39    pub fn with_disk(root: PathBuf) -> std::io::Result<Self> {
40        fs::create_dir_all(&root)?;
41        Ok(Self {
42            memory: HashMap::new(),
43            disk_root: Some(root),
44            stats: CacheStats::default(),
45        })
46    }
47
48    /// Check if a blob is cached.
49    pub fn has_blob(&self, hash: &FunctionHash) -> bool {
50        if self.memory.contains_key(hash) {
51            return true;
52        }
53        if let Some(ref root) = self.disk_root {
54            let path = Self::blob_path(root, hash);
55            return path.exists();
56        }
57        false
58    }
59
60    /// Get a cached blob by hash.
61    pub fn get_blob(&mut self, hash: &FunctionHash) -> Option<FunctionBlob> {
62        // Check memory first
63        if let Some(blob) = self.memory.get(hash) {
64            self.stats.memory_hits += 1;
65            return Some(blob.clone());
66        }
67
68        // Check disk
69        if let Some(ref root) = self.disk_root {
70            let path = Self::blob_path(root, hash);
71            if let Ok(data) = fs::read(&path) {
72                if let Ok(blob) = rmp_serde::from_slice::<FunctionBlob>(&data) {
73                    self.stats.disk_hits += 1;
74                    self.memory.insert(*hash, blob.clone());
75                    return Some(blob);
76                }
77            }
78        }
79
80        self.stats.misses += 1;
81        None
82    }
83
84    /// Store a blob in the cache.
85    pub fn put_blob(&mut self, blob: &FunctionBlob) {
86        let hash = blob.content_hash;
87        self.memory.insert(hash, blob.clone());
88        self.stats.insertions += 1;
89
90        // Write to disk if enabled
91        if let Some(ref root) = self.disk_root {
92            let path = Self::blob_path(root, &hash);
93            if let Some(parent) = path.parent() {
94                let _ = fs::create_dir_all(parent);
95            }
96            if let Ok(data) = rmp_serde::to_vec(blob) {
97                let _ = fs::write(&path, data);
98            }
99        }
100    }
101
102    /// Get cache statistics.
103    pub fn stats(&self) -> &CacheStats {
104        &self.stats
105    }
106
107    /// Clear the in-memory cache.
108    pub fn clear_memory(&mut self) {
109        self.memory.clear();
110    }
111
112    /// Number of blobs in memory cache.
113    pub fn memory_size(&self) -> usize {
114        self.memory.len()
115    }
116
117    fn blob_path(root: &PathBuf, hash: &FunctionHash) -> PathBuf {
118        let hex = hex::encode(hash.0);
119        root.join(&hex[..2]).join(format!("{}.blob", &hex[2..]))
120    }
121}
122
123/// JIT code cache - in-memory only (Cranelift JITModule can't serialize).
124pub struct JitCodeCache {
125    entries: HashMap<FunctionHash, *const u8>,
126}
127
128// SAFETY: Function pointers from JIT are valid for the lifetime of the JITModule
129unsafe impl Send for JitCodeCache {}
130unsafe impl Sync for JitCodeCache {}
131
132impl JitCodeCache {
133    pub fn new() -> Self {
134        Self {
135            entries: HashMap::new(),
136        }
137    }
138
139    pub fn get(&self, hash: &FunctionHash) -> Option<*const u8> {
140        self.entries.get(hash).copied()
141    }
142
143    pub fn insert(&mut self, hash: FunctionHash, ptr: *const u8) {
144        self.entries.insert(hash, ptr);
145    }
146
147    pub fn contains(&self, hash: &FunctionHash) -> bool {
148        self.entries.contains_key(hash)
149    }
150
151    pub fn len(&self) -> usize {
152        self.entries.len()
153    }
154
155    pub fn is_empty(&self) -> bool {
156        self.entries.is_empty()
157    }
158}