fix(tvix/eval): Actually trace spans for thunks

Currently, the span on *all* thunk force errors is the span at which the
thunk is forced, which for recursive thunk forcing ends up just being
the same span over and over again. This changes the span on thunk force
errors to be the span at which point the thunk is *created*, which is a
bit more helpful (though the printing atm is a little... crowded). To
make this work, we have to thread through the span at which a thunk is
created into a field on the thunk itself.

Change-Id: I81474810a763046e2eb3a8f07acf7d8ec708824a
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6932
Autosubmit: grfn <grfn@gws.fyi>
Reviewed-by: Adam Joseph <adam@westernsemico.com>
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

tvix/eval/src/builtins/impure.rs

@@ -151,7 +151,7 @@ pub fn builtins_import(
 
             // Compilation succeeded, we can construct a thunk from whatever it spat
             // out and return that.
-            Ok(Value::Thunk(Thunk::new(result.lambda)))
+            Ok(Value::Thunk(Thunk::new(result.lambda, vm.current_span())))
         },
     )
 }

tvix/eval/src/compiler/mod.rs

@@ -940,7 +940,7 @@ impl Compiler<'_> {
         // Emit the thunk directly if it does not close over the
         // environment.
         if lambda.upvalue_count == 0 {
-            self.emit_constant(Value::Thunk(Thunk::new(lambda)), node);
+            self.emit_constant(Value::Thunk(Thunk::new(lambda, span)), node);
             return;
         }
 

tvix/eval/src/value/thunk.rs

@@ -24,8 +24,10 @@ use std::{
     rc::Rc,
 };
 
+use codemap::Span;
+
 use crate::{
-    errors::ErrorKind,
+    errors::{Error, ErrorKind},
     upvalues::{UpvalueCarrier, Upvalues},
     vm::VM,
     Value,
@@ -52,14 +54,20 @@ enum ThunkRepr {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct Thunk(Rc<RefCell<ThunkRepr>>);
+pub struct Thunk {
+    inner: Rc<RefCell<ThunkRepr>>,
+    span: Span,
+}
 
 impl Thunk {
-    pub fn new(lambda: Rc<Lambda>) -> Self {
-        Thunk(Rc::new(RefCell::new(ThunkRepr::Suspended {
-            upvalues: Upvalues::with_capacity(lambda.upvalue_count),
-            lambda: lambda.clone(),
-        })))
+    pub fn new(lambda: Rc<Lambda>, span: Span) -> Self {
+        Thunk {
+            inner: Rc::new(RefCell::new(ThunkRepr::Suspended {
+                upvalues: Upvalues::with_capacity(lambda.upvalue_count),
+                lambda: lambda.clone(),
+            })),
+            span,
+        }
     }
 
     /// Evaluate the content of a thunk, potentially repeatedly, until
@@ -71,11 +79,11 @@ impl Thunk {
     /// are replaced.
     pub fn force(&self, vm: &mut VM) -> Result<(), ErrorKind> {
         loop {
-            let mut thunk_mut = self.0.borrow_mut();
+            let mut thunk_mut = self.inner.borrow_mut();
 
             match *thunk_mut {
                 ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => {
-                    let inner_repr = inner_thunk.0.borrow().clone();
+                    let inner_repr = inner_thunk.inner.borrow().clone();
                     *thunk_mut = inner_repr;
                 }
 
@@ -87,10 +95,14 @@ impl Thunk {
                         std::mem::replace(&mut *thunk_mut, ThunkRepr::Blackhole)
                     {
                         drop(thunk_mut);
-                        vm.enter_frame(lambda, upvalues, 0)
-                            .map_err(|e| ErrorKind::ThunkForce(Box::new(e)))?;
+                        vm.enter_frame(lambda, upvalues, 0).map_err(|e| {
+                            ErrorKind::ThunkForce(Box::new(Error {
+                                span: self.span,
+                                ..e
+                            }))
+                        })?;
                         let evaluated = ThunkRepr::Evaluated(vm.pop());
-                        (*self.0.borrow_mut()) = evaluated;
+                        (*self.inner.borrow_mut()) = evaluated;
                     }
                 }
             }
@@ -104,7 +116,7 @@ impl Thunk {
     // difficult to represent in the type system without impacting the
     // API too much.
     pub fn value(&self) -> Ref<Value> {
-        Ref::map(self.0.borrow(), |thunk| {
+        Ref::map(self.inner.borrow(), |thunk| {
             if let ThunkRepr::Evaluated(value) = thunk {
                 return value;
             }
@@ -116,7 +128,7 @@ impl Thunk {
 
 impl UpvalueCarrier for Thunk {
     fn upvalue_count(&self) -> usize {
-        if let ThunkRepr::Suspended { lambda, .. } = &*self.0.borrow() {
+        if let ThunkRepr::Suspended { lambda, .. } = &*self.inner.borrow() {
             return lambda.upvalue_count;
         }
 
@@ -124,23 +136,23 @@ impl UpvalueCarrier for Thunk {
     }
 
     fn upvalues(&self) -> Ref<'_, Upvalues> {
-        Ref::map(self.0.borrow(), |thunk| match thunk {
+        Ref::map(self.inner.borrow(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
             _ => panic!("upvalues() on non-suspended thunk"),
         })
     }
 
     fn upvalues_mut(&self) -> RefMut<'_, Upvalues> {
-        RefMut::map(self.0.borrow_mut(), |thunk| match thunk {
+        RefMut::map(self.inner.borrow_mut(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
-            _ => panic!("upvalues() on non-suspended thunk"),
+            thunk => panic!("upvalues() on non-suspended thunk: {thunk:?}"),
         })
     }
 }
 
 impl Display for Thunk {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self.0.try_borrow() {
+        match self.inner.try_borrow() {
             Ok(repr) => match &*repr {
                 ThunkRepr::Evaluated(v) => v.fmt(f),
                 _ => f.write_str("internal[thunk]"),

tvix/eval/src/vm.rs

@@ -186,7 +186,7 @@ impl<'o> VM<'o> {
 
     /// Returns the source span of the instruction currently being
     /// executed.
-    fn current_span(&self) -> codemap::Span {
+    pub(crate) fn current_span(&self) -> codemap::Span {
         self.chunk().get_span(self.frame().ip - 1)
     }
 
@@ -637,7 +637,7 @@ impl<'o> VM<'o> {
                     };
 
                     let upvalue_count = blueprint.upvalue_count;
-                    let thunk = Thunk::new(blueprint);
+                    let thunk = Thunk::new(blueprint, self.current_span());
                     let upvalues = thunk.upvalues_mut();
 
                     self.push(Value::Thunk(thunk.clone()));