feat(tvix/eval): give generators human-readable names r/5978

This adds static strings to generator frames that describe the
generator in a human-readable fashion, which are then logged in
observers.

This makes runtime traces very precise, explaining exactly what is
being requested from where.

Change-Id: I695659a6bd0b7b0bdee75bc8049651f62b150e0c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8206
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>

5 files changed, 80 insertions, 51 deletions

diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index 29e4fb2d9c06..a606eb136362 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -49,16 +49,16 @@ pub trait RuntimeObserver {
     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, _stack: &[Value]) {}
+    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, _stack: &[Value]) {}
+    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, _stack: &[Value]) {}
+    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, _msg: &GeneratorRequest) {}
+    fn observe_generator_request(&mut self, _name: &str, _msg: &GeneratorRequest) {}
 
     /// Called when the runtime replaces the current call frame for a
     /// tail call.
@@ -219,34 +219,42 @@ impl<W: Write> RuntimeObserver for TracingObserver<W> {
         self.write_stack(stack);
     }
 
-    fn observe_enter_generator(&mut self, frame_at: usize, stack: &[Value]) {
+    fn observe_enter_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
         let _ = write!(
             &mut self.writer,
-            "=== entering generator frame {} ===\t",
-            frame_at
+            "=== entering generator frame '{}' [{}] ===\t",
+            name, frame_at,
         );
 
         self.write_stack(stack);
     }
 
-    fn observe_exit_generator(&mut self, frame_at: usize, stack: &[Value]) {
-        let _ = write!(&mut self.writer, "=== exiting generator {} ===\t", frame_at);
+    fn observe_exit_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
+        let _ = write!(
+            &mut self.writer,
+            "=== exiting generator '{}' [{}] ===\t",
+            name, frame_at
+        );
 
         self.write_stack(stack);
     }
 
