feat(tvix/eval): implement DisassemblingObserver for compiler r/4775

This type implements an observer that is called whenever the compiler
emits a chunk (after the toplevel, thunks, or lambdas) and prints the
output of the disassembler to its internal writer.

This replaces half of the uses of the `disassembler` feature, which
has been removed from the Cargo configuration.

Note that at this commit runtime tracing is not yet implemented as an
observer.

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

7 files changed, 111 insertions, 98 deletions

diff --git a/tvix/eval/Cargo.toml b/tvix/eval/Cargo.toml
index 226f3a98a3..ac01b595be 100644
--- a/tvix/eval/Cargo.toml
+++ b/tvix/eval/Cargo.toml
@@ -17,7 +17,7 @@ smol_str = "0.1"
 rustyline = { version = "10.0.0", optional = true }
 dirs = "4.0.0"
 path-clean = "0.1"
-tabwriter = { version = "1.2", optional = true }
+tabwriter = "1.2"
 rowan = "*" # pinned by rnix
 codemap = "0.1.3"
 
@@ -43,9 +43,6 @@ nix_tests = []
 # Enables building the binary (tvix-eval REPL)
 repl = [ "dep:rustyline" ]
 
-# Enables printing compiled code and tracing the stack state at runtime.
-disassembler = ["dep:tabwriter"]
-
 [[bench]]
 name = "eval"
 harness = false
diff --git a/tvix/eval/src/chunk.rs b/tvix/eval/src/chunk.rs
index a4d6be752b..4d653c2b22 100644
--- a/tvix/eval/src/chunk.rs
+++ b/tvix/eval/src/chunk.rs
@@ -28,9 +28,6 @@ pub struct Chunk {
     pub code: Vec<OpCode>,
     pub constants: Vec<Value>,
     spans: Vec<SourceSpan>,
-
-    #[cfg(feature = "disassembler")]
-    pub codemap: std::rc::Rc<codemap::CodeMap>,
 }
 
 impl Index<ConstantIdx> for Chunk {
@@ -93,11 +90,11 @@ impl Chunk {
     }
 
     /// Retrieve the line from which the instruction at `offset` was
-    /// compiled. Only available when the chunk carries a codemap,
-    /// i.e. when the disassembler is enabled.
-    #[cfg(feature = "disassembler")]
-    pub fn get_line(&self, offset: CodeIdx) -> usize {
+    /// compiled in the specified codemap.
+    pub fn get_line(&self, codemap: &codemap::CodeMap, offset: CodeIdx) -> usize {
         let span = self.get_span(offset);
-        self.codemap.look_up_span(span).begin.line + 1
+        // lines are 0-indexed in the codemap, but users probably want
+        // real line numbers
+        codemap.look_up_span(span).begin.line + 1
     }
 }
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index b588522994..448999fb38 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -25,6 +25,7 @@ use std::rc::Rc;
 
 use crate::chunk::Chunk;
 use crate::errors::{Error, ErrorKind, EvalResult};
+use crate::observer::Observer;
 use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, UpvalueIdx};
 use crate::value::{Closure, Lambda, Thunk, Value};
 use crate::warnings::{EvalWarning, WarningKind};
@@ -35,7 +36,7 @@ use self::scope::{LocalIdx, LocalPosition, Scope, Upvalue, UpvalueKind};
 /// compilation was successful, the resulting bytecode can be passed
 /// to the VM.
 pub struct CompilationOutput {
-    pub lambda: Lambda,
+    pub lambda: Rc<Lambda>,
     pub warnings: Vec<EvalWarning>,
     pub errors: Vec<Error>,
 }
@@ -54,21 +55,11 @@ impl LambdaCtx {
         }
     }
 
-    #[allow(clippy::let_and_return)] // due to disassembler
     fn inherit(&self) -> Self {
-        let ctx = LambdaCtx {
+        LambdaCtx {
             lambda: Lambda::new_anonymous(),
             scope: self.scope.inherit(),
-        };
-
-        #[cfg(feature = "disassembler")]
-        #[allow(clippy::redundant_closure_call)]
-        let ctx = (|mut c: Self| {
-            c.lambda.chunk.codemap = self.lambda.chunk.codemap.clone();
-            c
-        })(ctx);
-
-        ctx
+        }
     }
 }
 
