1use crate::error::{PackError, Result};
2use crate::format::PackedSnapshot;
3use crate::checkpoint::{Checkpoint, CheckpointManager};
4use std::collections::VecDeque;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum ReplayDirection {
8 Forward,
9 Backward,
10}
11
12pub struct ReplayEngine {
13 checkpoints: VecDeque<Checkpoint>,
14 current_index: usize,
15 loop_replay: bool,
16}
17
18impl ReplayEngine {
19 pub fn new() -> Self {
20 Self {
21 checkpoints: VecDeque::new(),
22 current_index: 0,
23 loop_replay: false,
24 }
25 }
26
27 pub fn with_loop(mut self, enabled: bool) -> Self {
28 self.loop_replay = enabled;
29 self
30 }
31
32 pub fn add_checkpoint(&mut self, checkpoint: Checkpoint) {
33 self.checkpoints.push_back(checkpoint);
34 }
35
36 pub fn load_from_manager(&mut self, manager: &mut CheckpointManager) -> Result<()> {
37 self.checkpoints.clear();
38
39 let chain = manager.get_checkpoint_chain().to_vec();
40 for id in chain {
41 let checkpoint = manager.load_checkpoint(&id)?;
42 self.checkpoints.push_back(checkpoint);
43 }
44
45 self.current_index = 0;
46
47 Ok(())
48 }
49
50 pub fn current(&self) -> Option<&Checkpoint> {
51 self.checkpoints.get(self.current_index)
52 }
53
54 pub fn next(&mut self) -> Option<&Checkpoint> {
55 if self.current_index + 1 < self.checkpoints.len() {
56 self.current_index += 1;
57 self.current()
58 } else if self.loop_replay && !self.checkpoints.is_empty() {
59 self.current_index = 0;
60 self.current()
61 } else {
62 None
63 }
64 }
65
66 pub fn previous(&mut self) -> Option<&Checkpoint> {
67 if self.current_index > 0 {
68 self.current_index -= 1;
69 self.current()
70 } else if self.loop_replay && !self.checkpoints.is_empty() {
71 self.current_index = self.checkpoints.len() - 1;
72 self.current()
73 } else {
74 None
75 }
76 }
77
78 pub fn seek(&mut self, index: usize) -> Result<&Checkpoint> {
79 if index >= self.checkpoints.len() {
80 return Err(PackError::InvalidCheckpoint(
81 format!("Index {} out of bounds", index)
82 ));
83 }
84
85 self.current_index = index;
86 self.current()
87 .ok_or_else(|| PackError::InvalidCheckpoint("No checkpoint at index".to_string()))
88 }
89
90 pub fn seek_to_start(&mut self) -> Option<&Checkpoint> {
91 self.current_index = 0;
92 self.current()
93 }
94
95 pub fn seek_to_end(&mut self) -> Option<&Checkpoint> {
96 if !self.checkpoints.is_empty() {
97 self.current_index = self.checkpoints.len() - 1;
98 }
99 self.current()
100 }
101
102 pub fn get_index(&self) -> usize {
103 self.current_index
104 }
105
106 pub fn len(&self) -> usize {
107 self.checkpoints.len()
108 }
109
110 pub fn is_empty(&self) -> bool {
111 self.checkpoints.is_empty()
112 }
113
114 pub fn is_at_start(&self) -> bool {
115 self.current_index == 0
116 }
117
118 pub fn is_at_end(&self) -> bool {
119 self.current_index == self.checkpoints.len().saturating_sub(1)
120 }
121
122 pub fn clear(&mut self) {
123 self.checkpoints.clear();
124 self.current_index = 0;
125 }
126}
127
128impl Default for ReplayEngine {
129 fn default() -> Self {
130 Self::new()
131 }
132}
133
134pub struct TimeTravel {
135 snapshots: Vec<(f64, PackedSnapshot)>,
136 current_time: f64,
137}
138
139impl TimeTravel {
140 pub fn new() -> Self {
141 Self {
142 snapshots: Vec::new(),
143 current_time: 0.0,
144 }
145 }
146
147 pub fn record(&mut self, time: f64, snapshot: PackedSnapshot) {
148 self.snapshots.push((time, snapshot));
149 self.snapshots.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
150 self.current_time = time;
151 }
152
153 pub fn seek_to_time(&mut self, target_time: f64) -> Option<&PackedSnapshot> {
154 let index = self.find_snapshot_at_time(target_time)?;
155 self.current_time = self.snapshots[index].0;
156 Some(&self.snapshots[index].1)
157 }
158
159 pub fn get_snapshot_at_time(&self, time: f64) -> Option<&PackedSnapshot> {
160 let index = self.find_snapshot_at_time(time)?;
161 Some(&self.snapshots[index].1)
162 }
163
164 pub fn get_current_snapshot(&self) -> Option<&PackedSnapshot> {
165 self.get_snapshot_at_time(self.current_time)
166 }
167
168 pub fn get_earliest_time(&self) -> Option<f64> {
169 self.snapshots.first().map(|(t, _)| *t)
170 }
171
172 pub fn get_latest_time(&self) -> Option<f64> {
173 self.snapshots.last().map(|(t, _)| *t)
174 }
175
176 pub fn get_current_time(&self) -> f64 {
177 self.current_time
178 }
179
180 pub fn fork_at_time(&self, time: f64) -> Option<PackedSnapshot> {
181 self.get_snapshot_at_time(time).cloned()
182 }
183
184 pub fn prune_before(&mut self, time: f64) {
185 self.snapshots.retain(|(t, _)| *t >= time);
186 }
187
188 pub fn prune_after(&mut self, time: f64) {
189 self.snapshots.retain(|(t, _)| *t <= time);
190 }
191
192 pub fn clear(&mut self) {
193 self.snapshots.clear();
194 self.current_time = 0.0;
195 }
196
197 pub fn len(&self) -> usize {
198 self.snapshots.len()
199 }
200
201 pub fn is_empty(&self) -> bool {
202 self.snapshots.is_empty()
203 }
204
205 fn find_snapshot_at_time(&self, target_time: f64) -> Option<usize> {
206 if self.snapshots.is_empty() {
207 return None;
208 }
209
210 let mut left = 0;
211 let mut right = self.snapshots.len();
212
213 while left < right {
214 let mid = (left + right) / 2;
215 if self.snapshots[mid].0 < target_time {
216 left = mid + 1;
217 } else {
218 right = mid;
219 }
220 }
221
222 if left > 0 && (left >= self.snapshots.len() ||
223 (self.snapshots[left].0 - target_time).abs() >
224 (target_time - self.snapshots[left - 1].0).abs()) {
225 Some(left - 1)
226 } else if left < self.snapshots.len() {
227 Some(left)
228 } else if left > 0 {
229 Some(left - 1)
230 } else {
231 None
232 }
233 }
234}
235
236impl Default for TimeTravel {
237 fn default() -> Self {
238 Self::new()
239 }
240}
241
242#[cfg(test)]
243mod tests {
244 use super::*;
245
246 #[test]
247 fn test_replay_engine() {
248 let mut engine = ReplayEngine::new();
249
250 for i in 0..5 {
251 let checkpoint = Checkpoint::new(format!("cp{}", i), PackedSnapshot::new());
252 engine.add_checkpoint(checkpoint);
253 }
254
255 assert_eq!(engine.len(), 5);
256 assert_eq!(engine.get_index(), 0);
257
258 engine.next();
259 assert_eq!(engine.get_index(), 1);
260
261 engine.previous();
262 assert_eq!(engine.get_index(), 0);
263
264 engine.seek_to_end();
265 assert_eq!(engine.get_index(), 4);
266 assert!(engine.is_at_end());
267
268 engine.seek_to_start();
269 assert_eq!(engine.get_index(), 0);
270 assert!(engine.is_at_start());
271 }
272
273 #[test]
274 fn test_replay_loop() {
275 let mut engine = ReplayEngine::new().with_loop(true);
276
277 for i in 0..3 {
278 let checkpoint = Checkpoint::new(format!("cp{}", i), PackedSnapshot::new());
279 engine.add_checkpoint(checkpoint);
280 }
281
282 engine.seek_to_end();
283 assert_eq!(engine.get_index(), 2);
284
285 engine.next();
286 assert_eq!(engine.get_index(), 0);
287
288 engine.seek_to_start();
289 assert_eq!(engine.get_index(), 0);
290
291 engine.previous();
292 assert_eq!(engine.get_index(), 2);
293 }
294
295 #[test]
296 fn test_time_travel() {
297 let mut tt = TimeTravel::new();
298
299 for i in 0..10 {
300 let snapshot = PackedSnapshot::new();
301 tt.record(i as f64 * 10.0, snapshot);
302 }
303
304 assert_eq!(tt.len(), 10);
305 assert_eq!(tt.get_earliest_time(), Some(0.0));
306 assert_eq!(tt.get_latest_time(), Some(90.0));
307
308 let snapshot = tt.seek_to_time(45.0);
309 assert!(snapshot.is_some());
310 assert_eq!(tt.get_current_time(), 50.0);
311
312 tt.prune_before(30.0);
313 assert_eq!(tt.len(), 7);
314 assert_eq!(tt.get_earliest_time(), Some(30.0));
315
316 tt.prune_after(70.0);
317 assert_eq!(tt.len(), 5);
318 assert_eq!(tt.get_latest_time(), Some(70.0));
319 }
320
321 #[test]
322 fn test_time_travel_fork() {
323 let mut tt = TimeTravel::new();
324
325 for i in 0..5 {
326 let snapshot = PackedSnapshot::new();
327 tt.record(i as f64 * 10.0, snapshot);
328 }
329
330 let forked = tt.fork_at_time(20.0);
331 assert!(forked.is_some());
332 }
333}