use std::path::PathBuf;
pub struct Checkpoint {
pub id: usize,
pub label: String,
files: Vec<(PathBuf, Option<String>)>,
}
#[derive(Default)]
pub struct CheckpointStore {
checkpoints: Vec<Checkpoint>,
next_id: usize,
}
impl CheckpointStore {
pub fn new() -> Self {
CheckpointStore::default()
}
pub fn snapshot(&mut self, label: impl Into<String>, paths: &[PathBuf]) -> usize {
let files = paths
.iter()
.map(|p| (p.clone(), std::fs::read_to_string(p).ok()))
.collect();
let id = self.next_id;
self.next_id += 1;
self.checkpoints.push(Checkpoint {
id,
label: label.into(),
files,
});
id
}
pub fn rewind(&mut self, id: usize) -> std::io::Result<usize> {
let Some(pos) = self.checkpoints.iter().position(|c| c.id == id) else {
return Ok(0);
};
let mut restored = 0;
for cp in self.checkpoints[pos..].iter().rev() {
for (path, content) in &cp.files {
match content {
Some(c) => std::fs::write(path, c)?,
None => {
let _ = std::fs::remove_file(path);
}
}
restored += 1;
}
}
self.checkpoints.truncate(pos);
Ok(restored)
}
pub fn rewind_last(&mut self, n: usize) -> std::io::Result<usize> {
if self.checkpoints.is_empty() || n == 0 {
return Ok(0);
}
let n = n.min(self.checkpoints.len());
let id = self.checkpoints[self.checkpoints.len() - n].id;
self.rewind(id)
}
pub fn len(&self) -> usize {
self.checkpoints.len()
}
pub fn is_empty(&self) -> bool {
self.checkpoints.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rewind_restores_prior_content() {
let dir = tempfile::tempdir().unwrap();
let file = dir.path().join("a.txt");
std::fs::write(&file, "v1").unwrap();
let mut store = CheckpointStore::new();
let id0 = store.snapshot("edit1", std::slice::from_ref(&file));
std::fs::write(&file, "v2").unwrap();
store.snapshot("edit2", std::slice::from_ref(&file));
std::fs::write(&file, "v3").unwrap();
store.rewind_last(1).unwrap();
assert_eq!(std::fs::read_to_string(&file).unwrap(), "v2");
store.rewind(id0).unwrap();
assert_eq!(std::fs::read_to_string(&file).unwrap(), "v1");
assert!(store.is_empty());
}
#[test]
fn rewind_removes_files_that_did_not_exist() {
let dir = tempfile::tempdir().unwrap();
let file = dir.path().join("new.txt");
let mut store = CheckpointStore::new();
store.snapshot("create", std::slice::from_ref(&file)); std::fs::write(&file, "created").unwrap();
assert!(file.exists());
store.rewind_last(1).unwrap();
assert!(
!file.exists(),
"rewind should remove a file that did not exist at snapshot time"
);
}
}