refactor(tvix/eval): introduce source::SourceCode type r/5035

This type hides away the lower-level handling of most codemap data
structures, especially to library consumers (see corresponding changes
in tvixbolt).

This will help with implement `import` by giving us central control
over how the codemap works.

Change-Id: Ifcea36776879725871b30c518aeb96ab5fda035a
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6855
Tested-by: BuildkiteCI
Reviewed-by: wpcarro <wpcarro@gmail.com>

7 files changed, 97 insertions, 48 deletions

diff --git a/tvix/eval/src/chunk.rs b/tvix/eval/src/chunk.rs
index 8810f4e1c4..9958b551a1 100644
--- a/tvix/eval/src/chunk.rs
+++ b/tvix/eval/src/chunk.rs
@@ -1,10 +1,9 @@
 use std::io::Write;
 use std::ops::Index;
 
-use codemap::CodeMap;
-
 use crate::opcode::{CodeIdx, ConstantIdx, OpCode};
 use crate::value::Value;
+use crate::SourceCode;
 
 /// Represents a source location from which one or more operations
 /// were compiled.
@@ -117,21 +116,12 @@ impl Chunk {
         panic!("compiler error: chunk missing span for offset {}", offset.0);
     }
 
-    /// Retrieve the line from which the instruction at `offset` was
-    /// compiled in the specified codemap.
-    pub fn get_line(&self, codemap: &codemap::CodeMap, offset: CodeIdx) -> usize {
-        let span = self.get_span(offset);
-        // lines are 0-indexed in the codemap, but users probably want
-        // real line numbers
-        codemap.look_up_span(span).begin.line + 1
-    }
-
     /// Write the disassembler representation of the operation at
     /// `idx` to the specified writer.
     pub fn disassemble_op<W: Write>(
         &self,
         writer: &mut W,
-        codemap: &CodeMap,
+        source: &SourceCode,
         width: usize,
         idx: CodeIdx,
     ) -> Result<(), std::io::Error> {
@@ -139,8 +129,8 @@ impl Chunk {
 
         // Print continuation character if the previous operation was at
         // the same line, otherwise print the line.
-        let line = self.get_line(codemap, idx);
-        if idx.0 > 0 && self.get_line(codemap, CodeIdx(idx.0 - 1)) == line {
+        let line = source.get_line(self.get_span(idx));
+        if idx.0 > 0 && source.get_line(self.get_span(CodeIdx(idx.0 - 1))) == line {
             write!(writer, "   |\t")?;
         } else {
             write!(writer, "{:4}\t", line)?;
diff --git a/tvix/eval/src/errors.rs b/tvix/eval/src/errors.rs
index 8ecad80535..75fac9c926 100644
--- a/tvix/eval/src/errors.rs
+++ b/tvix/eval/src/errors.rs
@@ -2,11 +2,11 @@ use crate::value::CoercionKind;
 use std::path::PathBuf;
 use std::{fmt::Display, num::ParseIntError};
 
-use codemap::{CodeMap, Span};
+use codemap::Span;
 use codemap_diagnostic::{ColorConfig, Diagnostic, Emitter, Level, SpanLabel, SpanStyle};
 use smol_str::SmolStr;
 
-use crate::Value;
+use crate::{SourceCode, Value};
 
 #[derive(Clone, Debug)]
 pub enum ErrorKind {
@@ -135,16 +135,16 @@ impl Display for Error {
 pub type EvalResult<T> = Result<T, Error>;
 
 impl Error {
-    pub fn fancy_format_str(&self, codemap: &CodeMap) -> String {
+    pub fn fancy_format_str(&self, source: &SourceCode) -> String {
         let mut out = vec![];
-        Emitter::vec(&mut out, Some(codemap)).emit(&[self.diagnostic(codemap)]);
+        Emitter::vec(&mut out, Some(&*source.codemap())).emit(&[self.diagnostic()]);
         String::from_utf8_lossy(&out).to_string()
     }
 
     /// Render a fancy, human-readable output of this error and print
     /// it to stderr.
-    pub fn fancy_format_stderr(&self, codemap: &CodeMap) {
-        Emitter::stderr(ColorConfig::Auto, Some(codemap)).emit(&[self.diagnostic(codemap)]);
+    pub fn fancy_format_stderr(&self, source: &SourceCode) {
+        Emitter::stderr(ColorConfig::Auto, Some(&*source.codemap())).emit(&[self.diagnostic()]);
     }
 
     /// Create the optional span label displayed as an annotation on
@@ -154,7 +154,7 @@ impl Error {
     }
 
     /// Create the primary error message displayed to users.
-    fn message(&self, codemap: &CodeMap) -> String {
+    fn message(&self) -> String {
         match &self.kind {
             ErrorKind::Throw(msg) => format!("error thrown: {}", msg),
             ErrorKind::Abort(msg) => format!("evaluation aborted: {}", msg),
@@ -224,7 +224,7 @@ to a missing value in the attribute set(s) included via `with`."#,
             // TODO(tazjin): trace through the whole chain of thunk
             // forcing errors with secondary spans, instead of just
             // delegating to the inner error
-            ErrorKind::ThunkForce(err) => err.message(codemap),
+            ErrorKind::ThunkForce(err) => err.message(),
 
             ErrorKind::NotCoercibleToString { kind, from } => {
                 let kindly = match kind {
@@ -316,7 +316,7 @@ to a missing value in the attribute set(s) included via `with`."#,
         }
     }
 
-    fn diagnostic(&self, codemap: &CodeMap) -> Diagnostic {
+    fn diagnostic(&self) -> Diagnostic {
         let span_label = SpanLabel {
             label: self.span_label(),
             span: self.span,
@@ -325,7 +325,7 @@ to a missing value in the attribute set(s) included via `with`."#,
 
         Diagnostic {
             level: Level::Error,
-            message: self.message(codemap),
+            message: self.message(),
             spans: vec![span_label],
             code: Some(self.code().into()),
         }
diff --git a/tvix/eval/src/eval.rs b/tvix/eval/src/eval.rs
index bc430e5804..7b7d3983d4 100644
--- a/tvix/eval/src/eval.rs
+++ b/tvix/eval/src/eval.rs
@@ -1,10 +1,11 @@
-use std::{path::PathBuf, rc::Rc};
+use std::path::PathBuf;
 
 use crate::{
     builtins::global_builtins,
     errors::{Error, ErrorKind, EvalResult},
     observer::{DisassemblingObserver, NoOpObserver, TracingObserver},
     value::Value,
+    SourceCode,
 };
 
 /// Runtime options for the Tvix interpreter
@@ -25,15 +26,14 @@ pub struct Options {
 }
 
 pub fn interpret(code: &str, location: Option<PathBuf>, options: Options) -> EvalResult<Value> {
-    let mut codemap = codemap::CodeMap::new();
-    let file = codemap.add_file(
+    let source = SourceCode::new();
+    let file = source.add_file(
         location
             .as_ref()
             .map(|p| p.to_string_lossy().to_string())
             .unwrap_or_else(|| "[tvix-repl]".into()),
         code.into(),
     );
-    let codemap = Rc::new(codemap);
 
     let parsed = rnix::ast::Root::parse(code);
     let errors = parsed.errors();
@@ -64,7 +64,7 @@ pub fn interpret(code: &str, location: Option<PathBuf>, options: Options) -> Eva
             location,
             file.clone(),
             global_builtins(),
-            &mut DisassemblingObserver::new(codemap.clone(), std::io::stderr()),
+            &mut DisassemblingObserver::new(source.clone(), std::io::stderr()),
         )
     } else {
         crate::compiler::compile(
@@ -77,11 +77,11 @@ pub fn interpret(code: &str, location: Option<PathBuf>, options: Options) -> Eva
     }?;
 
     for warning in result.warnings {
-        warning.fancy_format_stderr(&codemap);
+        warning.fancy_format_stderr(&source);
     }
 
     for error in &result.errors {
-        error.fancy_format_stderr(&codemap);
+        error.fancy_format_stderr(&source);
     }
 
     if let Some(err) = result.errors.last() {
@@ -95,7 +95,7 @@ pub fn interpret(code: &str, location: Option<PathBuf>, options: Options) -> Eva
     };
 
     if let Err(err) = &result {
-        err.fancy_format_stderr(&codemap);
+        err.fancy_format_stderr(&source);
     }
 
     result
diff --git a/tvix/eval/src/lib.rs b/tvix/eval/src/lib.rs
index d7255af14e..847447cc84 100644
--- a/tvix/eval/src/lib.rs
+++ b/tvix/eval/src/lib.rs
@@ -5,6 +5,7 @@ mod errors;
 mod eval;
 pub mod observer;
 mod opcode;
+mod source;
 mod upvalues;
 mod value;
 mod vm;
@@ -22,5 +23,6 @@ pub use crate::builtins::global_builtins;
 pub use crate::compiler::compile;
 pub use crate::errors::EvalResult;
 pub use crate::eval::{interpret, Options};
+pub use crate::source::SourceCode;
 pub use crate::value::Value;
 pub use crate::vm::run_lambda;
diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index ce237529f5..8572c2c1eb 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -6,7 +6,6 @@
 //!
 //! All methods are optional, that is, observers can implement only
 /// what they are interested in observing.
-use codemap::CodeMap;
 use std::io::Write;
 use std::rc::Rc;
 use tabwriter::TabWriter;
@@ -14,6 +13,7 @@ use tabwriter::TabWriter;
 use crate::chunk::Chunk;
 use crate::opcode::{CodeIdx, OpCode};
 use crate::value::Lambda;
+use crate::SourceCode;
 use crate::Value;
 
 /// Implemented by types that wish to observe internal happenings of
@@ -69,14 +69,14 @@ impl RuntimeObserver for NoOpObserver {}
 /// internal writer whenwever the compiler emits a toplevel function,
 /// closure or thunk.
 pub struct DisassemblingObserver<W: Write> {
-    codemap: Rc<CodeMap>,
+    source: SourceCode,
     writer: TabWriter<W>,
 }
 
 impl<W: Write> DisassemblingObserver<W> {
-    pub fn new(codemap: Rc<CodeMap>, writer: W) -> Self {
+    pub fn new(source: SourceCode, writer: W) -> Self {
         Self {
-            codemap,
+            source,
             writer: TabWriter::new(writer),
         }
     }
@@ -96,7 +96,7 @@ impl<W: Write> DisassemblingObserver<W> {
         let width = format!("{:#x}", chunk.code.len() - 1).len();
 
         for (idx, _) in chunk.code.iter().enumerate() {
-            let _ = chunk.disassemble_op(&mut self.writer, &self.codemap, width, CodeIdx(idx));
+            let _ = chunk.disassemble_op(&mut self.writer, &self.source, width, CodeIdx(idx));
         }
     }
 }
diff --git a/tvix/eval/src/source.rs b/tvix/eval/src/source.rs
new file mode 100644
index 0000000000..0104495282
--- /dev/null
+++ b/tvix/eval/src/source.rs
@@ -0,0 +1,57 @@
+//! This module contains utilities for dealing with the codemap that
+//! needs to be carried across different compiler instantiations in an
+//! evaluation.
+//!
+//! The data type `SourceCode` should be carried through all relevant
+//! places instead of copying the codemap structures directly.
+
+use std::{
+    cell::{Ref, RefCell, RefMut},
+    rc::Rc,
+    sync::Arc,
+};
+
+use codemap::{CodeMap, Span};
+
+/// Tracks all source code in a Tvix evaluation for accurate error
+/// reporting.
+#[derive(Clone)]
+pub struct SourceCode(Rc<RefCell<CodeMap>>);
+
+impl SourceCode {
+    /// Create a new SourceCode instance.
+    pub fn new() -> Self {
+        SourceCode(Rc::new(RefCell::new(CodeMap::new())))
+    }
+
+    /// Access a read-only reference to the codemap.
+    pub fn codemap(&self) -> Ref<CodeMap> {
+        self.0.borrow()
+    }
+
+    /// Access a writable reference to the codemap.
+    fn codemap_mut(&self) -> RefMut<CodeMap> {
+        self.0.borrow_mut()
+    }
+
+    /// Add a file to the codemap. The returned Arc is managed by the
+    /// codemap internally and can be used like a normal reference.
+    pub fn add_file(&self, name: String, code: String) -> Arc<codemap::File> {
+        self.codemap_mut().add_file(name, code)
+    }
+
+    /// Retrieve the line number of the given span. If it spans
+    /// multiple lines, the first line will be returned.
+    pub fn get_line(&self, span: Span) -> usize {
+        // lines are 0-indexed in the codemap, but users probably want
+        // real line numbers
+        self.codemap().look_up_span(span).begin.line + 1
+    }
+
+    /// Returns the literal source slice of the given span.
+    pub fn source_slice(&self, span: Span) -> Ref<str> {
+        Ref::map(self.codemap(), |c| {
+            c.find_file(span.low()).source_slice(span)
+        })
+    }
+}
diff --git a/tvix/eval/src/warnings.rs b/tvix/eval/src/warnings.rs
index ddb7510c8a..63574b5389 100644
--- a/tvix/eval/src/warnings.rs
+++ b/tvix/eval/src/warnings.rs
@@ -1,9 +1,10 @@
 //! Implements warnings that are emitted in cases where code passed to
 //! Tvix exhibits problems that the user could address.
 
-use codemap::CodeMap;
 use codemap_diagnostic::{ColorConfig, Diagnostic, Emitter, Level, SpanLabel, SpanStyle};
 
+use crate::SourceCode;
+
 #[derive(Debug)]
 pub enum WarningKind {
     DeprecatedLiteralURL,
@@ -27,17 +28,18 @@ impl EvalWarning {
     /// Render a fancy, human-readable output of this warning and
     /// return it as a String. Note that this version of the output
     /// does not include any colours or font styles.
-    pub fn fancy_format_str(&self, codemap: &CodeMap) -> String {
+    pub fn fancy_format_str(&self, source: &SourceCode) -> String {
         let mut out = vec![];
-        Emitter::vec(&mut out, Some(codemap)).emit(&[self.diagnostic(codemap)]);
+        Emitter::vec(&mut out, Some(&*source.codemap())).emit(&[self.diagnostic(source)]);
         String::from_utf8_lossy(&out).to_string()
     }
 
     /// Render a fancy, human-readable output of this warning and
     /// print it to stderr. If rendered in a terminal that supports
     /// colours and font styles, the output will include those.
-    pub fn fancy_format_stderr(&self, codemap: &CodeMap) {
-        Emitter::stderr(ColorConfig::Auto, Some(codemap)).emit(&[self.diagnostic(codemap)]);
+    pub fn fancy_format_stderr(&self, source: &SourceCode) {
+        Emitter::stderr(ColorConfig::Auto, Some(&*source.codemap()))
+            .emit(&[self.diagnostic(source)]);
     }
 
     /// Create the optional span label displayed as an annotation on
@@ -53,7 +55,7 @@ impl EvalWarning {
 
     /// Create the primary warning message displayed to users for a
     /// warning.
-    fn message(&self, codemap: &CodeMap) -> String {
+    fn message(&self, source: &SourceCode) -> String {
         match self.kind {
             WarningKind::DeprecatedLiteralURL => {
                 "URL literal syntax is deprecated, use a quoted string instead".to_string()
@@ -64,11 +66,9 @@ impl EvalWarning {
             }
 
             WarningKind::UnusedBinding => {
-                let file = codemap.find_file(self.span.low());
-
                 format!(
                     "variable '{}' is declared, but never used:",
-                    file.source_slice(self.span)
+                    source.source_slice(self.span)
                 )
             }
 
@@ -99,7 +99,7 @@ impl EvalWarning {
         }
     }
 
-    fn diagnostic(&self, codemap: &CodeMap) -> Diagnostic {
+    fn diagnostic(&self, source: &SourceCode) -> Diagnostic {
         let span_label = SpanLabel {
             label: self.span_label(),
             span: self.span,
@@ -108,7 +108,7 @@ impl EvalWarning {
 
         Diagnostic {
             level: Level::Warning,
-            message: self.message(codemap),
+            message: self.message(source),
             spans: vec![span_label],
             code: Some(self.code().into()),
         }