From cc526a2c873524faa83cad62bde2edda59ea7820 Mon Sep 17 00:00:00 2001
From: Vincent Ambo <mail@tazj.in>
Date: Fri, 2 Sep 2022 17:44:31 +0300
Subject: [PATCH] feat(tvix/eval): thread codemap through to disassembler

If the disassembler feature is enabled, make sure that an Rc of the
codemap is available through the chunk.

Change-Id: I700f27ab665a704f73457b19bd2d7efc93828a16
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6414
Reviewed-by: sterni <sternenseemann@systemli.org>
Tested-by: BuildkiteCI
---
 tvix/eval/src/chunk.rs        |  3 +++
 tvix/eval/src/compiler/mod.rs | 27 +++++++++++++++++++++++++--
 tvix/eval/src/disassembler.rs | 13 +++++++++++--
 tvix/eval/src/eval.rs         |  9 ++++++++-
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/tvix/eval/src/chunk.rs b/tvix/eval/src/chunk.rs
index f3df51f31..defa409f1 100644
--- a/tvix/eval/src/chunk.rs
+++ b/tvix/eval/src/chunk.rs
@@ -26,6 +26,9 @@ 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 Chunk {
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 2d4995640..3ea27f843 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -55,10 +55,18 @@ impl LambdaCtx {
     }
 
     fn inherit(&self) -> Self {
-        LambdaCtx {
+        let ctx = LambdaCtx {
             lambda: Lambda::new_anonymous(),
             scope: self.scope.inherit(),
-        }
+        };
+
+        #[cfg(feature = "disassembler")]
+        let ctx = (|mut c: Self| {
+            c.lambda.chunk.codemap = self.lambda.chunk.codemap.clone();
+            c
+        })(ctx);
+
+        ctx
     }
 }
 
@@ -84,6 +92,12 @@ struct Compiler<'code> {
     /// and is used to track the spans from which instructions where
     /// 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>,
 }
 
 // Helper functions for emitting code and metadata to the internal
@@ -1331,6 +1345,8 @@ pub fn compile<'code>(
     location: Option<PathBuf>,
     file: &'code codemap::File,
     globals: HashMap<&'static str, Value>,
+
+    #[cfg(feature = "disassembler")] codemap: Rc<codemap::CodeMap>,
 ) -> EvalResult<CompilationOutput> {
     let mut root_dir = match location {
         Some(dir) => Ok(dir),
@@ -1353,12 +1369,19 @@ pub fn compile<'code>(
     let mut c = Compiler {
         root_dir,
         file,
+        #[cfg(feature = "disassembler")]
+        codemap,
         globals: prepare_globals(globals),
         contexts: vec![LambdaCtx::new()],
         warnings: vec![],
         errors: vec![],
     };
 
+    #[cfg(feature = "disassembler")]
+    {
+        c.context_mut().lambda.chunk.codemap = c.codemap.clone();
+    }
+
     c.compile(None, expr.clone());
 
     // The final operation of any top-level Nix program must always be
diff --git a/tvix/eval/src/disassembler.rs b/tvix/eval/src/disassembler.rs
index bc5a64fe6..e3cc1f16e 100644
--- a/tvix/eval/src/disassembler.rs
+++ b/tvix/eval/src/disassembler.rs
@@ -5,7 +5,7 @@ use std::io::{Stderr, Write};
 use tabwriter::TabWriter;
 
 use crate::chunk::Chunk;
-use crate::opcode::OpCode;
+use crate::opcode::{CodeIdx, OpCode};
 use crate::value::Value;
 
 /// Helper struct to trace runtime values and automatically flush the
@@ -38,7 +38,16 @@ impl Drop for Tracer {
 }
 
 fn disassemble_op(tw: &mut TabWriter<Stderr>, chunk: &Chunk, width: usize, offset: usize) {
-    write!(tw, "{:0width$}\t ", width = width).ok();
+    write!(tw, "{:0width$}\t ", offset, width = width).ok();
+
+    let span = chunk.get_span(CodeIdx(offset));
+
+    if offset > 0 && chunk.get_span(CodeIdx(offset - 1)) == span {
+        write!(tw, "   |\t").unwrap();
+    } else {
+        let loc = chunk.codemap.look_up_span(span);
+        write!(tw, "{:4}\t", loc.begin.line + 1).unwrap();
+    }
 
     match chunk.code[offset] {
         OpCode::OpConstant(idx) => write!(tw, "OpConstant({})\n", chunk.constant(idx)).ok(),
diff --git a/tvix/eval/src/eval.rs b/tvix/eval/src/eval.rs
index 1362c0394..296887451 100644
--- a/tvix/eval/src/eval.rs
+++ b/tvix/eval/src/eval.rs
@@ -39,7 +39,14 @@ pub fn interpret(code: &str, location: Option<PathBuf>) -> EvalResult<Value> {
         println!("{:?}", root_expr);
     }
 
-    let result = crate::compiler::compile(root_expr, location, &file, global_builtins())?;
+    let result = crate::compiler::compile(
+        root_expr,
+        location,
+        &file,
+        global_builtins(),
+        #[cfg(feature = "disassembler")]
+        std::rc::Rc::new(codemap),
+    )?;
 
     #[cfg(feature = "disassembler")]
     crate::disassembler::disassemble_chunk(&result.lambda.chunk);