-    fn observe_suspend_generator(&mut self, frame_at: usize, stack: &[Value]) {
+    fn observe_suspend_generator(&mut self, frame_at: usize, name: &str, stack: &[Value]) {
         let _ = write!(
             &mut self.writer,
-            "=== suspending generator {} ===\t",
-            frame_at
+            "=== suspending generator '{}' [{}] ===\t",
+            name, frame_at
         );
 
         self.write_stack(stack);
     }
 
-    fn observe_generator_request(&mut self, msg: &GeneratorRequest) {
-        let _ = writeln!(&mut self.writer, "=== generator requested {} ===", msg);
+    fn observe_generator_request(&mut self, name: &str, msg: &GeneratorRequest) {
+        let _ = writeln!(
+            &mut self.writer,
+            "=== generator '{}' requested {} ===",
+            name, msg
+        );
     }
 
     fn observe_enter_builtin(&mut self, name: &'static str) {
diff --git a/tvix/eval/src/value/builtin.rs b/tvix/eval/src/value/builtin.rs
index 0577111030a2..6d08ebf9506d 100644
--- a/tvix/eval/src/value/builtin.rs
+++ b/tvix/eval/src/value/builtin.rs
@@ -41,7 +41,7 @@ pub enum BuiltinResult {
     Partial(Builtin),
 
     /// Builtin was called and constructed a generator that the VM must run.
-    Called(Generator),
+    Called(&'static str, Generator),
 }
 
 /// Represents a single built-in function which directly executes Rust
@@ -105,7 +105,7 @@ impl Builtin {
     /// applied or return the builtin if it is partially applied.
     pub fn call(self) -> BuiltinResult {
         if self.0.partials.len() == self.0.arg_count {
-            BuiltinResult::Called((self.0.func)(self.0.partials))
+            BuiltinResult::Called(self.0.name, (self.0.func)(self.0.partials))
         } else {
             BuiltinResult::Partial(self)
         }
diff --git a/tvix/eval/src/vm/generators.rs b/tvix/eval/src/vm/generators.rs
index 1a9743dbd960..7a1ba0153d24 100644
--- a/tvix/eval/src/vm/generators.rs
+++ b/tvix/eval/src/vm/generators.rs
@@ -217,8 +217,9 @@ pub fn pin_generator(
 impl<'o> VM<'o> {
     /// Helper function to re-enqueue the current generator while it
     /// is awaiting a value.
-    fn reenqueue_generator(&mut self, span: LightSpan, generator: Generator) {
+    fn reenqueue_generator(&mut self, name: &'static str, span: LightSpan, generator: Generator) {
         self.frames.push(Frame::Generator {
+            name,
             generator,
             span,
             state: GeneratorState::AwaitingValue,
@@ -226,12 +227,13 @@ impl<'o> VM<'o> {
     }
 
     /// Helper function to enqueue a new generator.
-    pub(super) fn enqueue_generator<F, G>(&mut self, span: LightSpan, gen: G)
+    pub(super) fn enqueue_generator<F, G>(&mut self, name: &'static str, span: LightSpan, gen: G)
     where
         F: Future<Output = Result<Value, ErrorKind>> + 'static,
         G: FnOnce(GenCo) -> F,
     {
         self.frames.push(Frame::Generator {
+            name,
             span,
             state: GeneratorState::Running,
             generator: Gen::new(|co| pin_generator(gen(co))),
@@ -245,6 +247,7 @@ impl<'o> VM<'o> {
     /// or was suspended (false).
     pub(crate) fn run_generator(
         &mut self,
+        name: &'static str,
         span: LightSpan,
         frame_id: usize,
         state: GeneratorState,
@@ -266,7 +269,7 @@ impl<'o> VM<'o> {
                 // If the generator yields, it contains an instruction
                 // for what the VM should do.
                 genawaiter::GeneratorState::Yielded(request) => {
-                    self.observer.observe_generator_request(&request);
+                    self.observer.observe_generator_request(name, &request);
 
                     match request {
                         GeneratorRequest::StackPush(value) => {
@@ -282,15 +285,17 @@ impl<'o> VM<'o> {
                         // this function prepares the frame stack and yields
                         // back to the outer VM loop.
                         GeneratorRequest::ForceValue(value) => {
-                            self.reenqueue_generator(span.clone(), generator);
-                            self.enqueue_generator(span, |co| value.force(co));
+                            self.reenqueue_generator(name, span.clone(), generator);
+                            self.enqueue_generator("force", span, |co| value.force(co));
                             return Ok(false);
                         }
 
                         // Generator has requested a deep-force.
                         GeneratorRequest::DeepForceValue(value, thunk_set) => {
-                            self.reenqueue_generator(span.clone(), generator);
-                            self.enqueue_generator(span, |co| value.deep_force(co, thunk_set));
+                            self.reenqueue_generator(name, span.clone(), generator);
+                            self.enqueue_generator("deep_force", span, |co| {
+                                value.deep_force(co, thunk_set)
+                            });
                             return Ok(false);
                         }
 
@@ -298,10 +303,10 @@ impl<'o> VM<'o> {
                         // Logic is similar to `ForceValue`, except with the
                         // value being taken from that stack.
                         GeneratorRequest::WithValue(idx) => {
-                            self.reenqueue_generator(span.clone(), generator);
+                            self.reenqueue_generator(name, span.clone(), generator);
 
                             let value = self.stack[self.with_stack[idx]].clone();
-                            self.enqueue_generator(span, |co| value.force(co));
+                            self.enqueue_generator("force", span, |co| value.force(co));
 
                             return Ok(false);
                         }
@@ -310,34 +315,36 @@ impl<'o> VM<'o> {
                         // with-stack. Logic is same as above, except for the
                         // value being from that stack.
                         GeneratorRequest::CapturedWithValue(idx) => {
-                            self.reenqueue_generator(span.clone(), generator);
+                            self.reenqueue_generator(name, span.clone(), generator);
 
                             let call_frame = self.last_call_frame()
                                 .expect("Tvix bug: generator requested captured with-value, but there is no call frame");
 
                             let value = call_frame.upvalues.with_stack().unwrap()[idx].clone();
-                            self.enqueue_generator(span, |co| value.force(co));
+                            self.enqueue_generator("force", span, |co| value.force(co));
 
                             return Ok(false);
                         }
 
                         GeneratorRequest::NixEquality(values, ptr_eq) => {
                             let values = *values;
-                            self.reenqueue_generator(span.clone(), generator);
-                            self.enqueue_generator(span, |co| {
+                            self.reenqueue_generator(name, span.clone(), generator);
+                            self.enqueue_generator("nix_eq", span, |co| {
                                 values.0.nix_eq(values.1, co, ptr_eq)
                             });
                             return Ok(false);
                         }
 
                         GeneratorRequest::StringCoerce(val, kind) => {
-                            self.reenqueue_generator(span.clone(), generator);
-                            self.enqueue_generator(span, |co| val.coerce_to_string(co, kind));
+                            self.reenqueue_generator(name, span.clone(), generator);
+                            self.enqueue_generator("coerce_to_string", span, |co| {
+                                val.coerce_to_string(co, kind)
+                            });
                             return Ok(false);
                         }
 
                         GeneratorRequest::Call(callable) => {
-                            self.reenqueue_generator(span.clone(), generator);
+                            self.reenqueue_generator(name, span.clone(), generator);
                             self.tail_call_value(span, None, callable)?;
                             return Ok(false);
                         }
@@ -347,7 +354,7 @@ impl<'o> VM<'o> {
                             upvalues,
                             light_span,
                         } => {
-                            self.reenqueue_generator(span, generator);
+                            self.reenqueue_generator(name, span, generator);
 
                             self.frames.push(Frame::CallFrame {
                                 span: light_span,
@@ -423,14 +430,14 @@ impl<'o> VM<'o> {
 
                         GeneratorRequest::TryForce(value) => {
                             self.try_eval_frames.push(frame_id);
-                            self.reenqueue_generator(span.clone(), generator);
+                            self.reenqueue_generator(name, span.clone(), generator);
 
                             debug_assert!(
                                 self.frames.len() == frame_id + 1,
                                 "generator should be reenqueued with the same frame ID"
                             );
 
-                            self.enqueue_generator(span, |co| value.force(co));
+                            self.enqueue_generator("force", span, |co| value.force(co));
                             return Ok(false);
                         }
                     }
diff --git a/tvix/eval/src/vm/macros.rs b/tvix/eval/src/vm/macros.rs
index 4c027b0f6440..969b50807766 100644
--- a/tvix/eval/src/vm/macros.rs
+++ b/tvix/eval/src/vm/macros.rs
@@ -48,7 +48,7 @@ macro_rules! cmp_op {
 
         let gen_span = $frame.current_light_span();
         $vm.push_call_frame($span, $frame);
-        $vm.enqueue_generator(gen_span, |co| compare(a, b, co));
+        $vm.enqueue_generator("compare", gen_span, |co| compare(a, b, co));
         return Ok(false);
     }};
 
diff --git a/tvix/eval/src/vm/mod.rs b/tvix/eval/src/vm/mod.rs
index e104c3d2e7f9..78dc566ae729 100644
--- a/tvix/eval/src/vm/mod.rs
+++ b/tvix/eval/src/vm/mod.rs
@@ -125,6 +125,9 @@ enum Frame {
     /// be suspended while waiting for the VM to do something (e.g.
     /// thunk forcing), and resume at the same point.
     Generator {
+        /// human-readable description of the generator,
+        name: &'static str,
+
         /// Span from which the generator was launched.
         span: LightSpan,
 
@@ -296,11 +299,13 @@ impl<'o> VM<'o> {
                 // Handle generator frames, which can request thunk forcing
                 // during their execution.
                 Frame::Generator {
+                    name,
                     span,
                     state,
                     generator,
                 } => {
-                    self.observer.observe_enter_generator(frame_id, &self.stack);
+                    self.observer
+                        .observe_enter_generator(frame_id, name, &self.stack);
 
                     let initial_msg = if catchable_error_occurred {
                         catchable_error_occurred = false;
@@ -309,16 +314,24 @@ impl<'o> VM<'o> {
                         None
                     };
 
-                    match self.run_generator(span, frame_id, state, generator, initial_msg) {
-                        Ok(true) => self.observer.observe_exit_generator(frame_id, &self.stack),
-                        Ok(false) => self
-                            .observer
-                            .observe_suspend_generator(frame_id, &self.stack),
+                    match self.run_generator(name, span, frame_id, state, generator, initial_msg) {
+                        Ok(true) => {
+                            self.observer
+                                .observe_exit_generator(frame_id, name, &self.stack)
+                        }
+                        Ok(false) => {
+                            self.observer
+                                .observe_suspend_generator(frame_id, name, &self.stack)
+                        }
 
                         Err(err) => {
                             if let Some(catching_frame_idx) = self.try_eval_frames.pop() {
                                 if err.kind.is_catchable() {
-                                    self.observer.observe_exit_generator(frame_id, &self.stack);
+                                    self.observer.observe_exit_generator(
+                                        frame_id,
+                                        name,
+                                        &self.stack,
+                                    );
                                     catchable_error_occurred = true;
 
                                     // truncate the frame stack back to the
@@ -395,7 +408,7 @@ impl<'o> VM<'o> {
                     // OpAdd can add not just numbers, but also string-like
                     // things, which requires more VM logic. This operation is
                     // evaluated in a generator frame.
-                    self.enqueue_generator(gen_span, |co| add_values(co, a, b));
+                    self.enqueue_generator("add_values", gen_span, |co| add_values(co, a, b));
                     return Ok(false);
                 }
 
@@ -450,7 +463,7 @@ impl<'o> VM<'o> {
                     let a = self.stack_pop();
                     let gen_span = frame.current_light_span();
                     self.push_call_frame(span, frame);
-                    self.enqueue_generator(gen_span, |co| {
+                    self.enqueue_generator("nix_eq", gen_span, |co| {
                         a.nix_eq(b, co, PointerEquality::ForbidAll)
                     });
                     return Ok(false);
@@ -548,7 +561,7 @@ impl<'o> VM<'o> {
                     let gen_span = frame.current_light_span();
                     self.push_call_frame(span, frame);
 
-                    self.enqueue_generator(gen_span, |co| {
+                    self.enqueue_generator("coerce_to_string", gen_span, |co| {
                         value.coerce_to_string(co, CoercionKind::Weak)
                     });
 
@@ -665,7 +678,7 @@ impl<'o> VM<'o> {
                         .map(|frame| frame.upvalues.with_stack_len())
                         .unwrap_or(0);
 
-                    self.enqueue_generator(op_span, |co| {
+                    self.enqueue_generator("resolve_with", op_span, |co| {
                         resolve_with(
                             co,
                             ident.as_str().to_owned(),
@@ -752,7 +765,7 @@ impl<'o> VM<'o> {
                         let gen_span = frame.current_light_span();
 
                         self.push_call_frame(span, frame);
-                        self.enqueue_generator(gen_span, |co| thunk.force(co));
+                        self.enqueue_generator("force", gen_span, |co| thunk.force(co));
                         return Ok(false);
                     }
                 }
@@ -873,9 +886,10 @@ impl<'o> VM<'o> {
             BuiltinResult::Partial(partial) => self.stack.push(Value::Builtin(partial)),
 
             // Builtin is fully applied and the generator needs to be run by the VM.
-            BuiltinResult::Called(generator) => self.frames.push(Frame::Generator {
+            BuiltinResult::Called(name, generator) => self.frames.push(Frame::Generator {
                 generator,
                 span,
+                name,
                 state: GeneratorState::Running,
             }),
         }
@@ -922,7 +936,7 @@ impl<'o> VM<'o> {
                     .map(|p| p.current_light_span())
                     .unwrap_or_else(|| self.reasonable_light_span());
 
-                self.enqueue_generator(gen_span, |co| call_functor(co, val));
+                self.enqueue_generator("__functor call", gen_span, |co| call_functor(co, val));
                 Ok(())
             }
             v => Err(self.error(ErrorKind::NotCallable(v.type_of()))),
@@ -1104,7 +1118,7 @@ pub fn run_lambda(
 
     // Synthesise a frame that will instruct the VM to deep-force the final
     // value before returning it.
-    vm.enqueue_generator(root_span.into(), final_deep_force);
+    vm.enqueue_generator("final_deep_force", root_span.into(), final_deep_force);
 
     vm.frames.push(Frame::CallFrame {
         span: root_span.into(),