dot/
snapshot.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3
4pub struct SnapshotManager {
5    originals: HashMap<PathBuf, Option<String>>,
6    checkpoints: Vec<HashMap<PathBuf, Option<String>>>,
7}
8
9impl Default for SnapshotManager {
10    fn default() -> Self {
11        Self::new()
12    }
13}
14
15impl SnapshotManager {
16    pub fn new() -> Self {
17        Self {
18            originals: HashMap::new(),
19            checkpoints: Vec::new(),
20        }
21    }
22
23    pub fn before_write(&mut self, path: &str) {
24        let path = PathBuf::from(path);
25        if self.originals.contains_key(&path) {
26            return;
27        }
28        let content = std::fs::read_to_string(&path).ok();
29        self.originals.insert(path, content);
30    }
31
32    pub fn list_changes(&self) -> Vec<(String, ChangeKind)> {
33        let mut changes = Vec::new();
34        for (path, original) in &self.originals {
35            let current = std::fs::read_to_string(path).ok();
36            let kind = match (original, &current) {
37                (None, Some(_)) => ChangeKind::Created,
38                (Some(_), None) => ChangeKind::Deleted,
39                (Some(old), Some(new)) if old != new => ChangeKind::Modified,
40                (None, None) => continue,
41                _ => continue,
42            };
43            changes.push((path.display().to_string(), kind));
44        }
45        changes.sort_by(|a, b| a.0.cmp(&b.0));
46        changes
47    }
48
49    pub fn restore(&self, path: &str) -> anyhow::Result<String> {
50        let key = PathBuf::from(path);
51        let original = self
52            .originals
53            .get(&key)
54            .ok_or_else(|| anyhow::anyhow!("no snapshot for {}", path))?;
55        match original {
56            None => {
57                if key.exists() {
58                    std::fs::remove_file(&key)?;
59                    Ok(format!("deleted {} (was created this session)", path))
60                } else {
61                    Ok(format!("{} already gone", path))
62                }
63            }
64            Some(content) => {
65                std::fs::write(&key, content)?;
66                Ok(format!("restored {}", path))
67            }
68        }
69    }
70
71    pub fn restore_all(&self) -> anyhow::Result<Vec<String>> {
72        let mut restored = Vec::new();
73        for (path, original) in &self.originals {
74            match original {
75                None => {
76                    if path.exists() {
77                        std::fs::remove_file(path)?;
78                        restored.push(format!("deleted {}", path.display()));
79                    }
80                }
81                Some(content) => {
82                    std::fs::write(path, content)?;
83                    restored.push(format!("restored {}", path.display()));
84                }
85            }
86        }
87        Ok(restored)
88    }
89
90    pub fn file_count(&self) -> usize {
91        self.originals.len()
92    }
93
94    pub fn clear(&mut self) {
95        self.originals.clear();
96        self.checkpoints.clear();
97    }
98
99    pub fn checkpoint(&mut self) {
100        let mut snap = HashMap::new();
101        for path in self.originals.keys() {
102            let content = std::fs::read_to_string(path).ok();
103            snap.insert(path.clone(), content);
104        }
105        self.checkpoints.push(snap);
106    }
107
108    pub fn checkpoint_count(&self) -> usize {
109        self.checkpoints.len()
110    }
111
112    pub fn restore_to_checkpoint(&self, idx: usize) -> anyhow::Result<Vec<String>> {
113        let snap = self
114            .checkpoints
115            .get(idx)
116            .ok_or_else(|| anyhow::anyhow!("no checkpoint at index {}", idx))?;
117        let mut restored = Vec::new();
118        for (path, content) in snap {
119            match content {
120                None => {
121                    if path.exists() {
122                        std::fs::remove_file(path)?;
123                        restored.push(format!("deleted {}", path.display()));
124                    }
125                }
126                Some(c) => {
127                    std::fs::write(path, c)?;
128                    restored.push(format!("restored {}", path.display()));
129                }
130            }
131        }
132        for (path, original) in &self.originals {
133            if !snap.contains_key(path) {
134                match original {
135                    None => {
136                        if path.exists() {
137                            std::fs::remove_file(path)?;
138                            restored.push(format!("deleted {}", path.display()));
139                        }
140                    }
141                    Some(c) => {
142                        std::fs::write(path, c)?;
143                        restored.push(format!("restored {}", path.display()));
144                    }
145                }
146            }
147        }
148        Ok(restored)
149    }
150
151    pub fn truncate_checkpoints(&mut self, count: usize) {
152        self.checkpoints.truncate(count);
153    }
154}
155
156pub enum ChangeKind {
157    Created,
158    Modified,
159    Deleted,
160}
161
162impl ChangeKind {
163    pub fn label(&self) -> &'static str {
164        match self {
165            Self::Created => "created",
166            Self::Modified => "modified",
167            Self::Deleted => "deleted",
168        }
169    }
170
171    pub fn icon(&self) -> &'static str {
172        match self {
173            Self::Created => "+",
174            Self::Modified => "~",
175            Self::Deleted => "-",
176        }
177    }
178}
dot/snapshot.rs

dot/
snapshot.rs
