use std::fs::{File, create_dir_all};
use std::io::BufWriter;
use std::path::{Path, PathBuf};
use crate::dag::arena::DagArena;
use crate::dag::packed::{BorrowedArenaView, PackedArenaImage};
use crate::zerocopy::MmapBuffer;
#[derive(Debug, Clone)]
pub struct DiskCache {
cache_dir: PathBuf,
}
impl DiskCache {
pub fn new<P: AsRef<Path>>(cache_dir: P) -> std::io::Result<Self> {
let path = cache_dir.as_ref().to_path_buf();
create_dir_all(&path)?;
Ok(Self { cache_dir: path })
}
fn key_path(&self, key: &str) -> PathBuf {
self.cache_dir.join(format!("{key}.bin"))
}
pub fn spill(&self, key: &str, arena: &DagArena) -> Result<(), String> {
let image = PackedArenaImage::from_arena(arena);
let filepath = self.key_path(key);
let file = File::create(&filepath)
.map_err(|e| format!("Failed to create cache file {}: {e:?}", filepath.display()))?;
let buf_writer = BufWriter::new(file);
image
.write_to(buf_writer)
.map_err(|e| format!("Failed to stream arena to file: {e:?}"))
}
pub fn restore(&self, key: &str) -> Result<DagArena, String> {
let filepath = self.key_path(key);
let mm = MmapBuffer::open(&filepath)
.map_err(|e| format!("Failed to open cache file {}: {e:?}", filepath.display()))?;
let arena_owned = mm.with_view(|bytes| -> Result<DagArena, String> {
let view = crate::zerocopy::decode_zerocopy_raw::<BorrowedArenaView<'_>>(bytes)
.map_err(|e| format!("Packed arena decode failed: {e:?}"))?;
Ok(view.to_owned_arena())
})?;
Ok(arena_owned)
}
pub fn load_borrowed<R>(
&self,
key: &str,
f: impl FnOnce(BorrowedArenaView<'_>) -> R,
) -> Result<R, String> {
let filepath = self.key_path(key);
let mm = MmapBuffer::open(&filepath)
.map_err(|e| format!("Failed to open cache file {}: {e:?}", filepath.display()))?;
let out = mm.with_view(|bytes| -> Result<R, String> {
let view = crate::zerocopy::decode_zerocopy_raw::<BorrowedArenaView<'_>>(bytes)
.map_err(|e| format!("Packed arena decode failed: {e:?}"))?;
Ok(f(view))
})?;
Ok(out)
}
pub fn delete(&self, key: &str) -> Result<(), String> {
let filepath = self.key_path(key);
if filepath.exists() {
std::fs::remove_file(&filepath).map_err(|e| {
format!("Failed to delete cache file {}: {e:?}", filepath.display())
})?;
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct LazyDiskArena {
path: PathBuf,
}
impl LazyDiskArena {
#[must_use]
pub fn open<P: AsRef<Path>>(path: P) -> Self {
Self {
path: path.as_ref().to_path_buf(),
}
}
pub fn with_view<R>(&self, f: impl FnOnce(BorrowedArenaView<'_>) -> R) -> Result<R, String> {
let mm = MmapBuffer::open(&self.path)
.map_err(|e| format!("LazyDiskArena: cannot open {}: {e:?}", self.path.display()))?;
let out = mm.with_view(|bytes| -> Result<R, String> {
let view = crate::zerocopy::decode_zerocopy_raw::<BorrowedArenaView<'_>>(bytes)
.map_err(|e| format!("LazyDiskArena: decode failed: {e:?}"))?;
Ok(f(view))
})?;
Ok(out)
}
pub fn load(&self) -> Result<DagArena, String> {
self.with_view(|view| view.to_owned_arena())
}
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dag::builder::DagBuilder;
fn tmp_subdir(name: &str) -> PathBuf {
std::env::temp_dir().join(format!("rssn_diskcache_{}_{}", std::process::id(), name))
}
#[test]
fn round_trip_arena_via_mmap() {
let dir = tmp_subdir("rt");
let cache = DiskCache::new(&dir).expect("mkdir");
let mut b = DagBuilder::new();
let x = b.variable("x");
let y = b.variable("y");
let sum = b.add(x, y);
let c = b.constant(3.5);
let _ = b.mul(sum, c);
let original_len = b.arena().len();
cache.spill("k", b.arena()).expect("spill");
let restored = cache.restore("k").expect("restore");
assert_eq!(restored.len(), original_len);
for i in 0..original_len {
let orig = b
.arena()
.get(crate::dag::node::DagNodeId::new(i as u32))
.expect("orig");
let new = restored
.get(crate::dag::node::DagNodeId::new(i as u32))
.expect("new");
assert_eq!(orig.kind, new.kind);
assert_eq!(orig.children.len(), new.children.len());
assert_eq!(orig.children.as_slice(), new.children.as_slice());
}
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn load_borrowed_returns_view_without_copying_nodes() {
let dir = tmp_subdir("borrowed");
let cache = DiskCache::new(&dir).expect("mkdir");
let mut b = DagBuilder::new();
let x = b.variable("x");
let y = b.variable("y");
let _ = b.add(x, y);
cache.spill("k", b.arena()).expect("spill");
let len = cache.load_borrowed("k", |view| view.len()).expect("load");
assert_eq!(len, b.arena().len());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn lazy_disk_arena_with_view_and_load() {
let dir = tmp_subdir("lazy");
let cache = DiskCache::new(&dir).expect("mkdir");
let mut b = DagBuilder::new();
let x = b.variable("x");
let y = b.variable("y");
let _ = b.add(x, y);
cache.spill("k", b.arena()).expect("spill");
let path = dir.join("k.bin");
let lazy = LazyDiskArena::open(&path);
let len = lazy.with_view(|v| v.len()).expect("with_view");
assert_eq!(len, b.arena().len());
let arena = lazy.load().expect("load");
assert_eq!(arena.len(), b.arena().len());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn delete_removes_file() {
let dir = tmp_subdir("del");
let cache = DiskCache::new(&dir).expect("mkdir");
let mut b = DagBuilder::new();
let _ = b.variable("x");
cache.spill("k", b.arena()).expect("spill");
assert!(dir.join("k.bin").exists());
cache.delete("k").expect("delete");
assert!(!dir.join("k.bin").exists());
let _ = std::fs::remove_dir_all(&dir);
}
}