refactor(tvix/eval): add a LightSpan type for lighter span tracking r/5457

This type carries the information required for calculating a
span (i.e. the chunk and offset), instead of the span itself. The span
is then only calculated in cases where it is required (when throwing
errors).

This reduces the eval time for
`builtins.length (builtins.attrNames (import <nixpkgs> {}))` by *one
third*!

The data structure in chunks that carries span information reduces
in-memory size by trading off the speed of retrieving span
information. This is because the span information is only actually
required when throwing errors (or emitting warnings).

However, somewhere along the way we grew a dependency on carrying span
information in thunks (for correctly reporting error chains). Hitting
the code paths for span retrieval was expensive, and carrying the
spans in a different way would still be less cache-efficient. This
change is the best tradeoff I could come up with.

Refs: b/229.
Change-Id: I27d4c4b5c5f9be90ac47f2db61941e123a78a77b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7558
Reviewed-by: grfn <grfn@gws.fyi>
Tested-by: BuildkiteCI

5 files changed, 61 insertions, 12 deletions

diff --git a/tvix/eval/src/builtins/impure.rs b/tvix/eval/src/builtins/impure.rs
index ee6cb2afdee4..b086df2f843d 100644
--- a/tvix/eval/src/builtins/impure.rs
+++ b/tvix/eval/src/builtins/impure.rs
@@ -12,6 +12,7 @@ use crate::{
     compiler::GlobalsMap,
     errors::ErrorKind,
     observer::NoOpObserver,
+    spans::LightSpan,
     value::{Builtin, BuiltinArgument, NixAttrs, Thunk},
     vm::VM,
     SourceCode, Value,
@@ -176,7 +177,7 @@ pub fn builtins_import(globals: &Weak<GlobalsMap>, source: SourceCode) -> Builti
             let res = entry
                 .insert(Value::Thunk(Thunk::new_suspended(
                     result.lambda,
-                    current_span,
+                    LightSpan::new_actual(current_span),
                 )))
                 .clone();
 
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index b8d8fb11a247..ae31d2714513 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -28,6 +28,7 @@ use crate::chunk::Chunk;
 use crate::errors::{Error, ErrorKind, EvalResult};
 use crate::observer::CompilerObserver;
 use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, UpvalueIdx};
+use crate::spans::LightSpan;
 use crate::spans::ToSpan;
 use crate::value::{Closure, Formals, Lambda, Thunk, Value};
 use crate::warnings::{EvalWarning, WarningKind};
@@ -946,7 +947,7 @@ impl Compiler<'_> {
         if lambda.upvalue_count == 0 {
             self.emit_constant(
                 if is_suspended_thunk {
-                    Value::Thunk(Thunk::new_suspended(lambda, span))
+                    Value::Thunk(Thunk::new_suspended(lambda, LightSpan::new_actual(span)))
                 } else {
                     Value::Closure(Rc::new(Closure::new(lambda)))
                 },
diff --git a/tvix/eval/src/spans.rs b/tvix/eval/src/spans.rs
index 9998e438b220..d2a8e9badf86 100644
--- a/tvix/eval/src/spans.rs
+++ b/tvix/eval/src/spans.rs
@@ -1,9 +1,46 @@
 //! Utilities for dealing with span tracking in the compiler and in
 //! error reporting.
 
+use crate::opcode::CodeIdx;
+use crate::value::Lambda;
 use codemap::{File, Span};
 use rnix::ast;
 use rowan::ast::AstNode;
+use std::rc::Rc;
+
+/// Helper struct to carry information required for making a span, but
+/// without actually performing the (expensive) span lookup.
+///
+/// This is used for tracking spans across thunk boundaries, as they
+/// are frequently instantiated but spans are only used in error or
+/// warning cases.
+#[derive(Clone, Debug)]
+pub enum LightSpan {
+    /// The span has already been computed and can just be used right
+    /// away.
+    Actual { span: Span },
+
+    /// The span needs to be computed from the provided data, but only
+    /// when it is required.
+    Delayed { lambda: Rc<Lambda>, offset: CodeIdx },
+}
+
+impl LightSpan {
+    pub fn new_delayed(lambda: Rc<Lambda>, offset: CodeIdx) -> Self {
+        Self::Delayed { lambda, offset }
+    }
+
+    pub fn new_actual(span: Span) -> Self {
+        Self::Actual { span }
+    }
+
+    pub fn span(&self) -> Span {
+        match self {
+            LightSpan::Actual { span } => *span,
+            LightSpan::Delayed { lambda, offset } => lambda.chunk.get_span(*offset),
+        }
+    }
+}
 
 /// Trait implemented by all types from which we can retrieve a span.
 pub trait ToSpan {
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 420c5fe039aa..680cc9b1fbb1 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -24,11 +24,10 @@ use std::{
     rc::Rc,
 };
 
-use codemap::Span;
-
 use crate::{
     chunk::Chunk,
     errors::{Error, ErrorKind},
+    spans::LightSpan,
     upvalues::Upvalues,
     value::{Builtin, Closure},
     vm::VM,
@@ -49,7 +48,7 @@ enum ThunkRepr {
     Suspended {
         lambda: Rc<Lambda>,
         upvalues: Rc<Upvalues>,
-        span: Span,
+        light_span: LightSpan,
     },
 
     /// Thunk currently under-evaluation; encountering a blackhole
@@ -77,11 +76,11 @@ impl Thunk {
         )))))
     }
 
