1use std::collections::HashMap;
42use std::fs::File;
43use std::io::Write as IoWrite;
44use std::path::{Path, PathBuf};
45use std::sync::OnceLock;
46use std::time::{SystemTime, UNIX_EPOCH};
47
48use memmap2::Mmap;
49use parking_lot::Mutex;
50use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
51use std::os::unix::fs::MetadataExt;
52
53pub const SHARD_MAGIC: u32 = 0x5A525343;
56pub const SHARD_FORMAT_VERSION: u32 = 1;
58#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
60#[archive(check_bytes)]
61pub struct ShardHeader {
62 pub magic: u32,
64 pub format_version: u32,
66 pub zshrs_version: String,
68 pub pointer_width: u32,
70 pub built_at_secs: u64,
72}
73#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
75#[archive(check_bytes)]
76pub struct ScriptEntry {
77 pub mtime_secs: i64,
79 pub mtime_nsecs: i64,
81 pub binary_mtime_at_cache: i64,
83 pub cached_at_secs: i64,
85 pub chunk_blob: Vec<u8>,
87}
88#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
90#[archive(check_bytes)]
91pub struct ScriptShard {
92 pub header: ShardHeader,
94 pub entries: HashMap<String, ScriptEntry>,
96}
97
98pub struct MmappedShard {
101 _mmap: Mmap,
103 archived: *const ArchivedScriptShard,
105}
106
107unsafe impl Send for MmappedShard {}
110unsafe impl Sync for MmappedShard {}
111
112impl MmappedShard {
113 pub fn open(path: &Path) -> Option<Self> {
115 let file = File::open(path).ok()?;
116 let mmap = unsafe { Mmap::map(&file).ok()? };
117 let archived = rkyv::check_archived_root::<ScriptShard>(&mmap[..]).ok()?;
118 let archived_ptr = archived as *const ArchivedScriptShard;
119 Some(Self {
120 _mmap: mmap,
121 archived: archived_ptr,
122 })
123 }
124
125 fn shard(&self) -> &ArchivedScriptShard {
126 unsafe { &*self.archived }
128 }
129
130 fn header_ok(&self) -> bool {
131 let h = &self.shard().header;
132 let magic: u32 = h.magic.into();
133 let fv: u32 = h.format_version.into();
134 let pw: u32 = h.pointer_width.into();
135 magic == SHARD_MAGIC
136 && fv == SHARD_FORMAT_VERSION
137 && pw as usize == std::mem::size_of::<usize>()
138 && h.zshrs_version.as_str() == env!("CARGO_PKG_VERSION")
139 }
140
141 fn lookup(&self, path: &str) -> Option<&ArchivedScriptEntry> {
142 self.shard().entries.get(path)
143 }
144
145 fn entry_count(&self) -> usize {
146 self.shard().entries.len()
147 }
148}
149
150pub struct ScriptCache {
152 path: PathBuf,
154 lock_path: PathBuf,
156 mmap: Mutex<Option<MmappedShard>>,
158}
159
160impl ScriptCache {
161 pub fn open(path: &Path) -> std::io::Result<Self> {
163 if let Some(parent) = path.parent() {
164 std::fs::create_dir_all(parent)?;
165 }
166 let parent = path.parent().unwrap_or_else(|| Path::new("/tmp"));
167 let lock_path = parent.join(format!(
168 "{}.lock",
169 path.file_name()
170 .and_then(|s| s.to_str())
171 .unwrap_or("scripts.rkyv")
172 ));
173 Ok(Self {
174 path: path.to_path_buf(),
175 lock_path,
176 mmap: Mutex::new(None),
177 })
178 }
179
180 fn ensure_mmap(&self) {
181 let mut guard = self.mmap.lock();
182 if guard.is_none() {
183 *guard = MmappedShard::open(&self.path);
184 }
185 }
186
187 fn invalidate_mmap(&self) {
188 let mut guard = self.mmap.lock();
189 *guard = None;
190 }
191
192 pub fn get(&self, path: &str, mtime_secs: i64, mtime_nsecs: i64) -> Option<Vec<u8>> {
195 self.ensure_mmap();
196 let guard = self.mmap.lock();
197 let shard = guard.as_ref()?;
198 if !shard.header_ok() {
199 return None;
200 }
201 let entry = shard.lookup(path)?;
202
203 let entry_mtime_s: i64 = entry.mtime_secs.into();
204 let entry_mtime_ns: i64 = entry.mtime_nsecs.into();
205 if entry_mtime_s != mtime_secs || entry_mtime_ns != mtime_nsecs {
206 return None;
207 }
208
209 if let Some(bin_mtime) = current_binary_mtime_secs() {
210 let cached_bin_mtime: i64 = entry.binary_mtime_at_cache.into();
211 if cached_bin_mtime < bin_mtime {
212 return None;
213 }
214 }
215
216 Some(entry.chunk_blob.as_slice().to_vec())
217 }
218
219 pub fn put(
221 &self,
222 path: &str,
223 mtime_secs: i64,
224 mtime_nsecs: i64,
225 chunk_blob: Vec<u8>,
226 ) -> Result<(), String> {
227 let _lock = match acquire_lock(&self.lock_path) {
228 Some(l) => l,
229 None => return Ok(()),
230 };
231
232 let mut shard = match read_owned_shard(&self.path) {
233 Some(s)
234 if s.header.zshrs_version == env!("CARGO_PKG_VERSION")
235 && s.header.pointer_width as usize == std::mem::size_of::<usize>()
236 && s.header.format_version == SHARD_FORMAT_VERSION =>
237 {
238 s
239 }
240 _ => fresh_shard(),
241 };
242
243 let bin_mtime = current_binary_mtime_secs().unwrap_or(0);
244 let entry = ScriptEntry {
245 mtime_secs,
246 mtime_nsecs,
247 binary_mtime_at_cache: bin_mtime,
248 cached_at_secs: now_secs(),
249 chunk_blob,
250 };
251 shard.entries.insert(path.to_string(), entry);
252 shard.header.built_at_secs = now_secs() as u64;
253
254 write_shard_atomic(&self.path, &shard)?;
255 self.invalidate_mmap();
256 Ok(())
257 }
258
259 pub fn stats(&self) -> (i64, i64) {
261 self.ensure_mmap();
262 let guard = self.mmap.lock();
263 let Some(shard) = guard.as_ref() else {
264 return (0, 0);
265 };
266 let count = shard.entry_count() as i64;
267 let bytes: i64 = shard
268 .shard()
269 .entries
270 .values()
271 .map(|e| e.chunk_blob.len() as i64)
272 .sum();
273 (count, bytes)
274 }
275
276 pub fn list_scripts(&self) -> Vec<(String, f64, String, String)> {
279 self.ensure_mmap();
280 let guard = self.mmap.lock();
281 let Some(shard) = guard.as_ref() else {
282 return Vec::new();
283 };
284 let v = shard.shard().header.zshrs_version.as_str().to_string();
285 let mut out: Vec<(String, f64, String, String, i64)> = shard
286 .shard()
287 .entries
288 .iter()
289 .map(|(k, e)| {
290 let chunk_kb = e.chunk_blob.len() as f64 / 1024.0;
291 let cached_at: i64 = e.cached_at_secs.into();
292 let ts = format_local_ts(cached_at);
293 (k.as_str().to_string(), chunk_kb, v.clone(), ts, cached_at)
294 })
295 .collect();
296 out.sort_by_key(|x| std::cmp::Reverse(x.4));
297 out.into_iter()
298 .map(|(p, ck, ver, ts, _)| (p, ck, ver, ts))
299 .collect()
300 }
301
302 pub fn evict_stale(&self) -> usize {
304 let _lock = match acquire_lock(&self.lock_path) {
305 Some(l) => l,
306 None => return 0,
307 };
308 let mut shard = match read_owned_shard(&self.path) {
309 Some(s) => s,
310 None => return 0,
311 };
312 let before = shard.entries.len();
313 shard.entries.retain(|p, e| match file_mtime(Path::new(p)) {
314 Some((s, ns)) => s == e.mtime_secs && ns == e.mtime_nsecs,
315 None => false,
316 });
317 let evicted = before - shard.entries.len();
318 if evicted > 0 {
319 let _ = write_shard_atomic(&self.path, &shard);
320 self.invalidate_mmap();
321 }
322 evicted
323 }
324 pub fn clear(&self) -> std::io::Result<()> {
326 let _lock = acquire_lock(&self.lock_path);
327 let res = match std::fs::remove_file(&self.path) {
328 Ok(()) => Ok(()),
329 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
330 Err(e) => Err(e),
331 };
332 self.invalidate_mmap();
333 res
334 }
335}
336
337fn acquire_lock(path: &Path) -> Option<nix::fcntl::Flock<File>> {
338 let f = File::options()
339 .read(true)
340 .write(true)
341 .create(true)
342 .truncate(false)
343 .open(path)
344 .ok()?;
345 nix::fcntl::Flock::lock(f, nix::fcntl::FlockArg::LockExclusive).ok()
346}
347
348fn fresh_shard() -> ScriptShard {
349 ScriptShard {
350 header: ShardHeader {
351 magic: SHARD_MAGIC,
352 format_version: SHARD_FORMAT_VERSION,
353 zshrs_version: env!("CARGO_PKG_VERSION").to_string(),
354 pointer_width: std::mem::size_of::<usize>() as u32,
355 built_at_secs: now_secs() as u64,
356 },
357 entries: HashMap::new(),
358 }
359}
360
361fn read_owned_shard(path: &Path) -> Option<ScriptShard> {
362 let bytes = std::fs::read(path).ok()?;
363 let archived = rkyv::check_archived_root::<ScriptShard>(&bytes[..]).ok()?;
364 archived.deserialize(&mut rkyv::Infallible).ok()
365}
366
367fn write_shard_atomic(path: &Path, shard: &ScriptShard) -> Result<(), String> {
368 let bytes = rkyv::to_bytes::<_, 4096>(shard).map_err(|e| format!("rkyv serialize: {}", e))?;
369
370 let parent = path.parent().expect("cache path has parent");
371 let _ = std::fs::create_dir_all(parent);
372
373 let pid = std::process::id();
374 let nanos = SystemTime::now()
375 .duration_since(UNIX_EPOCH)
376 .map(|d| d.as_nanos())
377 .unwrap_or(0);
378 let tmp_path = parent.join(format!(
379 "{}.tmp.{}.{}",
380 path.file_name()
381 .and_then(|s| s.to_str())
382 .unwrap_or("scripts.rkyv"),
383 pid,
384 nanos
385 ));
386
387 {
388 let mut f = File::create(&tmp_path).map_err(|e| e.to_string())?;
389 f.write_all(&bytes).map_err(|e| e.to_string())?;
390 f.sync_all().map_err(|e| e.to_string())?;
391 }
392
393 std::fs::rename(&tmp_path, path).map_err(|e| e.to_string())?;
394 Ok(())
395}
396
397fn now_secs() -> i64 {
398 SystemTime::now()
399 .duration_since(UNIX_EPOCH)
400 .map(|d| d.as_secs() as i64)
401 .unwrap_or(0)
402}
403
404fn format_local_ts(secs: i64) -> String {
405 let dt = chrono::DateTime::<chrono::Local>::from(
406 UNIX_EPOCH + std::time::Duration::from_secs(secs.max(0) as u64),
407 );
408 dt.format("%Y-%m-%d %H:%M:%S").to_string()
409}
410pub fn file_mtime(path: &Path) -> Option<(i64, i64)> {
412 let meta = std::fs::metadata(path).ok()?;
413 Some((meta.mtime(), meta.mtime_nsec()))
414}
415
416fn current_binary_mtime_secs() -> Option<i64> {
417 static BIN_MTIME: OnceLock<Option<i64>> = OnceLock::new();
418 *BIN_MTIME.get_or_init(|| {
419 let exe = std::env::current_exe().ok()?;
420 let (secs, _) = file_mtime(&exe)?;
421 Some(secs)
422 })
423}
424
425pub fn default_cache_path() -> PathBuf {
427 dirs::home_dir()
428 .unwrap_or_else(|| PathBuf::from("/tmp"))
429 .join(".zshrs/scripts.rkyv")
430}
431
432pub static CACHE_DISABLED: std::sync::atomic::AtomicBool =
437 std::sync::atomic::AtomicBool::new(false);
438
439pub fn cache_enabled() -> bool {
442 if CACHE_DISABLED.load(std::sync::atomic::Ordering::Relaxed) {
443 return false;
444 }
445 !matches!(
446 std::env::var("ZSHRS_CACHE").as_deref(),
447 Ok("0") | Ok("false") | Ok("no")
448 )
449}
450
451pub static CACHE: once_cell::sync::Lazy<Option<ScriptCache>> = once_cell::sync::Lazy::new(|| {
454 if !cache_enabled() {
455 return None;
456 }
457 ScriptCache::open(&default_cache_path()).ok()
458});
459
460pub fn try_load_bytes(path: &Path) -> Option<Vec<u8>> {
462 let cache = CACHE.as_ref()?;
463 let canonical = path.canonicalize().ok()?;
464 let path_str = canonical.to_string_lossy();
465 let (mtime_s, mtime_ns) = file_mtime(&canonical)?;
466 cache.get(&path_str, mtime_s, mtime_ns)
467}
468
469pub fn try_save_bytes(path: &Path, chunk_blob: &[u8]) -> Result<(), String> {
473 let Some(cache) = CACHE.as_ref() else {
474 return Ok(());
475 };
476 let canonical = match path.canonicalize() {
477 Ok(p) => p,
478 Err(_) => return Ok(()),
479 };
480 let path_str = canonical.to_string_lossy();
481 let (mtime_s, mtime_ns) = match file_mtime(&canonical) {
482 Some(m) => m,
483 None => return Ok(()),
484 };
485 cache.put(&path_str, mtime_s, mtime_ns, chunk_blob.to_vec())
486}
487pub fn stats() -> Option<(i64, i64)> {
489 CACHE.as_ref().map(|c| c.stats())
490}
491pub fn evict_stale() -> usize {
493 CACHE.as_ref().map(|c| c.evict_stale()).unwrap_or(0)
494}
495pub fn clear() -> bool {
497 CACHE.as_ref().map(|c| c.clear().is_ok()).unwrap_or(false)
498}
499
500#[cfg(test)]
501mod tests {
502 use super::*;
503 use tempfile::tempdir;
504
505 #[test]
506 fn round_trip() {
507 let _g = crate::test_util::global_state_lock();
508 let dir = tempdir().unwrap();
509 let cache_path = dir.path().join("scripts.rkyv");
510 let cache = ScriptCache::open(&cache_path).unwrap();
511
512 let script_path = dir.path().join("test.zsh");
513 std::fs::write(&script_path, "echo hi").unwrap();
514
515 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
516 let path_str = script_path.to_string_lossy().to_string();
517
518 let blob = vec![1u8, 2, 3, 4, 5];
519 cache
520 .put(&path_str, mtime_s, mtime_ns, blob.clone())
521 .unwrap();
522
523 let loaded = cache.get(&path_str, mtime_s, mtime_ns).unwrap();
524 assert_eq!(loaded, blob);
525
526 let (count, _bytes) = cache.stats();
527 assert_eq!(count, 1);
528 }
529
530 #[test]
531 fn mtime_invalidation() {
532 let _g = crate::test_util::global_state_lock();
533 let dir = tempdir().unwrap();
534 let cache_path = dir.path().join("scripts.rkyv");
535 let cache = ScriptCache::open(&cache_path).unwrap();
536
537 let script_path = dir.path().join("test.zsh");
538 std::fs::write(&script_path, "echo hi").unwrap();
539
540 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
541 let path_str = script_path.to_string_lossy().to_string();
542 cache.put(&path_str, mtime_s, mtime_ns, vec![9u8]).unwrap();
543
544 assert!(cache.get(&path_str, mtime_s + 1, mtime_ns).is_none());
545 }
546
547 #[test]
548 fn second_put_replaces_first() {
549 let _g = crate::test_util::global_state_lock();
550 let dir = tempdir().unwrap();
551 let cache_path = dir.path().join("scripts.rkyv");
552 let cache = ScriptCache::open(&cache_path).unwrap();
553
554 let p1 = dir.path().join("a.zsh");
555 let p2 = dir.path().join("b.zsh");
556 std::fs::write(&p1, "1").unwrap();
557 std::fs::write(&p2, "2").unwrap();
558
559 let (m1s, m1n) = file_mtime(&p1).unwrap();
560 let (m2s, m2n) = file_mtime(&p2).unwrap();
561
562 cache
563 .put(&p1.to_string_lossy(), m1s, m1n, vec![1u8])
564 .unwrap();
565 cache
566 .put(&p2.to_string_lossy(), m2s, m2n, vec![2u8])
567 .unwrap();
568
569 let (count, _) = cache.stats();
570 assert_eq!(count, 2);
571 assert!(cache.get(&p1.to_string_lossy(), m1s, m1n).is_some());
572 assert!(cache.get(&p2.to_string_lossy(), m2s, m2n).is_some());
573 }
574
575 #[test]
576 fn corrupt_file_returns_no_mmap() {
577 let _g = crate::test_util::global_state_lock();
578 let dir = tempdir().unwrap();
579 let cache_path = dir.path().join("scripts.rkyv");
580 std::fs::write(&cache_path, b"this is not a valid rkyv archive").unwrap();
581 let cache = ScriptCache::open(&cache_path).unwrap();
582 assert!(cache.get("/nope", 0, 0).is_none());
583 }
584
585 #[test]
586 fn clear_removes_file() {
587 let _g = crate::test_util::global_state_lock();
588 let dir = tempdir().unwrap();
589 let cache_path = dir.path().join("scripts.rkyv");
590 let cache = ScriptCache::open(&cache_path).unwrap();
591
592 let script_path = dir.path().join("test.zsh");
593 std::fs::write(&script_path, "echo hi").unwrap();
594 let (mtime_s, mtime_ns) = file_mtime(&script_path).unwrap();
595 cache
596 .put(&script_path.to_string_lossy(), mtime_s, mtime_ns, vec![7u8])
597 .unwrap();
598 assert!(cache_path.exists());
599
600 cache.clear().unwrap();
601 assert!(!cache_path.exists());
602 }
603
604 #[test]
609 fn now_secs_is_positive_and_within_realistic_range() {
610 let _g = crate::test_util::global_state_lock();
611 let n = now_secs();
612 assert!(
614 (1_577_836_800..4_102_444_800).contains(&n),
615 "now_secs out of plausible range: {}",
616 n
617 );
618 }
619
620 #[test]
621 fn now_secs_does_not_go_backwards_in_quick_succession() {
622 let _g = crate::test_util::global_state_lock();
623 let a = now_secs();
624 let b = now_secs();
625 assert!(b >= a, "now_secs went backwards: {} -> {}", a, b);
626 }
627
628 #[test]
633 fn format_local_ts_includes_year_and_punctuation() {
634 let _g = crate::test_util::global_state_lock();
635 let s = format_local_ts(1_704_067_200);
638 assert!(s.starts_with("202"), "expected 21st century year: {}", s);
639 assert!(s.contains('-'), "expected dash separator: {}", s);
640 assert!(s.contains(':'), "expected colon separator: {}", s);
641 }
642
643 #[test]
644 fn format_local_ts_length_matches_pattern() {
645 let _g = crate::test_util::global_state_lock();
646 let s = format_local_ts(1_700_000_000);
647 assert_eq!(s.len(), 19, "unexpected width: {}", s);
649 }
650
651 #[test]
652 fn format_local_ts_handles_zero_secs_via_clamp() {
653 let _g = crate::test_util::global_state_lock();
654 let s = format_local_ts(0);
656 assert_eq!(s.len(), 19);
657 assert!(s.starts_with("19"), "expected 1970-ish year: {}", s);
658 }
659
660 #[test]
661 fn format_local_ts_negative_clamped_to_zero() {
662 let _g = crate::test_util::global_state_lock();
663 let s = format_local_ts(-1_000_000);
665 assert_eq!(s.len(), 19);
666 }
667
668 #[test]
673 fn file_mtime_returns_some_for_real_file() {
674 let _g = crate::test_util::global_state_lock();
675 let dir = tempdir().unwrap();
676 let p = dir.path().join("foo.zsh");
677 std::fs::write(&p, b"x").unwrap();
678 let (s, _ns) = file_mtime(&p).unwrap();
679 assert!(s > 0);
680 }
681
682 #[test]
683 fn file_mtime_returns_none_for_missing_path() {
684 let _g = crate::test_util::global_state_lock();
685 assert!(file_mtime(Path::new("/nonexistent/zshrs/script_cache_missing.bin")).is_none());
686 }
687
688 #[test]
693 fn default_cache_path_ends_in_scripts_rkyv() {
694 let _g = crate::test_util::global_state_lock();
695 let p = default_cache_path();
696 assert_eq!(p.file_name().and_then(|s| s.to_str()), Some("scripts.rkyv"));
697 }
698
699 #[test]
700 fn cache_enabled_true_when_env_unset() {
701 let _g = crate::test_util::global_state_lock();
702 let prev = std::env::var_os("ZSHRS_CACHE");
703 std::env::remove_var("ZSHRS_CACHE");
704 let on = cache_enabled();
705 if let Some(v) = prev {
706 std::env::set_var("ZSHRS_CACHE", v);
707 }
708 assert!(on, "cache should be enabled when ZSHRS_CACHE is unset");
709 }
710
711 #[test]
712 fn cache_enabled_false_when_env_is_zero_false_or_no() {
713 let _g = crate::test_util::global_state_lock();
714 let prev = std::env::var_os("ZSHRS_CACHE");
715 for v in ["0", "false", "no"] {
716 std::env::set_var("ZSHRS_CACHE", v);
717 assert!(!cache_enabled(), "ZSHRS_CACHE={} must disable cache", v);
718 }
719 if let Some(v) = prev {
720 std::env::set_var("ZSHRS_CACHE", v);
721 } else {
722 std::env::remove_var("ZSHRS_CACHE");
723 }
724 }
725
726 #[test]
727 fn cache_enabled_true_for_other_env_values() {
728 let _g = crate::test_util::global_state_lock();
731 let prev = std::env::var_os("ZSHRS_CACHE");
732 for v in ["1", "true", "yes", "on", ""] {
733 std::env::set_var("ZSHRS_CACHE", v);
734 assert!(
735 cache_enabled(),
736 "ZSHRS_CACHE={:?} must NOT disable cache",
737 v
738 );
739 }
740 if let Some(v) = prev {
741 std::env::set_var("ZSHRS_CACHE", v);
742 } else {
743 std::env::remove_var("ZSHRS_CACHE");
744 }
745 }
746}