refactor(tvix/eval): add SourceCode directly into error types r/7571

With this change it's no longer necessary to track the SourceCode
struct separately from the evaluation for error reporting: It's just
stored directly in the errors.

This also ends up resolving an issue in compiler::bindings, where we
cloned the Arc containing file references way too often. In fact those
clones probably compensate for all additional SourceCode clones during
error construction now.

Change-Id: Ice93bf161e61f8ea3d48103435e20c53e6aa8c3a
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10986
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

6 files changed, 67 insertions, 40 deletions

diff --git a/tvix/eval/src/compiler/bindings.rs b/tvix/eval/src/compiler/bindings.rs
index 236b9dd2c3..634cc54022 100644
--- a/tvix/eval/src/compiler/bindings.rs
+++ b/tvix/eval/src/compiler/bindings.rs
@@ -261,10 +261,10 @@ impl TrackedBindings {
 trait HasEntryProxy {
     fn inherits(&self) -> Box<dyn Iterator<Item = ast::Inherit>>;
 
-    fn attributes(
+    fn attributes<'a>(
         &self,
-        file: Arc<codemap::File>,
-    ) -> Box<dyn Iterator<Item = (Span, PeekableAttrs, ast::Expr)>>;
+        file: &'a codemap::File,
+    ) -> Box<dyn Iterator<Item = (Span, PeekableAttrs, ast::Expr)> + 'a>;
 }
 
 impl<N: HasEntry> HasEntryProxy for N {
@@ -272,13 +272,13 @@ impl<N: HasEntry> HasEntryProxy for N {
         Box::new(ast::HasEntry::inherits(self))
     }
 
-    fn attributes(
+    fn attributes<'a>(
         &self,
-        file: Arc<codemap::File>,
-    ) -> Box<dyn Iterator<Item = (Span, PeekableAttrs, ast::Expr)>> {
+        file: &'a codemap::File,
+    ) -> Box<dyn Iterator<Item = (Span, PeekableAttrs, ast::Expr)> + 'a> {
         Box::new(ast::HasEntry::attrpath_values(self).map(move |entry| {
             (
-                entry.span_for(&file),
+                entry.span_for(file),
                 entry.attrpath().unwrap().attrs().peekable(),
                 entry.value().unwrap(),
             )
@@ -291,16 +291,16 @@ impl HasEntryProxy for AttributeSet {
         Box::new(self.inherits.clone().into_iter())
     }
 
-    fn attributes(
+    fn attributes<'a>(
         &self,
-        _: Arc<codemap::File>,
-    ) -> Box<dyn Iterator<Item = (Span, PeekableAttrs, ast::Expr)>> {
+        _: &'a codemap::File,
+    ) -> Box<dyn Iterator<Item = (Span, PeekableAttrs, ast::Expr)> + 'a> {
         Box::new(self.entries.clone().into_iter())
     }
 }
 
 /// AST-traversing functions related to bindings.
-impl Compiler<'_> {
+impl Compiler<'_, '_> {
     /// Compile all inherits of a node with entries that do *not* have a
     /// namespace to inherit from, and return the remaining ones that do.
     fn compile_plain_inherits<N>(
@@ -465,7 +465,7 @@ impl Compiler<'_> {
     ) where
         N: ToSpan + HasEntryProxy,
     {
-        for (span, mut path, value) in node.attributes(self.file.clone()) {
+        for (span, mut path, value) in node.attributes(self.file) {
             let key = path.next().unwrap();
 
             if bindings.try_merge(self, span, &key, path.clone(), value.clone()) {
@@ -765,7 +765,7 @@ impl Compiler<'_> {
 }
 
 /// Private compiler helpers related to bindings.
-impl Compiler<'_> {
+impl Compiler<'_, '_> {
     fn resolve_upvalue<N: ToSpan>(
         &mut self,
         ctx_idx: usize,
diff --git a/tvix/eval/src/compiler/import.rs b/tvix/eval/src/compiler/import.rs
index 6774c7b892..c56909e958 100644
--- a/tvix/eval/src/compiler/import.rs
+++ b/tvix/eval/src/compiler/import.rs
@@ -58,13 +58,14 @@ async fn import_impl(
     let result = crate::compiler::compile(
         &parsed.tree().expr().unwrap(),
         Some(path.clone()),
-        file,
         // The VM must ensure that a strong reference to the globals outlives
         // any self-references (which are weak) embedded within the globals. If
         // the expect() below panics, it means that did not happen.
         globals
             .upgrade()
             .expect("globals dropped while still in use"),
+        &source,
+        &file,
         &mut NoOpObserver::default(),
     )
     .map_err(|err| ErrorKind::ImportCompilerError {
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index d64202c2ad..0f9accc13e 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -24,7 +24,6 @@ use smol_str::SmolStr;
 use std::collections::{BTreeMap, HashMap};
 use std::path::{Path, PathBuf};
 use std::rc::{Rc, Weak};
-use std::sync::Arc;
 
 use crate::chunk::Chunk;
 use crate::errors::{CatchableErrorKind, Error, ErrorKind, EvalResult};
@@ -151,7 +150,7 @@ const GLOBAL_BUILTINS: &[&str] = &[
     "__curPos",
 ];
 
-pub struct Compiler<'observer> {
+pub struct Compiler<'source, 'observer> {
     contexts: Vec<LambdaCtx>,
     warnings: Vec<EvalWarning>,
     errors: Vec<Error>,
@@ -165,10 +164,13 @@ pub struct Compiler<'observer> {
     /// and a function that should emit code for the token.
     globals: Rc<GlobalsMap>,
 
-    /// File reference in the codemap contains all known source code
-    /// and is used to track the spans from which instructions where
-    /// derived.
-    file: Arc<codemap::File>,
+    /// Reference to the struct holding all of the source code, which
+    /// is used for error creation.
+    source: &'source SourceCode,
+
+    /// File reference in the source map for the current file, which
+    /// is used for creating spans.
+    file: &'source codemap::File,
 
     /// Carry an observer for the compilation process, which is called
     /// whenever a chunk is emitted.
@@ -180,18 +182,19 @@ pub struct Compiler<'observer> {
     dead_scope: usize,
 }
 
-impl Compiler<'_> {
+impl Compiler<'_, '_> {
     pub(super) fn span_for<S: ToSpan>(&self, to_span: &S) -> Span {
-        to_span.span_for(&self.file)
+        to_span.span_for(self.file)
     }
 }
 
 /// Compiler construction
-impl<'observer> Compiler<'observer> {
+impl<'source, 'observer> Compiler<'source, 'observer> {
     pub(crate) fn new(
         location: Option<PathBuf>,
-        file: Arc<codemap::File>,
         globals: Rc<GlobalsMap>,
+        source: &'source SourceCode,
+        file: &'source codemap::File,
         observer: &'observer mut dyn CompilerObserver,
     ) -> EvalResult<Self> {
         let mut root_dir = match location {
@@ -204,6 +207,7 @@ impl<'observer> Compiler<'observer> {
                             e
                         )),
                         file.span,
+                        source.clone(),
                     )
                 })?;
                 if let Some(dir) = location {
@@ -226,6 +230,7 @@ impl<'observer> Compiler<'observer> {
 
         Ok(Self {
             root_dir,
+            source,
             file,
             observer,
             globals,
@@ -239,7 +244,7 @@ impl<'observer> Compiler<'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]
     }
@@ -285,7 +290,7 @@ impl Compiler<'_> {
 }
 
 // Actual code-emitting AST traversal methods.
-impl Compiler<'_> {
+impl Compiler<'_, '_> {
     fn compile(&mut self, slot: LocalIdx, expr: ast::Expr) {
         let expr = optimiser::optimise_expr(self, slot, expr);
 
@@ -1473,7 +1478,8 @@ impl Compiler<'_> {
 
     fn emit_error<N: ToSpan>(&mut self, node: &N, kind: ErrorKind) {
         let span = self.span_for(node);
-        self.errors.push(Error::new(kind, span))
+        self.errors
+            .push(Error::new(kind, span, self.source.clone()))
     }
 }
 
@@ -1519,7 +1525,7 @@ fn expr_static_attr_str(node: &ast::Attr) -> Option<SmolStr> {
 fn compile_src_builtin(
     name: &'static str,
     code: &str,
-    source: &SourceCode,
+    source: SourceCode,
     weak: &Weak<GlobalsMap>,
 ) -> Value {
     use std::fmt::Write;
@@ -1542,8 +1548,9 @@ fn compile_src_builtin(
         let result = compile(
             &parsed.tree().expr().unwrap(),
             None,
-            file.clone(),
             weak.upgrade().unwrap(),
+            &source,
+            &file,
             &mut crate::observer::NoOpObserver {},
         )
         .map_err(|e| ErrorKind::NativeError {
@@ -1621,7 +1628,7 @@ pub fn prepare_globals(
         // If "source builtins" were supplied, compile them and insert
         // them.
         builtins.extend(src_builtins.into_iter().map(move |(name, code)| {
-            let compiled = compile_src_builtin(name, code, &source, weak);
+            let compiled = compile_src_builtin(name, code, source.clone(), weak);
             (name, compiled)
         }));
 
@@ -1647,11 +1654,12 @@ pub fn prepare_globals(
 pub fn compile(
     expr: &ast::Expr,
     location: Option<PathBuf>,
-    file: Arc<codemap::File>,
     globals: Rc<GlobalsMap>,
+    source: &SourceCode,
+    file: &codemap::File,
     observer: &mut dyn CompilerObserver,
 ) -> EvalResult<CompilationOutput> {
-    let mut c = Compiler::new(location, file, globals.clone(), observer)?;
+    let mut c = Compiler::new(location, globals.clone(), source, file, observer)?;
 
     let root_span = c.span_for(expr);
     let root_slot = c.scope_mut().declare_phantom(root_span, false);
diff --git a/tvix/eval/src/errors.rs b/tvix/eval/src/errors.rs
index b59ee675dd..0f17aafe87 100644
--- a/tvix/eval/src/errors.rs
+++ b/tvix/eval/src/errors.rs
@@ -293,10 +293,11 @@ pub struct Error {
     pub kind: ErrorKind,
     pub span: Span,
     pub contexts: Vec<String>,
+    pub source: SourceCode,
 }
 
 impl Error {
-    pub fn new(mut kind: ErrorKind, span: Span) -> Self {
+    pub fn new(mut kind: ErrorKind, span: Span, source: SourceCode) -> Self {
         let mut contexts = vec![];
         while let ErrorKind::WithContext {
             context,
@@ -311,6 +312,7 @@ impl Error {
             kind,
             span,
             contexts,
+            source,
         }
     }
 }
diff --git a/tvix/eval/src/lib.rs b/tvix/eval/src/lib.rs
index 5bb48baf05..9c298df60a 100644
--- a/tvix/eval/src/lib.rs
+++ b/tvix/eval/src/lib.rs
@@ -261,7 +261,7 @@ where
             code.as_ref(),
             file.clone(),
             location,
-            source,
+            source.clone(),
             self.builtins,
             self.src_builtins,
             self.enable_import,
@@ -295,6 +295,7 @@ where
             nix_path,
             self.io_handle,
             runtime_observer,
+            source,
             globals,
             lambda,
             self.strict,
@@ -335,6 +336,7 @@ fn parse_compile_internal(
         result.errors.push(Error::new(
             ErrorKind::ParseErrors(parse_errors.to_vec()),
             file.span,
+            source,
         ));
         return None;
     }
@@ -344,13 +346,15 @@ fn parse_compile_internal(
     // the result, in case the caller needs it for something.
     result.expr = parsed.tree().expr();
 
-    let builtins = crate::compiler::prepare_globals(builtins, src_builtins, source, enable_import);
+    let builtins =
+        crate::compiler::prepare_globals(builtins, src_builtins, source.clone(), enable_import);
 
     let compiler_result = match compiler::compile(
         result.expr.as_ref().unwrap(),
         location,
-        file,
         builtins,
+        &source,
+        &file,
         compiler_observer,
     ) {
         Ok(result) => result,
diff --git a/tvix/eval/src/vm/mod.rs b/tvix/eval/src/vm/mod.rs
index 0ca72bdf30..32f0d0456f 100644
--- a/tvix/eval/src/vm/mod.rs
+++ b/tvix/eval/src/vm/mod.rs
@@ -36,7 +36,7 @@ use crate::{
     },
     vm::generators::GenCo,
     warnings::{EvalWarning, WarningKind},
-    NixString,
+    NixString, SourceCode,
 };
 
 use generators::{call_functor, Generator, GeneratorState};
@@ -85,15 +85,18 @@ impl<T, S: GetSpan, IO> WithSpan<T, S, IO> for Result<T, ErrorKind> {
         match self {
             Ok(something) => Ok(something),
             Err(kind) => {
-                let mut error = Error::new(kind, top_span.get_span());
+                let mut error = Error::new(kind, top_span.get_span(), vm.source.clone());
 
                 // Wrap the top-level error in chaining errors for each element
                 // of the frame stack.
                 for frame in vm.frames.iter().rev() {
                     match frame {
                         Frame::CallFrame { span, .. } => {
-                            error =
-                                Error::new(ErrorKind::BytecodeError(Box::new(error)), span.span());
+                            error = Error::new(
+                                ErrorKind::BytecodeError(Box::new(error)),
+                                span.span(),
+                                vm.source.clone(),
+                            );
                         }
                         Frame::Generator { name, span, .. } => {
                             error = Error::new(
@@ -102,6 +105,7 @@ impl<T, S: GetSpan, IO> WithSpan<T, S, IO> for Result<T, ErrorKind> {
                                     gen_type: name,
                                 },
                                 span.span(),
+                                vm.source.clone(),
                             );
                         }
                     }
@@ -275,6 +279,10 @@ struct VM<'o, IO> {
     // TODO: should probably be based on a file hash
     pub import_cache: ImportCache,
 
+    /// Data structure holding all source code evaluated in this VM,
+    /// used for pretty error reporting.
+    source: SourceCode,
+
     /// Parsed Nix search path, which is used to resolve `<...>`
     /// references.
     nix_search_path: NixSearchPath,
@@ -333,6 +341,7 @@ where
         nix_search_path: NixSearchPath,
         io_handle: IO,
         observer: &'o mut dyn RuntimeObserver,
+        source: SourceCode,
         globals: Rc<GlobalsMap>,
         reasonable_span: LightSpan,
     ) -> Self {
@@ -342,6 +351,7 @@ where
             observer,
             globals,
             reasonable_span,
+            source,
             frames: vec![],
             stack: vec![],
             with_stack: vec![],
@@ -1315,6 +1325,7 @@ pub fn run_lambda<IO>(
     nix_search_path: NixSearchPath,
     io_handle: IO,
     observer: &mut dyn RuntimeObserver,
+    source: SourceCode,
     globals: Rc<GlobalsMap>,
     lambda: Rc<Lambda>,
     strict: bool,
@@ -1334,6 +1345,7 @@ where
         nix_search_path,
         io_handle,
         observer,
+        source,
         globals,
         root_span.into(),
     );