@@ -76,7 +67,7 @@ impl LambdaCtx {
 /// implicitly be resolvable in the global scope.
 type GlobalsMap = HashMap<&'static str, Rc<dyn Fn(&mut Compiler, rnix::ast::Ident)>>;
 
-struct Compiler<'code> {
+struct Compiler<'code, 'observer> {
     contexts: Vec<LambdaCtx>,
     warnings: Vec<EvalWarning>,
     errors: Vec<Error>,
@@ -95,16 +86,14 @@ struct Compiler<'code> {
     /// derived.
     file: &'code codemap::File,
 
-    #[cfg(feature = "disassembler")]
-    /// Carry a reference to the codemap around when the disassembler
-    /// is enabled, to allow displaying lines and other source
-    /// information in the disassembler output.
-    codemap: Rc<codemap::CodeMap>,
+    /// Carry an observer for the compilation process, which is called
+    /// whenever a chunk is emitted.
+    observer: &'observer mut dyn Observer,
 }
 
 // Helper functions for emitting code and metadata to the internal
 // structures of the compiler.
-impl Compiler<'_> {
+impl Compiler<'_, '_> {
     fn context(&self) -> &LambdaCtx {
         &self.contexts[self.contexts.len() - 1]
     }
@@ -150,7 +139,7 @@ impl Compiler<'_> {
 }
 
 // Actual code-emitting AST traversal methods.
-impl Compiler<'_> {
+impl Compiler<'_, '_> {
     fn compile(&mut self, slot: LocalIdx, expr: ast::Expr) {
         match expr {
             ast::Expr::Literal(literal) => self.compile_literal(literal),
@@ -913,11 +902,7 @@ impl Compiler<'_> {
         // lambda as a constant.
         let compiled = self.contexts.pop().unwrap();
         let lambda = Rc::new(compiled.lambda);
-
-        #[cfg(feature = "disassembler")]
-        {
-            crate::disassembler::disassemble_lambda(lambda.clone());
-        }
+        self.observer.observe_compiled_lambda(&lambda);
 
         // If the function is not a closure, just emit it directly and
         // move on.
@@ -964,11 +949,7 @@ impl Compiler<'_> {
 
         let thunk = self.contexts.pop().unwrap();
         let lambda = Rc::new(thunk.lambda);
-
-        #[cfg(feature = "disassembler")]
-        {
-            crate::disassembler::disassemble_lambda(lambda.clone());
-        }
+        self.observer.observe_compiled_thunk(&lambda);
 
         // Emit the thunk directly if it does not close over the
         // environment.
@@ -1369,8 +1350,7 @@ pub fn compile(
     location: Option<PathBuf>,
     file: &codemap::File,
     globals: HashMap<&'static str, Value>,
-
-    #[cfg(feature = "disassembler")] codemap: Rc<codemap::CodeMap>,
+    observer: &mut dyn Observer,
 ) -> EvalResult<CompilationOutput> {
     let mut root_dir = match location {
         Some(dir) => Ok(dir),
@@ -1393,19 +1373,13 @@ pub fn compile(
     let mut c = Compiler {
         root_dir,
         file,
-        #[cfg(feature = "disassembler")]
-        codemap,
+        observer,
         globals: prepare_globals(globals),
         contexts: vec![LambdaCtx::new()],
         warnings: vec![],
         errors: vec![],
     };
 
-    #[cfg(feature = "disassembler")]
-    {
-        c.context_mut().lambda.chunk.codemap = c.codemap.clone();
-    }
-
     let root_span = c.span_for(&expr);
     let root_slot = c.scope_mut().declare_phantom(root_span);
     c.compile(root_slot, expr.clone());
@@ -1416,8 +1390,11 @@ pub fn compile(
     // thunk might be returned).
     c.emit_force(&expr);
 
+    let lambda = Rc::new(c.contexts.pop().unwrap().lambda);
+    c.observer.observe_compiled_toplevel(&lambda);
+
     Ok(CompilationOutput {
-        lambda: c.contexts.pop().unwrap().lambda,
+        lambda,
         warnings: c.warnings,
         errors: c.errors,
     })
diff --git a/tvix/eval/src/disassembler.rs b/tvix/eval/src/disassembler.rs
index 9894dea652..d7c9c8895c 100644
--- a/tvix/eval/src/disassembler.rs
+++ b/tvix/eval/src/disassembler.rs
@@ -1,13 +1,13 @@
 //! Implements methods for disassembling and printing a representation
 //! of compiled code, as well as tracing the runtime stack during
 //! execution.
+use codemap::CodeMap;
 use std::io::{Stderr, Write};
-use std::rc::Rc;
 use tabwriter::TabWriter;
 
 use crate::chunk::Chunk;
 use crate::opcode::{CodeIdx, OpCode};
-use crate::value::{Lambda, Value};
+use crate::value::Value;
 
 /// Helper struct to trace runtime values and automatically flush the
 /// output after the value is dropped (i.e. in both success and
@@ -42,38 +42,26 @@ impl Drop for Tracer {
     }
 }
 
-fn disassemble_op(tw: &mut TabWriter<Stderr>, chunk: &Chunk, width: usize, offset: usize) {
-    let _ = write!(tw, "{:0width$}\t ", offset, width = width);
+pub fn disassemble_op<W: Write>(
+    tw: &mut W,
+    codemap: &CodeMap,
+    chunk: &Chunk,
+    width: usize,
+    idx: CodeIdx,
+) {
+    let _ = write!(tw, "{:#width$x}\t ", idx.0, width = width);
 
-    let line = chunk.get_line(CodeIdx(offset));
-
-    if offset > 0 && chunk.get_line(CodeIdx(offset - 1)) == line {
+    // Print continuation character if the previous operation was at
+    // the same line, otherwise print the line.
+    let line = chunk.get_line(codemap, idx);
+    if idx.0 > 0 && chunk.get_line(codemap, CodeIdx(idx.0 - 1)) == line {
         write!(tw, "   |\t").unwrap();
     } else {
         write!(tw, "{:4}\t", line).unwrap();
     }
 
-    let _ = match chunk.code[offset] {
-        OpCode::OpConstant(idx) => writeln!(tw, "OpConstant({})", chunk.constant(idx)),
+    let _ = match chunk[idx] {
+        OpCode::OpConstant(idx) => writeln!(tw, "OpConstant({}@{})", chunk[idx], idx.0),
         op => writeln!(tw, "{:?}", op),
     };
 }
-
-/// 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 code (@{:p}, {} ops) ===",
-        lambda,
-        lambda.chunk.code.len()
-    );
-
-    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 c2fe7f1513..f510cb6a5e 100644
--- a/tvix/eval/src/eval.rs
+++ b/tvix/eval/src/eval.rs
@@ -3,6 +3,7 @@ use std::{path::PathBuf, rc::Rc};
 use crate::{
     builtins::global_builtins,
     errors::{Error, ErrorKind, EvalResult},
+    observer::DisassemblingObserver,
     value::Value,
 };
 
@@ -15,6 +16,7 @@ pub fn interpret(code: &str, location: Option<PathBuf>) -> EvalResult<Value> {
             .unwrap_or_else(|| "<repl>".into()),
         code.into(),
     );
+    let codemap = Rc::new(codemap);
 
     let parsed = rnix::ast::Root::parse(code);
     let errors = parsed.errors();
@@ -39,18 +41,10 @@ pub fn interpret(code: &str, location: Option<PathBuf>) -> EvalResult<Value> {
         println!("{:?}", root_expr);
     }
 
-    let result = crate::compiler::compile(
-        root_expr,
-        location,
-        &file,
-        global_builtins(),
-        #[cfg(feature = "disassembler")]
-        Rc::new(codemap),
-    )?;
-    let lambda = Rc::new(result.lambda);
+    let mut observer = DisassemblingObserver::new(codemap.clone(), std::io::stderr());
 
-    #[cfg(feature = "disassembler")]
-    crate::disassembler::disassemble_lambda(lambda.clone());
+    let result =
+        crate::compiler::compile(root_expr, location, &file, global_builtins(), &mut observer)?;
 
     for warning in result.warnings {
         eprintln!(
@@ -74,5 +68,5 @@ pub fn interpret(code: &str, location: Option<PathBuf>) -> EvalResult<Value> {
         return Err(err.clone());
     }
 
-    crate::vm::run_lambda(lambda)
+    crate::vm::run_lambda(result.lambda)
 }
diff --git a/tvix/eval/src/lib.rs b/tvix/eval/src/lib.rs
index e088ef3ee4..e1a5ceaed8 100644
--- a/tvix/eval/src/lib.rs
+++ b/tvix/eval/src/lib.rs
@@ -1,6 +1,7 @@
 mod builtins;
 mod chunk;
 mod compiler;
+mod disassembler;
 mod errors;
 mod eval;
 mod observer;
@@ -10,9 +11,6 @@ mod value;
 mod vm;
 mod warnings;
 
-#[cfg(feature = "disassembler")]
-mod disassembler;
-
 #[cfg(test)]
 mod tests;
 
diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index bf05295ac3..5aeb344ee9 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -4,9 +4,15 @@
 //! This can be used to gain insights from compilation, to trace the
 //! runtime, and so on.
 
-use crate::value::Lambda;
-
+use codemap::CodeMap;
+use std::io::Write;
 use std::rc::Rc;
+use tabwriter::TabWriter;
+
+use crate::chunk::Chunk;
+use crate::disassembler::disassemble_op;
+use crate::opcode::CodeIdx;
+use crate::value::Lambda;
 
 /// Implemented by types that wish to observe internal happenings of
 /// Tvix.
@@ -34,3 +40,59 @@ pub trait Observer {
 pub struct NoOpObserver {}
 
 impl Observer for NoOpObserver {}
+
+/// An observer that prints disassembled chunk information to its
+/// internal writer whenwever the compiler emits a toplevel function,
+/// closure or thunk.
+pub struct DisassemblingObserver<W: Write> {
+    codemap: Rc<CodeMap>,
+    writer: TabWriter<W>,
+}
+
+impl<W: Write> DisassemblingObserver<W> {
+    pub fn new(codemap: Rc<CodeMap>, writer: W) -> Self {
+        Self {
+            codemap,
+            writer: TabWriter::new(writer),
+        }
+    }
+
+    fn lambda_header(&mut self, kind: &str, lambda: &Rc<Lambda>) {
+        let _ = writeln!(
+            &mut self.writer,
+            "=== compiled {} @ {:p} ({} ops) ===",
+            kind,
+            lambda,
+            lambda.chunk.code.len()
+        );
+    }
+
+    fn disassemble_chunk(&mut self, chunk: &Chunk) {
+        // 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() {
+            disassemble_op(&mut self.writer, &self.codemap, chunk, width, CodeIdx(idx));
+        }
+    }
+}
+
+impl<W: Write> Observer for DisassemblingObserver<W> {
+    fn observe_compiled_toplevel(&mut self, lambda: &Rc<Lambda>) {
+        self.lambda_header("toplevel", lambda);
+        self.disassemble_chunk(&lambda.chunk);
+        let _ = self.writer.flush();
+    }
+
+    fn observe_compiled_lambda(&mut self, lambda: &Rc<Lambda>) {
+        self.lambda_header("lambda", lambda);
+        self.disassemble_chunk(&lambda.chunk);
+        let _ = self.writer.flush();
+    }
+
+    fn observe_compiled_thunk(&mut self, lambda: &Rc<Lambda>) {
+        self.lambda_header("thunk", lambda);
+        self.disassemble_chunk(&lambda.chunk);
+        let _ = self.writer.flush();
+    }
+}