1use crate::error::UndoRedoError;
4use chrono::{DateTime, Utc};
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use uuid::Uuid;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct Checkpoint {
12 pub id: String,
14 pub name: String,
16 pub description: String,
18 pub created_at: DateTime<Utc>,
20 pub changes_count: usize,
22 pub file_states: HashMap<String, String>,
24}
25
26impl Checkpoint {
27 pub fn new(
29 name: impl Into<String>,
30 description: impl Into<String>,
31 file_states: HashMap<String, String>,
32 ) -> Result<Self, UndoRedoError> {
33 let name = name.into();
34 let description = description.into();
35
36 if name.is_empty() {
38 return Err(UndoRedoError::validation_error("Checkpoint name cannot be empty"));
39 }
40
41 if file_states.is_empty() {
43 return Err(UndoRedoError::validation_error(
44 "Checkpoint must contain at least one file state",
45 ));
46 }
47
48 Ok(Checkpoint {
49 id: Uuid::new_v4().to_string(),
50 name,
51 description,
52 created_at: Utc::now(),
53 changes_count: file_states.len(),
54 file_states,
55 })
56 }
57
58 pub fn validate(&self) -> Result<(), UndoRedoError> {
60 if self.name.is_empty() {
61 return Err(UndoRedoError::validation_error("Checkpoint name cannot be empty"));
62 }
63
64 if self.file_states.is_empty() {
65 return Err(UndoRedoError::validation_error(
66 "Checkpoint must contain at least one file state",
67 ));
68 }
69
70 if self.changes_count != self.file_states.len() {
71 return Err(UndoRedoError::validation_error(
72 "Checkpoint changes_count does not match file_states length",
73 ));
74 }
75
76 Ok(())
77 }
78}
79
80pub struct CheckpointManager {
82 checkpoints: HashMap<String, Checkpoint>,
83 current_state: HashMap<String, String>,
84}
85
86impl CheckpointManager {
87 pub fn new() -> Self {
89 CheckpointManager {
90 checkpoints: HashMap::new(),
91 current_state: HashMap::new(),
92 }
93 }
94
95 pub fn create_checkpoint(
97 &mut self,
98 name: impl Into<String>,
99 description: impl Into<String>,
100 file_states: HashMap<String, String>,
101 ) -> Result<String, UndoRedoError> {
102 let checkpoint = Checkpoint::new(name, description, file_states)?;
103 let id = checkpoint.id.clone();
104 self.checkpoints.insert(id.clone(), checkpoint);
105 Ok(id)
106 }
107
108 pub fn list_checkpoints(&self) -> Vec<Checkpoint> {
110 self.checkpoints.values().cloned().collect()
111 }
112
113 pub fn get_checkpoint(&self, checkpoint_id: &str) -> Result<Checkpoint, UndoRedoError> {
115 self.checkpoints
116 .get(checkpoint_id)
117 .cloned()
118 .ok_or_else(|| UndoRedoError::checkpoint_not_found(checkpoint_id))
119 }
120
121 pub fn delete_checkpoint(&mut self, checkpoint_id: &str) -> Result<(), UndoRedoError> {
123 self.checkpoints
124 .remove(checkpoint_id)
125 .ok_or_else(|| UndoRedoError::checkpoint_not_found(checkpoint_id))?;
126 Ok(())
127 }
128
129 pub fn rollback_to(&mut self, checkpoint_id: &str) -> Result<(), UndoRedoError> {
135 let checkpoint = self.get_checkpoint(checkpoint_id)?;
137 checkpoint.validate()?;
138
139 let pre_rollback_state = self.current_state.clone();
141
142 let mut updates = Vec::new();
145 for (file_path, content) in &checkpoint.file_states {
146 if file_path.is_empty() {
148 self.current_state = pre_rollback_state;
150 return Err(UndoRedoError::validation_error("File path cannot be empty"));
151 }
152 updates.push((file_path.clone(), content.clone()));
153 }
154
155 for (file_path, content) in updates {
157 self.current_state.insert(file_path, content);
158 }
159
160 Ok(())
163 }
164
165 pub fn verify_rollback(&self, checkpoint_id: &str) -> Result<bool, UndoRedoError> {
167 let checkpoint = self.get_checkpoint(checkpoint_id)?;
168
169 for (file_path, expected_content) in &checkpoint.file_states {
171 match self.current_state.get(file_path) {
172 Some(actual_content) => {
173 if actual_content != expected_content {
174 return Ok(false);
175 }
176 }
177 None => return Ok(false),
178 }
179 }
180
181 Ok(true)
182 }
183
184 pub fn restore_pre_rollback_state(&mut self, pre_rollback_state: HashMap<String, String>) {
186 self.current_state = pre_rollback_state;
187 }
188
189 pub fn get_current_state(&self) -> HashMap<String, String> {
191 self.current_state.clone()
192 }
193
194 pub fn set_current_state(&mut self, state: HashMap<String, String>) {
196 self.current_state = state;
197 }
198
199 pub fn checkpoint_count(&self) -> usize {
201 self.checkpoints.len()
202 }
203}
204
205impl Default for CheckpointManager {
206 fn default() -> Self {
207 Self::new()
208 }
209}
210
211#[cfg(test)]
212mod tests {
213 use super::*;
214
215 #[test]
216 fn test_checkpoint_create_valid() {
217 let mut file_states = HashMap::new();
218 file_states.insert("file1.txt".to_string(), "content1".to_string());
219 let checkpoint = Checkpoint::new("Test Checkpoint", "A test checkpoint", file_states);
220 assert!(checkpoint.is_ok());
221 let checkpoint = checkpoint.unwrap();
222 assert_eq!(checkpoint.name, "Test Checkpoint");
223 assert_eq!(checkpoint.changes_count, 1);
224 }
225
226 #[test]
227 fn test_checkpoint_empty_name() {
228 let mut file_states = HashMap::new();
229 file_states.insert("file1.txt".to_string(), "content1".to_string());
230 let checkpoint = Checkpoint::new("", "description", file_states);
231 assert!(checkpoint.is_err());
232 }
233
234 #[test]
235 fn test_checkpoint_empty_file_states() {
236 let file_states = HashMap::new();
237 let checkpoint = Checkpoint::new("Test", "description", file_states);
238 assert!(checkpoint.is_err());
239 }
240
241 #[test]
242 fn test_checkpoint_manager_create() {
243 let mut manager = CheckpointManager::new();
244 let mut file_states = HashMap::new();
245 file_states.insert("file1.txt".to_string(), "content1".to_string());
246 let result = manager.create_checkpoint("Test", "description", file_states);
247 assert!(result.is_ok());
248 assert_eq!(manager.checkpoint_count(), 1);
249 }
250
251 #[test]
252 fn test_checkpoint_manager_list() {
253 let mut manager = CheckpointManager::new();
254 let mut file_states1 = HashMap::new();
255 file_states1.insert("file1.txt".to_string(), "content1".to_string());
256 let mut file_states2 = HashMap::new();
257 file_states2.insert("file2.txt".to_string(), "content2".to_string());
258
259 manager
260 .create_checkpoint("Checkpoint 1", "desc1", file_states1)
261 .unwrap();
262 manager
263 .create_checkpoint("Checkpoint 2", "desc2", file_states2)
264 .unwrap();
265
266 let checkpoints = manager.list_checkpoints();
267 assert_eq!(checkpoints.len(), 2);
268 }
269
270 #[test]
271 fn test_checkpoint_manager_get() {
272 let mut manager = CheckpointManager::new();
273 let mut file_states = HashMap::new();
274 file_states.insert("file1.txt".to_string(), "content1".to_string());
275 let id = manager
276 .create_checkpoint("Test", "description", file_states)
277 .unwrap();
278
279 let checkpoint = manager.get_checkpoint(&id);
280 assert!(checkpoint.is_ok());
281 assert_eq!(checkpoint.unwrap().name, "Test");
282 }
283
284 #[test]
285 fn test_checkpoint_manager_get_not_found() {
286 let manager = CheckpointManager::new();
287 let checkpoint = manager.get_checkpoint("nonexistent");
288 assert!(checkpoint.is_err());
289 }
290
291 #[test]
292 fn test_checkpoint_manager_delete() {
293 let mut manager = CheckpointManager::new();
294 let mut file_states = HashMap::new();
295 file_states.insert("file1.txt".to_string(), "content1".to_string());
296 let id = manager
297 .create_checkpoint("Test", "description", file_states)
298 .unwrap();
299
300 assert_eq!(manager.checkpoint_count(), 1);
301 let result = manager.delete_checkpoint(&id);
302 assert!(result.is_ok());
303 assert_eq!(manager.checkpoint_count(), 0);
304 }
305
306 #[test]
307 fn test_checkpoint_serialization() {
308 let mut file_states = HashMap::new();
309 file_states.insert("file1.txt".to_string(), "content1".to_string());
310 let checkpoint = Checkpoint::new("Test", "description", file_states).unwrap();
311 let json = serde_json::to_string(&checkpoint).unwrap();
312 let deserialized: Checkpoint = serde_json::from_str(&json).unwrap();
313 assert_eq!(checkpoint.id, deserialized.id);
314 assert_eq!(checkpoint.name, deserialized.name);
315 }
316
317 #[test]
318 fn test_checkpoint_manager_rollback_to() {
319 let mut manager = CheckpointManager::new();
320
321 let mut initial_state = HashMap::new();
323 initial_state.insert("file1.txt".to_string(), "initial content".to_string());
324 manager.set_current_state(initial_state);
325
326 let mut checkpoint_state = HashMap::new();
328 checkpoint_state.insert("file1.txt".to_string(), "checkpoint content".to_string());
329 checkpoint_state.insert("file2.txt".to_string(), "new file".to_string());
330
331 let checkpoint_id = manager
332 .create_checkpoint("Checkpoint 1", "desc", checkpoint_state.clone())
333 .unwrap();
334
335 let result = manager.rollback_to(&checkpoint_id);
337 assert!(result.is_ok());
338
339 let current_state = manager.get_current_state();
341 assert_eq!(current_state.get("file1.txt"), Some(&"checkpoint content".to_string()));
342 assert_eq!(current_state.get("file2.txt"), Some(&"new file".to_string()));
343 }
344
345 #[test]
346 fn test_checkpoint_manager_rollback_not_found() {
347 let mut manager = CheckpointManager::new();
348 let result = manager.rollback_to("nonexistent");
349 assert!(result.is_err());
350 }
351
352 #[test]
353 fn test_checkpoint_manager_rollback_isolation() {
354 let mut manager = CheckpointManager::new();
355
356 let mut state1 = HashMap::new();
358 state1.insert("file.txt".to_string(), "state1".to_string());
359 let id1 = manager
360 .create_checkpoint("Checkpoint 1", "desc1", state1)
361 .unwrap();
362
363 let mut state2 = HashMap::new();
364 state2.insert("file.txt".to_string(), "state2".to_string());
365 let id2 = manager
366 .create_checkpoint("Checkpoint 2", "desc2", state2)
367 .unwrap();
368
369 manager.rollback_to(&id1).unwrap();
371 let current = manager.get_current_state();
372 assert_eq!(current.get("file.txt"), Some(&"state1".to_string()));
373
374 let cp2 = manager.get_checkpoint(&id2).unwrap();
376 assert_eq!(cp2.file_states.get("file.txt"), Some(&"state2".to_string()));
377
378 manager.rollback_to(&id2).unwrap();
380 let current = manager.get_current_state();
381 assert_eq!(current.get("file.txt"), Some(&"state2".to_string()));
382
383 let cp1 = manager.get_checkpoint(&id1).unwrap();
385 assert_eq!(cp1.file_states.get("file.txt"), Some(&"state1".to_string()));
386 }
387
388 #[test]
389 fn test_checkpoint_manager_restore_pre_rollback_state() {
390 let mut manager = CheckpointManager::new();
391
392 let mut initial_state = HashMap::new();
394 initial_state.insert("file.txt".to_string(), "initial".to_string());
395 manager.set_current_state(initial_state.clone());
396
397 let mut checkpoint_state = HashMap::new();
399 checkpoint_state.insert("file.txt".to_string(), "checkpoint".to_string());
400 let checkpoint_id = manager
401 .create_checkpoint("CP", "desc", checkpoint_state)
402 .unwrap();
403
404 manager.rollback_to(&checkpoint_id).unwrap();
406 assert_eq!(
407 manager.get_current_state().get("file.txt"),
408 Some(&"checkpoint".to_string())
409 );
410
411 manager.restore_pre_rollback_state(initial_state);
413 assert_eq!(
414 manager.get_current_state().get("file.txt"),
415 Some(&"initial".to_string())
416 );
417 }
418
419 #[test]
420 fn test_checkpoint_manager_verify_rollback_success() {
421 let mut manager = CheckpointManager::new();
422
423 let mut checkpoint_state = HashMap::new();
425 checkpoint_state.insert("file1.txt".to_string(), "content1".to_string());
426 checkpoint_state.insert("file2.txt".to_string(), "content2".to_string());
427 let checkpoint_id = manager
428 .create_checkpoint("CP", "desc", checkpoint_state)
429 .unwrap();
430
431 manager.rollback_to(&checkpoint_id).unwrap();
433
434 let result = manager.verify_rollback(&checkpoint_id);
436 assert!(result.is_ok());
437 assert!(result.unwrap());
438 }
439
440 #[test]
441 fn test_checkpoint_manager_verify_rollback_failure() {
442 let mut manager = CheckpointManager::new();
443
444 let mut checkpoint_state = HashMap::new();
446 checkpoint_state.insert("file.txt".to_string(), "content".to_string());
447 let checkpoint_id = manager
448 .create_checkpoint("CP", "desc", checkpoint_state)
449 .unwrap();
450
451 manager.rollback_to(&checkpoint_id).unwrap();
453
454 manager
456 .current_state
457 .insert("file.txt".to_string(), "modified".to_string());
458
459 let result = manager.verify_rollback(&checkpoint_id);
461 assert!(result.is_ok());
462 assert!(!result.unwrap());
463 }
464
465 #[test]
466 fn test_checkpoint_manager_rollback_failure_recovery() {
467 let mut manager = CheckpointManager::new();
468
469 let mut initial_state = HashMap::new();
471 initial_state.insert("file.txt".to_string(), "initial".to_string());
472 manager.set_current_state(initial_state.clone());
473
474 let mut checkpoint_state = HashMap::new();
476 checkpoint_state.insert("file.txt".to_string(), "checkpoint".to_string());
477 let checkpoint_id = manager
478 .create_checkpoint("CP", "desc", checkpoint_state)
479 .unwrap();
480
481 let result = manager.rollback_to(&checkpoint_id);
483 assert!(result.is_ok());
484
485 assert_eq!(
487 manager.get_current_state().get("file.txt"),
488 Some(&"checkpoint".to_string())
489 );
490
491 manager.restore_pre_rollback_state(initial_state);
493 assert_eq!(
494 manager.get_current_state().get("file.txt"),
495 Some(&"initial".to_string())
496 );
497 }
498
499 #[test]
500 fn test_checkpoint_manager_rollback_atomic_all_or_nothing() {
501 let mut manager = CheckpointManager::new();
502
503 let mut checkpoint_state = HashMap::new();
505 checkpoint_state.insert("file1.txt".to_string(), "content1".to_string());
506 checkpoint_state.insert("file2.txt".to_string(), "content2".to_string());
507 checkpoint_state.insert("file3.txt".to_string(), "content3".to_string());
508 let checkpoint_id = manager
509 .create_checkpoint("CP", "desc", checkpoint_state)
510 .unwrap();
511
512 let result = manager.rollback_to(&checkpoint_id);
514 assert!(result.is_ok());
515
516 let current_state = manager.get_current_state();
518 assert_eq!(current_state.len(), 3);
519 assert_eq!(current_state.get("file1.txt"), Some(&"content1".to_string()));
520 assert_eq!(current_state.get("file2.txt"), Some(&"content2".to_string()));
521 assert_eq!(current_state.get("file3.txt"), Some(&"content3".to_string()));
522 }
523
524 #[test]
525 fn test_checkpoint_isolation_independent_storage() {
526 let mut manager = CheckpointManager::new();
527
528 let mut state1 = HashMap::new();
530 state1.insert("file.txt".to_string(), "state1".to_string());
531 let id1 = manager
532 .create_checkpoint("CP1", "desc1", state1)
533 .unwrap();
534
535 let mut state2 = HashMap::new();
537 state2.insert("file.txt".to_string(), "state2".to_string());
538 let id2 = manager
539 .create_checkpoint("CP2", "desc2", state2)
540 .unwrap();
541
542 let cp1 = manager.get_checkpoint(&id1).unwrap();
544 let cp2 = manager.get_checkpoint(&id2).unwrap();
545
546 assert_eq!(cp1.file_states.get("file.txt"), Some(&"state1".to_string()));
547 assert_eq!(cp2.file_states.get("file.txt"), Some(&"state2".to_string()));
548
549 manager.rollback_to(&id1).unwrap();
551
552 let cp2_after = manager.get_checkpoint(&id2).unwrap();
554 assert_eq!(
555 cp2_after.file_states.get("file.txt"),
556 Some(&"state2".to_string())
557 );
558
559 assert_eq!(
561 manager.get_current_state().get("file.txt"),
562 Some(&"state1".to_string())
563 );
564 }
565
566 #[test]
567 fn test_checkpoint_isolation_prevent_corruption() {
568 let mut manager = CheckpointManager::new();
569
570 let mut checkpoint_state = HashMap::new();
572 checkpoint_state.insert("file1.txt".to_string(), "content1".to_string());
573 checkpoint_state.insert("file2.txt".to_string(), "content2".to_string());
574 let checkpoint_id = manager
575 .create_checkpoint("CP", "desc", checkpoint_state)
576 .unwrap();
577
578 let checkpoint = manager.get_checkpoint(&checkpoint_id).unwrap();
580 assert!(checkpoint.validate().is_ok());
581
582 manager.rollback_to(&checkpoint_id).unwrap();
584
585 let checkpoint_after = manager.get_checkpoint(&checkpoint_id).unwrap();
587 assert!(checkpoint_after.validate().is_ok());
588 assert_eq!(checkpoint.id, checkpoint_after.id);
589 assert_eq!(checkpoint.name, checkpoint_after.name);
590 assert_eq!(checkpoint.file_states, checkpoint_after.file_states);
591 }
592}
593
594#[cfg(test)]
595mod property_tests {
596 use super::*;
597 use proptest::prelude::*;
598
599 fn checkpoint_name_strategy() -> impl Strategy<Value = String> {
601 r"[a-zA-Z0-9\s\-_]{1,50}"
602 .prop_map(|s| s.to_string())
603 }
604
605 fn file_path_strategy() -> impl Strategy<Value = String> {
607 r"[a-zA-Z0-9_\-./]{1,50}\.rs"
608 .prop_map(|s| s.to_string())
609 }
610
611 fn content_strategy() -> impl Strategy<Value = String> {
613 r"[a-zA-Z0-9\s]{1,100}"
614 .prop_map(|s| s.to_string())
615 }
616
617 proptest! {
618 #[test]
623 fn prop_rollback_atomicity(
624 checkpoint_files in prop::collection::hash_map(
625 file_path_strategy(),
626 content_strategy(),
627 1..10
628 ),
629 name in checkpoint_name_strategy(),
630 ) {
631 let mut manager = CheckpointManager::new();
632
633 let checkpoint_id = manager
635 .create_checkpoint(&name, "description", checkpoint_files.clone())
636 .ok();
637
638 prop_assert!(checkpoint_id.is_some(), "Checkpoint creation should succeed");
639 let checkpoint_id = checkpoint_id.unwrap();
640
641 let rollback_result = manager.rollback_to(&checkpoint_id);
643 prop_assert!(rollback_result.is_ok(), "Rollback should succeed");
644
645 let current_state = manager.get_current_state();
647 for (file_path, expected_content) in &checkpoint_files {
648 let actual_content = current_state.get(file_path);
649 prop_assert_eq!(
650 actual_content,
651 Some(expected_content),
652 "File {} should be restored to checkpoint state",
653 file_path
654 );
655 }
656
657 prop_assert_eq!(
659 current_state.len(),
660 checkpoint_files.len(),
661 "Current state should have exactly the checkpoint files"
662 );
663 }
664
665 #[test]
670 fn prop_checkpoint_isolation(
671 checkpoint_data in prop::collection::vec(
672 (checkpoint_name_strategy(), prop::collection::hash_map(
673 file_path_strategy(),
674 content_strategy(),
675 1..5
676 )),
677 2..5
678 ),
679 ) {
680 let mut manager = CheckpointManager::new();
681 let mut checkpoint_ids = Vec::new();
682
683 for (name, files) in checkpoint_data.iter() {
685 prop_assume!(!files.is_empty());
686 if let Ok(id) = manager.create_checkpoint(name, "desc", files.clone()) {
687 checkpoint_ids.push((id, files.clone()));
688 }
689 }
690
691 prop_assume!(checkpoint_ids.len() >= 2);
692
693 for (rollback_idx, (rollback_id, _)) in checkpoint_ids.iter().enumerate() {
695 let rollback_result = manager.rollback_to(rollback_id);
697 prop_assert!(rollback_result.is_ok(), "Rollback should succeed");
698
699 for (other_idx, (other_id, other_files)) in checkpoint_ids.iter().enumerate() {
701 if rollback_idx != other_idx {
702 let checkpoint = manager.get_checkpoint(other_id);
703 prop_assert!(checkpoint.is_ok(), "Other checkpoint should still exist");
704
705 let checkpoint = checkpoint.unwrap();
706 for (file_path, expected_content) in other_files {
707 let actual_content = checkpoint.file_states.get(file_path);
708 prop_assert_eq!(
709 actual_content,
710 Some(expected_content),
711 "Other checkpoint file {} should be unchanged",
712 file_path
713 );
714 }
715 }
716 }
717 }
718 }
719
720 #[test]
724 fn prop_single_checkpoint_rollback(
725 files in prop::collection::hash_map(
726 file_path_strategy(),
727 content_strategy(),
728 1..5
729 ),
730 name in checkpoint_name_strategy(),
731 ) {
732 prop_assume!(!files.is_empty());
733
734 let mut manager = CheckpointManager::new();
735
736 let checkpoint_id = manager
738 .create_checkpoint(&name, "desc", files.clone())
739 .unwrap();
740
741 let checkpoint = manager.get_checkpoint(&checkpoint_id).unwrap();
743 prop_assert_eq!(
744 checkpoint.file_states.len(),
745 files.len(),
746 "Checkpoint should contain all files"
747 );
748
749 manager.rollback_to(&checkpoint_id).unwrap();
751
752 let current_state = manager.get_current_state();
754 prop_assert_eq!(
755 current_state.len(),
756 files.len(),
757 "Current state should have same number of files"
758 );
759
760 for (file_path, expected_content) in &files {
761 let actual_content = current_state.get(file_path);
762 prop_assert_eq!(
763 actual_content,
764 Some(expected_content),
765 "File {} should match checkpoint state",
766 file_path
767 );
768 }
769 }
770 }
771}