-    pub fn new_suspended(lambda: Rc<Lambda>, span: Span) -> Self {
+    pub fn new_suspended(lambda: Rc<Lambda>, light_span: LightSpan) -> Self {
         Thunk(Rc::new(RefCell::new(ThunkRepr::Suspended {
             upvalues: Rc::new(Upvalues::with_capacity(lambda.upvalue_count)),
             lambda: lambda.clone(),
-            span,
+            light_span,
         })))
     }
 
@@ -124,7 +123,7 @@ impl Thunk {
         Thunk(Rc::new(RefCell::new(ThunkRepr::Suspended {
             lambda: Rc::new(lambda),
             upvalues: Rc::new(Upvalues::with_capacity(0)),
-            span: span,
+            light_span: LightSpan::new_actual(span),
         })))
     }
 
@@ -151,12 +150,16 @@ impl Thunk {
                     if let ThunkRepr::Suspended {
                         lambda,
                         upvalues,
-                        span,
+                        light_span,
                     } = std::mem::replace(&mut *thunk_mut, ThunkRepr::Blackhole)
                     {
                         drop(thunk_mut);
-                        vm.enter_frame(lambda, upvalues, 0)
-                            .map_err(|e| ErrorKind::ThunkForce(Box::new(Error { span, ..e })))?;
+                        vm.enter_frame(lambda, upvalues, 0).map_err(|e| {
+                            ErrorKind::ThunkForce(Box::new(Error {
+                                span: light_span.span(),
+                                ..e
+                            }))
+                        })?;
                         (*self.0.borrow_mut()) = ThunkRepr::Evaluated(vm.pop())
                     }
                 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index f0764c314fb3..37d30a6c0178 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -10,6 +10,7 @@ use crate::{
     nix_search_path::NixSearchPath,
     observer::RuntimeObserver,
     opcode::{CodeIdx, Count, JumpOffset, OpCode, StackIdx, UpvalueIdx},
+    spans::LightSpan,
     unwrap_or_clone_rc,
     upvalues::Upvalues,
     value::{Builtin, Closure, CoercionKind, Lambda, NixAttrs, NixList, Thunk, Value},
@@ -211,6 +212,12 @@ impl<'o> VM<'o> {
         self.chunk().get_span(self.frame().ip - 1)
     }
 
+    /// Returns the information needed to calculate the current span,
+    /// but without performing that calculation.
+    fn current_light_span(&self) -> LightSpan {
+        LightSpan::new_delayed(self.frame().lambda.clone(), self.frame().ip - 1)
+    }
+
     /// Construct an error from the given ErrorKind and the source
     /// span of the current instruction.
     pub fn error(&self, kind: ErrorKind) -> Error {
@@ -884,7 +891,7 @@ impl<'o> VM<'o> {
                     );
                     Thunk::new_closure(blueprint)
                 } else {
-                    Thunk::new_suspended(blueprint, self.current_span())
+                    Thunk::new_suspended(blueprint, self.current_light_span())
                 };
                 let upvalues = thunk.upvalues_mut();
                 self.push(Value::Thunk(thunk.clone()));