1use std::collections::{HashMap, HashSet};
2use std::io;
3use std::path::Path;
4use std::sync::{Arc, Mutex};
5
6use crate::vm::{Program, Value, Vm, VmExecutionFrameSnapshot, VmFrameContinuation, VmStatus};
7
8#[derive(Clone, Debug, PartialEq)]
9pub struct VmRecordingFrame {
10 pub ip: usize,
11 pub call_depth: usize,
12 pub execution_frames: Vec<VmExecutionFrameSnapshot>,
13 pub stack: Vec<Value>,
14 pub locals: Vec<Value>,
15}
16
17#[derive(Clone, Debug)]
18pub struct VmRecording {
19 pub program: Program,
20 pub frames: Vec<VmRecordingFrame>,
21 pub terminal_status: Option<VmStatus>,
22}
23
24#[derive(Clone, Debug, Default)]
25pub struct VmRecordingReplayState {
26 pub cursor: usize,
27 pub offset_breakpoints: HashSet<usize>,
28 pub line_breakpoints: HashSet<u32>,
29}
30
31#[derive(Clone, Debug)]
32pub struct VmRecordingReplayResponse {
33 pub output: String,
34 pub current_line: Option<u32>,
35 pub at_end: bool,
36 pub exited: bool,
37}
38
39#[derive(Debug)]
40pub enum VmRecordingError {
41 Io(io::Error),
42 Wire(crate::vmbc::WireError),
43 InvalidFormat(&'static str),
44 Message(String),
45}
46
47pub(super) struct VmRecordingBuilder {
48 recording: VmRecording,
49}
50
51impl VmRecordingFrame {
52 pub(super) fn from_vm(vm: &Vm) -> Self {
53 Self {
54 ip: vm.ip(),
55 call_depth: vm.call_depth(),
56 execution_frames: vm.execution_frames(),
57 stack: vm.stack().to_vec(),
58 locals: vm.locals().to_vec(),
59 }
60 }
61}
62
63impl VmRecordingBuilder {
64 pub(super) fn new(program: Program) -> Self {
65 Self {
66 recording: VmRecording {
67 program,
68 frames: Vec::new(),
69 terminal_status: None,
70 },
71 }
72 }
73
74 pub(super) fn record_state(&mut self, vm: &Vm) {
75 let frame = VmRecordingFrame::from_vm(vm);
76 if self.recording.frames.last() == Some(&frame) {
77 return;
78 }
79 self.recording.frames.push(frame);
80 }
81
82 pub(super) fn on_terminal_status(&mut self, vm: &Vm, status: VmStatus) {
83 self.record_state(vm);
84 self.recording.terminal_status = Some(status);
85 }
86
87 pub(super) fn finish(self) -> VmRecording {
88 self.recording
89 }
90}
91
92impl VmRecording {
93 pub fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<(), VmRecordingError> {
94 let bytes = self.encode()?;
95 std::fs::write(path, bytes).map_err(VmRecordingError::Io)
96 }
97
98 pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self, VmRecordingError> {
99 let bytes = std::fs::read(path).map_err(VmRecordingError::Io)?;
100 Self::decode(&bytes)
101 }
102
103 pub fn encode(&self) -> Result<Vec<u8>, VmRecordingError> {
104 const MAGIC: [u8; 4] = *b"PDRC";
105 const VERSION: u16 = 5;
106
107 let mut out = Vec::new();
108 out.extend_from_slice(&MAGIC);
109 out.extend_from_slice(&VERSION.to_le_bytes());
110
111 let program_bytes =
112 crate::vmbc::encode_program(&self.program).map_err(VmRecordingError::Wire)?;
113 write_u32_len(program_bytes.len(), &mut out)?;
114 out.extend_from_slice(&program_bytes);
115
116 let status_tag = match self.terminal_status {
117 Some(VmStatus::Halted) => 1u8,
118 Some(VmStatus::Yielded) => 2u8,
119 Some(VmStatus::Waiting(_)) => 3u8,
120 None => 0u8,
121 };
122 out.push(status_tag);
123 if let Some(VmStatus::Waiting(op_id)) = self.terminal_status {
124 out.extend_from_slice(&op_id.to_le_bytes());
125 }
126
127 write_u32_len(self.frames.len(), &mut out)?;
128 let mut value_context = ValueEncodeContext::default();
129 for frame in &self.frames {
130 write_u32_from_usize(frame.ip, &mut out)?;
131 write_u32_from_usize(frame.call_depth, &mut out)?;
132 write_u32_len(frame.execution_frames.len(), &mut out)?;
133 for execution_frame in &frame.execution_frames {
134 match execution_frame.continuation {
135 VmFrameContinuation::Halt => out.push(0),
136 VmFrameContinuation::ResumeBytecode { return_ip } => {
137 out.push(1);
138 write_u32_from_usize(return_ip, &mut out)?;
139 }
140 VmFrameContinuation::ReturnToHost => out.push(2),
141 }
142 write_u32_from_usize(execution_frame.operand_stack_base, &mut out)?;
143 write_u32_from_usize(execution_frame.local_base, &mut out)?;
144 write_u32_from_usize(execution_frame.local_count, &mut out)?;
145 match execution_frame.prototype_id {
146 Some(prototype_id) => {
147 out.push(1);
148 out.extend_from_slice(&prototype_id.to_le_bytes());
149 }
150 None => out.push(0),
151 }
152 }
153
154 write_u32_len(frame.stack.len(), &mut out)?;
155 for value in &frame.stack {
156 encode_value(value, &mut out, &mut value_context)?;
157 }
158
159 write_u32_len(frame.locals.len(), &mut out)?;
160 for value in &frame.locals {
161 encode_value(value, &mut out, &mut value_context)?;
162 }
163 }
164
165 Ok(out)
166 }
167
168 pub fn decode(bytes: &[u8]) -> Result<Self, VmRecordingError> {
169 const MAGIC: [u8; 4] = *b"PDRC";
170 const VERSION: u16 = 5;
171
172 let mut cursor = RecordingCursor::new(bytes);
173
174 let magic = cursor.read_exact(4)?;
175 if magic != MAGIC {
176 return Err(VmRecordingError::InvalidFormat("invalid recording magic"));
177 }
178
179 let version = cursor.read_u16()?;
180 if version != VERSION {
181 return Err(VmRecordingError::Message(format!(
182 "unsupported recording version {version}"
183 )));
184 }
185
186 let program_len = cursor.read_u32()? as usize;
187 let program_bytes = cursor.read_exact(program_len)?;
188 let program = crate::vmbc::decode_program(program_bytes).map_err(VmRecordingError::Wire)?;
189
190 let terminal_status = match cursor.read_u8()? {
191 0 => None,
192 1 => Some(VmStatus::Halted),
193 2 => Some(VmStatus::Yielded),
194 3 if version >= VERSION => {
195 let op_id = cursor.read_u64()?;
196 Some(VmStatus::Waiting(op_id))
197 }
198 _ => {
199 return Err(VmRecordingError::InvalidFormat(
200 "invalid terminal status tag",
201 ));
202 }
203 };
204
205 let frame_count = cursor.read_u32()? as usize;
206 let mut frames = Vec::with_capacity(frame_count);
207 let mut value_context = ValueDecodeContext::default();
208 for _ in 0..frame_count {
209 let ip = cursor.read_u32()? as usize;
210 let call_depth = cursor.read_u32()? as usize;
211 let execution_frame_count = cursor.read_u32()? as usize;
212 let mut execution_frames = Vec::with_capacity(execution_frame_count);
213 for _ in 0..execution_frame_count {
214 let continuation = match cursor.read_u8()? {
215 0 => VmFrameContinuation::Halt,
216 1 => VmFrameContinuation::ResumeBytecode {
217 return_ip: cursor.read_u32()? as usize,
218 },
219 2 => VmFrameContinuation::ReturnToHost,
220 _ => {
221 return Err(VmRecordingError::InvalidFormat(
222 "invalid frame continuation tag",
223 ));
224 }
225 };
226 let operand_stack_base = cursor.read_u32()? as usize;
227 let local_base = cursor.read_u32()? as usize;
228 let local_count = cursor.read_u32()? as usize;
229 let prototype_id = match cursor.read_u8()? {
230 0 => None,
231 1 => Some(cursor.read_u32()?),
232 _ => {
233 return Err(VmRecordingError::InvalidFormat(
234 "invalid frame prototype tag",
235 ));
236 }
237 };
238 execution_frames.push(VmExecutionFrameSnapshot {
239 continuation,
240 operand_stack_base,
241 local_base,
242 local_count,
243 prototype_id,
244 });
245 }
246
247 let stack_len = cursor.read_u32()? as usize;
248 let mut stack = Vec::with_capacity(stack_len);
249 for _ in 0..stack_len {
250 stack.push(decode_value(&mut cursor, &mut value_context)?);
251 }
252
253 let locals_len = cursor.read_u32()? as usize;
254 let mut locals = Vec::with_capacity(locals_len);
255 for _ in 0..locals_len {
256 locals.push(decode_value(&mut cursor, &mut value_context)?);
257 }
258
259 frames.push(VmRecordingFrame {
260 ip,
261 call_depth,
262 execution_frames,
263 stack,
264 locals,
265 });
266 }
267
268 if !cursor.is_at_end() {
269 return Err(VmRecordingError::InvalidFormat(
270 "trailing bytes in recording payload",
271 ));
272 }
273
274 Ok(Self {
275 program,
276 frames,
277 terminal_status,
278 })
279 }
280}
281
282impl std::fmt::Display for VmRecordingError {
283 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
284 match self {
285 VmRecordingError::Io(err) => write!(f, "{err}"),
286 VmRecordingError::Wire(err) => write!(f, "{err}"),
287 VmRecordingError::InvalidFormat(message) => write!(f, "{message}"),
288 VmRecordingError::Message(message) => write!(f, "{message}"),
289 }
290 }
291}
292
293impl std::error::Error for VmRecordingError {}
294
295fn write_u32_len(len: usize, out: &mut Vec<u8>) -> Result<(), VmRecordingError> {
296 let value = u32::try_from(len)
297 .map_err(|_| VmRecordingError::Message(format!("length too large: {len}")))?;
298 out.extend_from_slice(&value.to_le_bytes());
299 Ok(())
300}
301
302fn write_u32_from_usize(value: usize, out: &mut Vec<u8>) -> Result<(), VmRecordingError> {
303 let value = u32::try_from(value)
304 .map_err(|_| VmRecordingError::Message(format!("value too large: {value}")))?;
305 out.extend_from_slice(&value.to_le_bytes());
306 Ok(())
307}
308
309#[derive(Default)]
310struct ValueEncodeContext {
311 environment_ids: HashMap<usize, u32>,
312 next_environment_id: u32,
313}
314
315#[derive(Default)]
316struct ValueDecodeContext {
317 environments: HashMap<u32, Arc<crate::CallableEnvironment>>,
318}
319
320fn encode_value(
321 value: &Value,
322 out: &mut Vec<u8>,
323 context: &mut ValueEncodeContext,
324) -> Result<(), VmRecordingError> {
325 match value {
326 Value::Null => out.push(0),
327 Value::Int(value) => {
328 out.push(1);
329 out.extend_from_slice(&value.to_le_bytes());
330 }
331 Value::Float(value) => {
332 out.push(2);
333 out.extend_from_slice(&value.to_le_bytes());
334 }
335 Value::Bool(false) => out.push(3),
336 Value::Bool(true) => out.push(4),
337 Value::String(value) => {
338 out.push(5);
339 write_u32_len(value.len(), out)?;
340 out.extend_from_slice(value.as_bytes());
341 }
342 Value::Bytes(value) => {
343 out.push(6);
344 write_u32_len(value.len(), out)?;
345 out.extend_from_slice(value.as_slice());
346 }
347 Value::Array(values) => {
348 out.push(7);
349 write_u32_len(values.len(), out)?;
350 for value in values.iter() {
351 encode_value(value, out, context)?;
352 }
353 }
354 Value::Map(entries) => {
355 out.push(8);
356 write_u32_len(entries.len(), out)?;
357 for (key, value) in entries.iter() {
358 encode_value(key, out, context)?;
359 encode_value(value, out, context)?;
360 }
361 }
362 Value::Callable(callable) => {
363 out.push(9);
364 out.extend_from_slice(&callable.prototype_id.to_le_bytes());
365 out.push(match callable.kind {
366 crate::CallableKind::FunctionItem => 0,
367 crate::CallableKind::Closure => 1,
368 crate::CallableKind::HostFunction => 2,
369 });
370 if let Some(env) = &callable.env {
371 let environment_key = Arc::as_ptr(env) as usize;
372 if let Some(environment_id) = context.environment_ids.get(&environment_key) {
373 out.push(2);
374 out.extend_from_slice(&environment_id.to_le_bytes());
375 } else {
376 let environment_id = context.next_environment_id;
377 context.next_environment_id =
378 context.next_environment_id.checked_add(1).ok_or(
379 VmRecordingError::InvalidFormat("too many callable environments"),
380 )?;
381 context
382 .environment_ids
383 .insert(environment_key, environment_id);
384 out.push(1);
385 out.extend_from_slice(&environment_id.to_le_bytes());
386 let cells = env.cells.lock().map_err(|_| {
387 VmRecordingError::InvalidFormat("poisoned callable environment")
388 })?;
389 write_u32_len(cells.len(), out)?;
390 for cell in cells.iter() {
391 let value = cell.lock().map_err(|_| {
392 VmRecordingError::InvalidFormat("poisoned callable capture cell")
393 })?;
394 encode_value(&value, out, context)?;
395 }
396 }
397 } else {
398 out.push(0);
399 }
400 }
401 }
402 Ok(())
403}
404
405fn decode_value(
406 cursor: &mut RecordingCursor<'_>,
407 context: &mut ValueDecodeContext,
408) -> Result<Value, VmRecordingError> {
409 match cursor.read_u8()? {
410 0 => Ok(Value::Null),
411 1 => Ok(Value::Int(cursor.read_i64()?)),
412 2 => Ok(Value::Float(cursor.read_f64()?)),
413 3 => Ok(Value::Bool(false)),
414 4 => Ok(Value::Bool(true)),
415 5 => {
416 let len = cursor.read_u32()? as usize;
417 let bytes = cursor.read_exact(len)?;
418 let value = std::str::from_utf8(bytes)
419 .map_err(|_| VmRecordingError::InvalidFormat("invalid utf-8 string"))?;
420 Ok(Value::string(value))
421 }
422 6 => {
423 let len = cursor.read_u32()? as usize;
424 Ok(Value::bytes(cursor.read_exact(len)?.to_vec()))
425 }
426 7 => {
427 let len = cursor.read_u32()? as usize;
428 let mut values = Vec::with_capacity(len);
429 for _ in 0..len {
430 values.push(decode_value(cursor, context)?);
431 }
432 Ok(Value::Array(values.into()))
433 }
434 8 => {
435 let len = cursor.read_u32()? as usize;
436 let mut entries = Vec::with_capacity(len);
437 for _ in 0..len {
438 let key = decode_value(cursor, context)?;
439 let value = decode_value(cursor, context)?;
440 entries.push((key, value));
441 }
442 Ok(Value::map(entries))
443 }
444 9 => {
445 let prototype_id = cursor.read_u32()?;
446 let kind = match cursor.read_u8()? {
447 0 => crate::CallableKind::FunctionItem,
448 1 => crate::CallableKind::Closure,
449 2 => crate::CallableKind::HostFunction,
450 _ => return Err(VmRecordingError::InvalidFormat("invalid callable kind")),
451 };
452 let env = match cursor.read_u8()? {
453 0 => None,
454 1 => {
455 let environment_id = cursor.read_u32()?;
456 if context.environments.contains_key(&environment_id) {
457 return Err(VmRecordingError::InvalidFormat(
458 "duplicate callable environment id",
459 ));
460 }
461 let environment = Arc::new(crate::CallableEnvironment {
462 cells: Mutex::new(Vec::new()),
463 });
464 context
465 .environments
466 .insert(environment_id, environment.clone());
467 let len = cursor.read_u32()? as usize;
468 let mut cells = Vec::with_capacity(len);
469 for _ in 0..len {
470 cells.push(Arc::new(Mutex::new(decode_value(cursor, context)?)));
471 }
472 *environment.cells.lock().map_err(|_| {
473 VmRecordingError::InvalidFormat("poisoned callable environment")
474 })? = cells;
475 Some(environment)
476 }
477 2 => {
478 let environment_id = cursor.read_u32()?;
479 Some(context.environments.get(&environment_id).cloned().ok_or(
480 VmRecordingError::InvalidFormat("unknown callable environment id"),
481 )?)
482 }
483 _ => return Err(VmRecordingError::InvalidFormat("invalid callable env flag")),
484 };
485 Ok(Value::Callable(Arc::new(crate::CallableValue {
486 prototype_id,
487 kind,
488 env,
489 })))
490 }
491 _ => Err(VmRecordingError::InvalidFormat("invalid value tag")),
492 }
493}
494
495struct RecordingCursor<'a> {
496 bytes: &'a [u8],
497 offset: usize,
498}
499
500impl<'a> RecordingCursor<'a> {
501 fn new(bytes: &'a [u8]) -> Self {
502 Self { bytes, offset: 0 }
503 }
504
505 fn is_at_end(&self) -> bool {
506 self.offset == self.bytes.len()
507 }
508
509 fn read_exact(&mut self, len: usize) -> Result<&'a [u8], VmRecordingError> {
510 if self.offset + len > self.bytes.len() {
511 return Err(VmRecordingError::InvalidFormat(
512 "unexpected end of recording payload",
513 ));
514 }
515 let bytes = &self.bytes[self.offset..self.offset + len];
516 self.offset += len;
517 Ok(bytes)
518 }
519
520 fn read_u8(&mut self) -> Result<u8, VmRecordingError> {
521 Ok(self.read_exact(1)?[0])
522 }
523
524 fn read_u16(&mut self) -> Result<u16, VmRecordingError> {
525 let mut bytes = [0u8; 2];
526 bytes.copy_from_slice(self.read_exact(2)?);
527 Ok(u16::from_le_bytes(bytes))
528 }
529
530 fn read_u32(&mut self) -> Result<u32, VmRecordingError> {
531 let mut bytes = [0u8; 4];
532 bytes.copy_from_slice(self.read_exact(4)?);
533 Ok(u32::from_le_bytes(bytes))
534 }
535
536 fn read_u64(&mut self) -> Result<u64, VmRecordingError> {
537 let mut bytes = [0u8; 8];
538 bytes.copy_from_slice(self.read_exact(8)?);
539 Ok(u64::from_le_bytes(bytes))
540 }
541
542 fn read_i64(&mut self) -> Result<i64, VmRecordingError> {
543 let mut bytes = [0u8; 8];
544 bytes.copy_from_slice(self.read_exact(8)?);
545 Ok(i64::from_le_bytes(bytes))
546 }
547
548 fn read_f64(&mut self) -> Result<f64, VmRecordingError> {
549 let mut bytes = [0u8; 8];
550 bytes.copy_from_slice(self.read_exact(8)?);
551 Ok(f64::from_le_bytes(bytes))
552 }
553}
554
555#[cfg(test)]
556mod tests {
557 use super::*;
558
559 #[test]
560 fn recording_preserves_callable_environment_aliases() {
561 let environment = Arc::new(crate::CallableEnvironment {
562 cells: Mutex::new(vec![Arc::new(Mutex::new(Value::Int(7)))]),
563 });
564 let callable = |prototype_id| {
565 Value::Callable(Arc::new(crate::CallableValue {
566 prototype_id,
567 kind: crate::CallableKind::Closure,
568 env: Some(environment.clone()),
569 }))
570 };
571 let recording = VmRecording {
572 program: Program::new(Vec::new(), vec![crate::OpCode::Ret as u8]),
573 frames: vec![VmRecordingFrame {
574 ip: 0,
575 call_depth: 0,
576 execution_frames: Vec::new(),
577 stack: vec![callable(1), callable(2)],
578 locals: Vec::new(),
579 }],
580 terminal_status: None,
581 };
582
583 let decoded = VmRecording::decode(&recording.encode().expect("encode")).expect("decode");
584 let Value::Callable(first) = &decoded.frames[0].stack[0] else {
585 panic!("first value should be callable");
586 };
587 let Value::Callable(second) = &decoded.frames[0].stack[1] else {
588 panic!("second value should be callable");
589 };
590 assert!(Arc::ptr_eq(
591 first.env.as_ref().expect("first env"),
592 second.env.as_ref().expect("second env")
593 ));
594 }
595
596 #[test]
597 fn recording_v5_rejects_legacy_versions() {
598 let recording = VmRecording {
599 program: Program::new(Vec::new(), vec![crate::OpCode::Ret as u8]),
600 frames: Vec::new(),
601 terminal_status: Some(VmStatus::Halted),
602 };
603 let bytes = recording.encode().expect("recording should encode");
604 assert_eq!(u16::from_le_bytes([bytes[4], bytes[5]]), 5);
605
606 for legacy in [1u16, 2u16, 3u16, 4u16] {
607 let mut legacy_bytes = bytes.clone();
608 legacy_bytes[4..6].copy_from_slice(&legacy.to_le_bytes());
609 assert!(matches!(
610 VmRecording::decode(&legacy_bytes),
611 Err(VmRecordingError::Message(message))
612 if message == format!("unsupported recording version {legacy}")
613 ));
614 }
615 }
616}