1 files changed, 146 insertions, 30 deletions
diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index 1617d61c8752..5e6526418b3b 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -8,10 +8,12 @@
 /// what they are interested in observing.
 use std::io::Write;
 use std::rc::Rc;
+use std::time::Instant;
 use tabwriter::TabWriter;
 
 use crate::chunk::Chunk;
-use crate::opcode::{CodeIdx, OpCode};
+use crate::generators::VMRequest;
+use crate::opcode::{CodeIdx, Op};
 use crate::value::Lambda;
 use crate::SourceCode;
 use crate::Value;
@@ -39,10 +41,25 @@ pub trait CompilerObserver {
 /// the Tvix virtual machine at runtime.
 pub trait RuntimeObserver {
     /// Called when the runtime enters a new call frame.
-    fn observe_enter_frame(&mut self, _arg_count: usize, _: &Rc<Lambda>, _call_depth: usize) {}
+    fn observe_enter_call_frame(&mut self, _arg_count: usize, _: &Rc<Lambda>, _call_depth: usize) {}
 
     /// Called when the runtime exits a call frame.
-    fn observe_exit_frame(&mut self, _frame_at: usize, _stack: &[Value]) {}
+    fn observe_exit_call_frame(&mut self, _frame_at: usize, _stack: &[Value]) {}
+
+    /// Called when the runtime suspends a call frame.
+    fn observe_suspend_call_frame(&mut self, _frame_at: usize, _stack: &[Value]) {}
+
+    /// Called when the runtime enters a generator frame.
+    fn observe_enter_generator(&mut self, _frame_at: usize, _name: &str, _stack: &[Value]) {}
+
+    /// Called when the runtime exits a generator frame.
+    fn observe_exit_generator(&mut self, _frame_at: usize, _name: &str, _stack: &[Value]) {}
+
+    /// Called when the runtime suspends a generator frame.
+    fn observe_suspend_generator(&mut self, _frame_at: usize, _name: &str, _stack: &[Value]) {}
+
+    /// Called when a generator requests an action from the VM.
+    fn observe_generator_request(&mut self, _name: &str, _msg: &VMRequest) {}
 
     /// Called when the runtime replaces the current call frame for a
     /// tail call.
@@ -56,7 +73,7 @@ pub trait RuntimeObserver {
 
     /// Called when the runtime *begins* executing an instruction. The
     /// provided stack is the state at the beginning of the operation.
-    fn observe_execute_op(&mut self, _ip: CodeIdx, _: &OpCode, _: &[Value]) {}
+    fn observe_execute_op(&mut self, _ip: CodeIdx, _: &Op, _: &[Value]) {}
 }
 
 #[derive(Default)]
@@ -95,8 +112,12 @@ impl<W: Write> DisassemblingObserver<W> {
         // calculate width of the widest address in the chunk
         let width = format!("{:#x}", chunk.code.len() - 1).len();
 
-        for (idx, _) in chunk.code.iter().enumerate() {
-            let _ = chunk.disassemble_op(&mut self.writer, &self.source, width, CodeIdx(idx));
+        let mut idx = 0;
+        while idx < chunk.code.len() {
+            let size = chunk
+                .disassemble_op(&mut self.writer, &self.source, width, CodeIdx(idx))
+                .expect("writing debug output should work");
+            idx += size;
         }
     }
 }
@@ -124,19 +145,74 @@ impl<W: Write> CompilerObserver for DisassemblingObserver<W> {
 /// An observer that collects a textual representation of an entire
 /// runtime execution.
 pub struct TracingObserver<W: Write> {
+    // If timing is enabled, contains the timestamp of the last-emitted trace event
+    last_event: Option<Instant>,
     writer: TabWriter<W>,
 }
 
 impl<W: Write> TracingObserver<W> {
     pub fn new(writer: W) -> Self {
         Self {
+            last_event: None,
             writer: TabWriter::new(writer),
         }
     }
+
+    /// Write the time of each runtime event, relative to when this method is called
+    pub fn enable_timing(&mut self) {
+        self.last_event = Some(Instant::now());
+    }
+
+    fn maybe_write_time(&mut self) {
+        if let Some(last_event) = &mut self.last_event {
+            let _ = write!(&mut self.writer, "+{}ns\t", last_event.elapsed().as_nanos());
+            *last_event = Instant::now();
+        }
+    }
+
+    fn write_value(&mut self, val: &Value) {
+        let _ = match val {
+            // Potentially large types which we only want to print
+            // the type of (and avoid recursing).
+            Value::List(l) => write!(&mut self.writer, "list[{}] ", l.len()),
+            Value::Attrs(a) => write!(&mut self.writer, "attrs[{}] ", a.len()),
+            Value::Thunk(t) if t.is_evaluated() => {
+                self.write_value(&t.value());
+                Ok(())
+            }
+
+            // For other value types, defer to the standard value printer.
+            _ => write!(&mut self.writer, "{} ", val),
+        };
+    }
+
+    fn write_stack(&mut self, stack: &[Value]) {
+        let _ = write!(&mut self.writer, "[ ");
+
+        // Print out a maximum of 6 values from the top of the stack,
+        // before abbreviating it to `...`.
+        for (i, val) in stack.iter().rev().enumerate() {
+            if i == 6 {
+                let _ = write!(&mut self.writer, "...");
+                break;
+            }
+
+            self.write_value(val);
+        }
+
+        let _ = writeln!(&mut self.writer, "]");
+    }
 }
 
 impl<W: Write> RuntimeObserver for TracingObserver<W> {
-    fn observe_enter_frame(&mut self, arg_count: usize, lambda: &Rc<Lambda>, call_depth: usize) {
+    fn observe_enter_call_frame(
+        &mut self,
+        arg_count: usize,
+        lambda: &Rc<Lambda>,
+        call_depth: usize,
+    ) {
+        self.maybe_write_time();
+
         let _ = write!(&mut self.writer, "=== entering ");
 
         let _ = if arg_count == 0 {
@@ -152,50 +228,90 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
         let _ = writeln!(
             &mut self.writer,
             "in frame[{}] @ {:p} ===",
-            call_depth, lambda
+            call_depth, *lambda
         );
     }
 
-    fn observe_exit_frame(&mut self, frame_at: usize, stack: &[Value]) {
-        let _ = write!(&mut self.writer, "=== exiting frame {} ===\t[ ", frame_at);
+    /// Called when the runtime exits a call frame.
+    fn observe_exit_call_frame(&mut self, frame_at: usize, stack: &[Value]) {
+        self.maybe_write_time();
+        let _ = write!(&mut self.writer, "=== exiting frame {} ===\t ", frame_at);
+        self.write_stack(stack);
+    }
 
-        for val in stack {
-            let _ = write!(&mut self.writer, "{} ", val);
-        }
+    fn observe_suspend_call_frame(&mut self, frame_at: usize, stack: &[Value]) {
+        self.maybe_write_time();
+        let _ = write!(&mut self.writer, "=== suspending frame {} ===\t", frame_at);
 
-        let _ = writeln!(&mut self.writer, "]");
+        self.write_stack(stack);
+    }
+
+    fn observe_enter_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        self.maybe_write_time();
+        let _ = write!(
+            &mut self.writer,
+            "=== entering generator frame '{}' [{}] ===\t",
+            name, frame_at,
+        );
+
+        self.write_stack(stack);
+    }
+
+    fn observe_exit_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        self.maybe_write_time();
+        let _ = write!(
+            &mut self.writer,
+            "=== exiting generator '{}' [{}] ===\t",
+            name, frame_at
+        );
+
+        self.write_stack(stack);
+    }
+
+    fn observe_suspend_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        self.maybe_write_time();
+        let _ = write!(
+            &mut self.writer,
+            "=== suspending generator '{}' [{}] ===\t",
+            name, frame_at
+        );
+
+        self.write_stack(stack);
+    }
+
+    fn observe_generator_request(&mut self, name: &str, msg: &VMRequest) {
+        self.maybe_write_time();
+        let _ = writeln!(
+            &mut self.writer,
+            "=== generator '{}' requested {} ===",
+            name, msg
+        );
     }
 
     fn observe_enter_builtin(&mut self, name: &'static str) {
+        self.maybe_write_time();
         let _ = writeln!(&mut self.writer, "=== entering builtin {} ===", name);
     }
 
     fn observe_exit_builtin(&mut self, name: &'static str, stack: &[Value]) {
-        let _ = write!(&mut self.writer, "=== exiting builtin {} ===\t[ ", name);
-
-        for val in stack {
-            let _ = write!(&mut self.writer, "{} ", val);
-        }
-
-        let _ = writeln!(&mut self.writer, "]");
+        self.maybe_write_time();
+        let _ = write!(&mut self.writer, "=== exiting builtin {} ===\t", name);
+        self.write_stack(stack);
     }
 
     fn observe_tail_call(&mut self, frame_at: usize, lambda: &Rc<Lambda>) {
+        self.maybe_write_time();
         let _ = writeln!(
             &mut self.writer,
             "=== tail-calling {:p} in frame[{}] ===",
-            lambda, frame_at
+            *lambda, frame_at
         );
     }
 
-    fn observe_execute_op(&mut self, ip: CodeIdx, op: &OpCode, stack: &[Value]) {
-        let _ = write!(&mut self.writer, "{:04} {:?}\t[ ", ip.0, op);
-
-        for val in stack {
-            let _ = write!(&mut self.writer, "{} ", val);
-        }
-
-        let _ = writeln!(&mut self.writer, "]");
+    fn observe_execute_op(&mut self, ip: CodeIdx, op: &Op, stack: &[Value]) {
+        self.maybe_write_time();
+        let _ = write!(&mut self.writer, "{:04} {:?}\t", ip.0, op);
+        self.write_stack(stack);
     }
 }