feat(tvix/eval): print lambda memory adresses in disassembler r/4768

This makes it easier to track exactly which lambda is which when
inspecting e.g. the concrete representation of a thunk.

At runtime all lambdas live in an Rc. To make this print the right
address, the construction of these Rcs had to be moved up right to the
point where the lambda is first emitted (and disassembled).

Change-Id: I6070e6c8ac55f0bd697966c4e7c5565c20d19106
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6435
Reviewed-by: sterni <sternenseemann@systemli.org>
Tested-by: BuildkiteCI

4 files changed, 30 insertions, 32 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index aacb7143d3b6..bae607d4a024 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -900,19 +900,17 @@ impl Compiler<'_> {
         // Pop the lambda context back off, and emit the finished
         // lambda as a constant.
         let compiled = self.contexts.pop().unwrap();
+        let lambda = Rc::new(compiled.lambda);
 
         #[cfg(feature = "disassembler")]
         {
-            crate::disassembler::disassemble_chunk(&compiled.lambda.chunk);
+            crate::disassembler::disassemble_lambda(lambda.clone());
         }
 
         // If the function is not a closure, just emit it directly and
         // move on.
-        if compiled.lambda.upvalue_count == 0 {
-            self.emit_constant(
-                Value::Closure(Closure::new(Rc::new(compiled.lambda))),
-                &node,
-            );
+        if lambda.upvalue_count == 0 {
+            self.emit_constant(Value::Closure(Closure::new(lambda)), &node);
             return;
         }
 
@@ -920,9 +918,7 @@ impl Compiler<'_> {
         // number of operands that allow the runtime to close over the
         // upvalues and leave a blueprint in the constant index from
         // which the runtime closure can be constructed.
-        let blueprint_idx = self
-            .chunk()
-            .push_constant(Value::Blueprint(Rc::new(compiled.lambda)));
+        let blueprint_idx = self.chunk().push_constant(Value::Blueprint(lambda));
 
         self.push_op(OpCode::OpClosure(blueprint_idx), &node);
         self.emit_upvalue_data(outer_slot, compiled.scope.upvalues);
@@ -955,23 +951,22 @@ impl Compiler<'_> {
         self.end_scope(node);
 
         let thunk = self.contexts.pop().unwrap();
+        let lambda = Rc::new(thunk.lambda);
 
         #[cfg(feature = "disassembler")]
         {
-            crate::disassembler::disassemble_chunk(&thunk.lambda.chunk);
+            crate::disassembler::disassemble_lambda(lambda.clone());
         }
 
         // Emit the thunk directly if it does not close over the
         // environment.
-        if thunk.lambda.upvalue_count == 0 {
-            self.emit_constant(Value::Thunk(Thunk::new(Rc::new(thunk.lambda))), node);
+        if lambda.upvalue_count == 0 {
+            self.emit_constant(Value::Thunk(Thunk::new(lambda)), node);
             return;
         }
 
         // Otherwise prepare for runtime construction of the thunk.
-        let blueprint_idx = self
-            .chunk()
-            .push_constant(Value::Blueprint(Rc::new(thunk.lambda)));
+        let blueprint_idx = self.chunk().push_constant(Value::Blueprint(lambda));
 
         self.push_op(OpCode::OpThunk(blueprint_idx), node);
         self.emit_upvalue_data(outer_slot, thunk.scope.upvalues);
diff --git a/tvix/eval/src/disassembler.rs b/tvix/eval/src/disassembler.rs
index 7948977e4b55..9894dea652e8 100644
--- a/tvix/eval/src/disassembler.rs
+++ b/tvix/eval/src/disassembler.rs
@@ -2,11 +2,12 @@
 //! of compiled code, as well as tracing the runtime stack during
 //! execution.
 use std::io::{Stderr, Write};
+use std::rc::Rc;
 use tabwriter::TabWriter;
 
 use crate::chunk::Chunk;
 use crate::opcode::{CodeIdx, OpCode};
-use crate::value::Value;
+use crate::value::{Lambda, Value};
 
 /// Helper struct to trace runtime values and automatically flush the
 /// output after the value is dropped (i.e. in both success and
@@ -58,20 +59,20 @@ fn disassemble_op(tw: &mut TabWriter<Stderr>, chunk: &Chunk, width: usize, offse
     };
 }
 
-/// Disassemble a chunk of code, printing out the operations in a
-/// reasonable, human-readable format.
-pub fn disassemble_chunk(chunk: &Chunk) {
+/// Disassemble an entire lambda, printing its address and its
+/// operations in human-readable format.
+pub fn disassemble_lambda(lambda: Rc<Lambda>) {
     let mut tw = TabWriter::new(std::io::stderr());
-
     let _ = writeln!(
         &mut tw,
-        "=== compiled bytecode ({} operations) ===",
-        chunk.code.len()
+        "=== compiled code (@{:p}, {} ops) ===",
+        lambda,
+        lambda.chunk.code.len()
     );
 
-    let width = format!("{}", chunk.code.len()).len();
-    for (idx, _) in chunk.code.iter().enumerate() {
-        disassemble_op(&mut tw, chunk, width, idx);
+    let width = format!("{}", lambda.chunk.code.len()).len();
+    for (idx, _) in lambda.chunk.code.iter().enumerate() {
+        disassemble_op(&mut tw, &lambda.chunk, width, idx);
     }
 
     let _ = tw.flush();
diff --git a/tvix/eval/src/eval.rs b/tvix/eval/src/eval.rs
index ecf0824871a0..c2fe7f151377 100644
--- a/tvix/eval/src/eval.rs
+++ b/tvix/eval/src/eval.rs
@@ -1,4 +1,4 @@
-use std::path::PathBuf;
+use std::{path::PathBuf, rc::Rc};
 
 use crate::{
     builtins::global_builtins,
@@ -45,11 +45,12 @@ pub fn interpret(code: &str, location: Option<PathBuf>) -> EvalResult<Value> {
         &file,
         global_builtins(),
         #[cfg(feature = "disassembler")]
-        std::rc::Rc::new(codemap),
+        Rc::new(codemap),
     )?;
+    let lambda = Rc::new(result.lambda);
 
     #[cfg(feature = "disassembler")]
-    crate::disassembler::disassemble_chunk(&result.lambda.chunk);
+    crate::disassembler::disassemble_lambda(lambda.clone());
 
     for warning in result.warnings {
         eprintln!(
@@ -73,5 +74,5 @@ pub fn interpret(code: &str, location: Option<PathBuf>) -> EvalResult<Value> {
         return Err(err.clone());
     }
 
-    crate::vm::run_lambda(result.lambda)
+    crate::vm::run_lambda(lambda)
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 5d26afe3f304..2c490200e1d1 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -173,8 +173,9 @@ impl VM {
     ) -> EvalResult<Value> {
         #[cfg(feature = "disassembler")]
         self.tracer.literal(&format!(
-            "=== entering closure/{} [{}] ===",
+            "=== entering closure/{} @ {:p} [{}] ===",
             arg_count,
+            lambda,
             self.frames.len()
         ));
 
@@ -706,9 +707,9 @@ fn unwrap_or_clone_rc<T: Clone>(rc: Rc<T>) -> T {
     Rc::try_unwrap(rc).unwrap_or_else(|rc| (*rc).clone())
 }
 
-pub fn run_lambda(lambda: Lambda) -> EvalResult<Value> {
+pub fn run_lambda(lambda: Rc<Lambda>) -> EvalResult<Value> {
     let mut vm = VM::default();
-    let value = vm.call(Rc::new(lambda), vec![], 0)?;
+    let value = vm.call(lambda, vec![], 0)?;
     vm.force_for_output(&value)?;
     Ok(value)
 }