1use std::{fs, path::PathBuf};
5
6use objects::{
7 fs_atomic::{sync_directory, temp_path, write_file_atomic},
8 fs_ops::remove_path_recursively,
9 lock::RepoLock,
10 object::ContentHash,
11};
12use serde::{Deserialize, Serialize};
13
14use crate::{Repository, Result};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
17pub struct StashId([u8; 16]);
18
19impl StashId {
20 fn generate() -> Self {
21 use std::sync::atomic::{AtomicU64, Ordering};
22
23 static NEXT: AtomicU64 = AtomicU64::new(1);
24 let mut hasher = blake3::Hasher::new();
25 hasher.update(
26 &chrono::Utc::now()
27 .timestamp_nanos_opt()
28 .unwrap_or_default()
29 .to_le_bytes(),
30 );
31 hasher.update(&NEXT.fetch_add(1, Ordering::Relaxed).to_le_bytes());
32 let mut bytes = [0; 16];
33 bytes.copy_from_slice(&hasher.finalize().as_bytes()[..16]);
34 Self(bytes)
35 }
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct StashEntry {
40 pub index: usize,
41 pub stash_id: StashId,
42 pub tree_hash: String,
43 pub parent_tree_hash: String,
44 pub message: Option<String>,
45 pub created_at: chrono::DateTime<chrono::Utc>,
46}
47
48pub struct StashManager {
49 stash_dir: PathBuf,
50 lock: RepoLock,
51}
52
53impl StashManager {
54 pub fn new(heddle_dir: impl AsRef<std::path::Path>) -> Self {
55 Self {
56 stash_dir: heddle_dir.as_ref().join("stashes"),
57 lock: RepoLock::at(heddle_dir.as_ref().join("locks/stash.lock")),
58 }
59 }
60
61 pub fn init(&self) -> Result<()> {
62 if !self.stash_dir.exists() {
63 fs::create_dir_all(&self.stash_dir)?;
64 }
65 Ok(())
66 }
67
68 pub fn push(
69 &self,
70 tree_hash: ContentHash,
71 parent_tree_hash: String,
72 message: Option<String>,
73 ) -> Result<StashEntry> {
74 let _lock = self.write_lock()?;
75 let stashes = self.list_unlocked()?;
76
77 let stash_id = StashId::generate();
78
79 let entry = StashEntry {
80 index: stashes.len(),
81 stash_id,
82 tree_hash: tree_hash.to_string(),
83 parent_tree_hash,
84 message,
85 created_at: chrono::Utc::now(),
86 };
87
88 let entry_path = self.stash_dir.join(format!("{}", entry.index));
89 let content = serde_json::to_string(&entry)?;
90 write_file_atomic(&entry_path, content.as_bytes())?;
91
92 Ok(entry)
93 }
94
95 pub fn list(&self) -> Result<Vec<StashEntry>> {
96 let _lock = self.read_lock()?;
97 self.list_unlocked()
98 }
99
100 fn list_unlocked(&self) -> Result<Vec<StashEntry>> {
101 if !self.stash_dir.exists() {
102 return Ok(Vec::new());
103 }
104
105 let mut stashes = Vec::new();
106
107 for entry in fs::read_dir(&self.stash_dir)? {
108 let entry = entry?;
109 let path = entry.path();
110
111 if path.extension().is_none()
112 && let Ok(content) = fs::read_to_string(&path)
113 && let Ok(stash) = serde_json::from_str::<StashEntry>(&content)
114 {
115 stashes.push(stash);
116 }
117 }
118
119 stashes.sort_by_key(|s| s.index);
120 Ok(stashes)
121 }
122
123 pub fn top(&self) -> Result<Option<StashEntry>> {
124 let stashes = self.list()?;
125 Ok(stashes.last().cloned())
126 }
127
128 pub fn drop(&self) -> Result<Option<StashEntry>> {
129 let _lock = self.write_lock()?;
130 let mut stashes = self.list_unlocked()?;
131
132 if stashes.is_empty() {
133 return Ok(None);
134 }
135
136 let removed = stashes.pop();
137 self.rewrite_unlocked(&mut stashes)?;
138 Ok(removed)
139 }
140
141 pub fn pop_with<F>(&self, apply: F) -> Result<Option<StashEntry>>
142 where
143 F: FnOnce(&StashEntry) -> Result<()>,
144 {
145 let _lock = self.write_lock()?;
146 let mut stashes = self.list_unlocked()?;
147
148 let Some(removed) = stashes.pop() else {
149 return Ok(None);
150 };
151
152 apply(&removed)?;
153 self.rewrite_unlocked(&mut stashes)?;
154 Ok(Some(removed))
155 }
156
157 fn rewrite_unlocked(&self, stashes: &mut [StashEntry]) -> Result<()> {
158 let parent = self
159 .stash_dir
160 .parent()
161 .ok_or_else(|| std::io::Error::other("invalid stash directory"))?;
162 fs::create_dir_all(parent)?;
163
164 let replacement_dir = temp_path(&self.stash_dir);
165 fs::create_dir_all(&replacement_dir)?;
166
167 for (new_index, entry) in stashes.iter_mut().enumerate() {
168 entry.index = new_index;
169 let path = replacement_dir.join(format!("{}", new_index));
170 let content = serde_json::to_string(entry)?;
171 write_file_atomic(&path, content.as_bytes())?;
172 }
173
174 sync_directory(&replacement_dir)?;
175
176 let backup_dir = self.stash_dir.with_extension("old");
177 remove_stash_path(&backup_dir)?;
178 fs::rename(&self.stash_dir, &backup_dir)?;
179 sync_directory(parent)?;
180 if let Err(error) = fs::rename(&replacement_dir, &self.stash_dir) {
181 fs::rename(&backup_dir, &self.stash_dir)?;
182 sync_directory(parent)?;
183 return Err(error.into());
184 }
185 sync_directory(parent)?;
186 remove_stash_path(&backup_dir)?;
187 sync_directory(parent)?;
188
189 Ok(())
190 }
191
192 pub fn clear(&self) -> Result<usize> {
193 let _lock = self.write_lock()?;
194 let stashes = self.list_unlocked()?;
195 let count = stashes.len();
196
197 if self.stash_dir.exists() {
198 if self.stash_dir.is_symlink() {
199 fs::remove_file(&self.stash_dir)?;
200 } else {
201 remove_path_recursively(&self.stash_dir)?;
202 }
203 }
204 fs::create_dir_all(&self.stash_dir)?;
205
206 Ok(count)
207 }
208
209 fn read_lock(&self) -> Result<objects::lock::ReadLockGuard> {
210 self.lock.read().map_err(|err| {
211 std::io::Error::other(format!("failed to acquire stash lock: {err}")).into()
212 })
213 }
214
215 fn write_lock(&self) -> Result<objects::lock::WriteLockGuard> {
216 self.lock.write().map_err(|err| {
217 std::io::Error::other(format!("failed to acquire stash lock: {err}")).into()
218 })
219 }
220}
221
222fn remove_stash_path(path: &std::path::Path) -> Result<()> {
223 if !path.exists() {
224 return Ok(());
225 }
226
227 if path.is_symlink() {
228 fs::remove_file(path)?;
229 } else {
230 remove_path_recursively(path)?;
231 }
232
233 Ok(())
234}
235
236impl Repository {
237 pub fn stash_manager(&self) -> StashManager {
238 StashManager::new(self.heddle_dir())
239 }
240}
241
242#[cfg(test)]
243mod tests {
244 use std::{
245 sync::{Arc, Barrier},
246 thread,
247 };
248
249 use tempfile::TempDir;
250
251 use super::*;
252
253 fn create_manager() -> (TempDir, StashManager) {
254 let temp_dir = TempDir::new().unwrap();
255 let heddle_dir = temp_dir.path().join(".heddle");
256 let manager = StashManager::new(&heddle_dir);
257 manager.init().unwrap();
258 (temp_dir, manager)
259 }
260
261 #[test]
262 fn test_drop_rewrites_remaining_entries() {
263 let (_temp_dir, manager) = create_manager();
264 let first = manager
265 .push(ContentHash::compute(b"one"), "parent-1".to_string(), None)
266 .unwrap();
267 let second = manager
268 .push(ContentHash::compute(b"two"), "parent-2".to_string(), None)
269 .unwrap();
270 let third = manager
271 .push(ContentHash::compute(b"three"), "parent-3".to_string(), None)
272 .unwrap();
273
274 let removed = manager.drop().unwrap().unwrap();
275 assert_eq!(removed.stash_id, third.stash_id);
276
277 let remaining = manager.list().unwrap();
278 assert_eq!(remaining.len(), 2);
279 assert_eq!(remaining[0].index, 0);
280 assert_eq!(remaining[0].stash_id, first.stash_id);
281 assert_eq!(remaining[1].index, 1);
282 assert_eq!(remaining[1].stash_id, second.stash_id);
283
284 let temp_entries = fs::read_dir(&manager.stash_dir)
285 .unwrap()
286 .filter_map(|entry| entry.ok())
287 .filter(|entry| entry.file_name().to_string_lossy().contains(".tmp-"))
288 .count();
289 assert_eq!(temp_entries, 0);
290 assert!(!manager.stash_dir.with_extension("old").exists());
291 }
292
293 #[test]
294 fn test_pop_with_drops_only_after_successful_apply() {
295 let (_temp_dir, manager) = create_manager();
296 let first = manager
297 .push(ContentHash::compute(b"one"), "parent-1".to_string(), None)
298 .unwrap();
299 let second = manager
300 .push(ContentHash::compute(b"two"), "parent-2".to_string(), None)
301 .unwrap();
302
303 let error = manager
304 .pop_with(|_| Err(std::io::Error::other("apply failed").into()))
305 .unwrap_err();
306 assert!(error.to_string().contains("apply failed"));
307 assert_eq!(manager.list().unwrap().len(), 2);
308
309 let applied = manager
310 .pop_with(|stash| {
311 assert_eq!(stash.stash_id, second.stash_id);
312 Ok(())
313 })
314 .unwrap()
315 .unwrap();
316 assert_eq!(applied.stash_id, second.stash_id);
317
318 let remaining = manager.list().unwrap();
319 assert_eq!(remaining.len(), 1);
320 assert_eq!(remaining[0].index, 0);
321 assert_eq!(remaining[0].stash_id, first.stash_id);
322 }
323
324 #[test]
325 fn test_concurrent_pushes_preserve_all_entries() {
326 let (_temp_dir, manager) = create_manager();
327 let manager = Arc::new(manager);
328 let barrier = Arc::new(Barrier::new(9));
329 let mut handles = Vec::new();
330
331 for i in 0..8 {
332 let manager = Arc::clone(&manager);
333 let barrier = Arc::clone(&barrier);
334 handles.push(thread::spawn(move || {
335 barrier.wait();
336 manager
337 .push(
338 ContentHash::compute(format!("tree-{i}").as_bytes()),
339 format!("parent-{i}"),
340 Some(format!("stash-{i}")),
341 )
342 .unwrap();
343 }));
344 }
345
346 barrier.wait();
347
348 for handle in handles {
349 handle.join().unwrap();
350 }
351
352 let stashes = manager.list().unwrap();
353 assert_eq!(stashes.len(), 8);
354
355 let mut indices: Vec<_> = stashes.iter().map(|entry| entry.index).collect();
356 indices.sort_unstable();
357 assert_eq!(indices, (0..8).collect::<Vec<_>>());
358
359 let change_ids: std::collections::HashSet<_> =
360 stashes.iter().map(|entry| entry.stash_id).collect();
361 assert_eq!(change_ids.len(), 8);
362 }
363}