feat(tvix/eval): track span of first force in a thunk blackhole r/6025

This is step 1 towards being able to use all 4 spans that we know when
dealing with infinite recursion. It tracks the span at which the
force of a thunk was first requested when constructing a blackhole, so
that we can highlight the spans of the first and second forces.

These are actually the least relevant spans, but the easiest to put in
place, more coming soon.

Change-Id: I4c7e82f6211b98756439d4148a4191457cc46807
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8269
Autosubmit: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

5 files changed, 48 insertions, 17 deletions

diff --git a/tvix/eval/src/errors.rs b/tvix/eval/src/errors.rs
index 5e890c68a74c..76a3da3ff52f 100644
--- a/tvix/eval/src/errors.rs
+++ b/tvix/eval/src/errors.rs
@@ -77,7 +77,9 @@ pub enum ErrorKind {
     NotCallable(&'static str),
 
     /// Infinite recursion encountered while forcing thunks.
-    InfiniteRecursion,
+    InfiniteRecursion {
+        first_force: Span,
+    },
 
     ParseErrors(Vec<rnix::parser::ParseError>),
 
@@ -354,7 +356,7 @@ to a missing value in the attribute set(s) included via `with`."#,
                 )
             }
 
-            ErrorKind::InfiniteRecursion => write!(f, "infinite recursion encountered"),
+            ErrorKind::InfiniteRecursion { .. } => write!(f, "infinite recursion encountered"),
 
             // Errors themselves ignored here & handled in Self::spans instead
             ErrorKind::ParseErrors(_) => write!(f, "failed to parse Nix code:"),
@@ -754,7 +756,7 @@ impl Error {
             | ErrorKind::UnknownDynamicVariable(_)
             | ErrorKind::VariableAlreadyDefined(_)
             | ErrorKind::NotCallable(_)
-            | ErrorKind::InfiniteRecursion
+            | ErrorKind::InfiniteRecursion { .. }
             | ErrorKind::ParseErrors(_)
             | ErrorKind::NativeError { .. }
             | ErrorKind::BytecodeError(_)
@@ -797,7 +799,7 @@ impl Error {
             ErrorKind::UnknownDynamicVariable(_) => "E011",
             ErrorKind::VariableAlreadyDefined(_) => "E012",
             ErrorKind::NotCallable(_) => "E013",
-            ErrorKind::InfiniteRecursion => "E014",
+            ErrorKind::InfiniteRecursion { .. } => "E014",
             ErrorKind::ParseErrors(_) => "E015",
             ErrorKind::DuplicateAttrsKey { .. } => "E016",
             ErrorKind::NotCoercibleToString { .. } => "E018",
@@ -869,6 +871,21 @@ impl Error {
                 ]
             }
 
+            ErrorKind::InfiniteRecursion { first_force } => {
+                vec![
+                    SpanLabel {
+                        label: Some("first requested here".into()),
+                        span: *first_force,
+                        style: SpanStyle::Secondary,
+                    },
+                    SpanLabel {
+                        label: Some("requested again here".into()),
+                        span: self.span,
+                        style: SpanStyle::Primary,
+                    },
+                ]
+            }
+
             // All other errors pretty much have the same shape.
             _ => {
                 vec![SpanLabel {
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index f57cfda5eecf..77d40b0e3906 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -22,6 +22,7 @@ mod thunk;
 
 use crate::errors::ErrorKind;
 use crate::opcode::StackIdx;
+use crate::spans::LightSpan;
 use crate::vm::generators::{self, GenCo};
 use crate::AddContext;
 pub use attrs::NixAttrs;
@@ -611,9 +612,9 @@ impl Value {
     }
 
     // TODO(amjoseph): de-asyncify this (when called directly by the VM)
-    pub async fn force(self, co: GenCo) -> Result<Value, ErrorKind> {
+    pub async fn force(self, co: GenCo, span: LightSpan) -> Result<Value, ErrorKind> {
         if let Value::Thunk(thunk) = self {
-            return thunk.force(co).await;
+            return thunk.force(co, span).await;
         }
 
         Ok(self)
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index f2959ba15c74..0290e73ebb7d 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -65,7 +65,7 @@ enum ThunkRepr {
 
     /// Thunk currently under-evaluation; encountering a blackhole
     /// value means that infinite recursion has occured.
-    Blackhole,
+    Blackhole { forced_at: LightSpan },
 
     /// Fully evaluated thunk.
     Evaluated(Value),
@@ -75,7 +75,7 @@ impl ThunkRepr {
     fn debug_repr(&self) -> String {
         match self {
             ThunkRepr::Evaluated(v) => format!("thunk(val|{})", v),
-            ThunkRepr::Blackhole => "thunk(blackhole)".to_string(),
+            ThunkRepr::Blackhole { .. } => "thunk(blackhole)".to_string(),
             ThunkRepr::Native(_) => "thunk(native)".to_string(),
             ThunkRepr::Suspended { lambda, .. } => format!("thunk({:p})", *lambda),
         }
@@ -114,7 +114,7 @@ impl Thunk {
     }
 
     // TODO(amjoseph): de-asyncify this
-    pub async fn force(self, co: GenCo) -> Result<Value, ErrorKind> {
+    pub async fn force(self, co: GenCo, span: LightSpan) -> Result<Value, ErrorKind> {
         // If the current thunk is already fully evaluated, return its evaluated
         // value. The VM will continue running the code that landed us here.
         if self.is_forced() {
@@ -124,12 +124,14 @@ impl Thunk {
         // Begin evaluation of this thunk by marking it as a blackhole, meaning
         // that any other forcing frame encountering this thunk before its
         // evaluation is completed detected an evaluation cycle.
-        let inner = self.0.replace(ThunkRepr::Blackhole);
+        let inner = self.0.replace(ThunkRepr::Blackhole { forced_at: span });
 
         match inner {
             // If there was already a blackhole in the thunk, this is an
             // evaluation cycle.
-            ThunkRepr::Blackhole => Err(ErrorKind::InfiniteRecursion),
+            ThunkRepr::Blackhole { forced_at } => Err(ErrorKind::InfiniteRecursion {
+                first_force: forced_at.span(),
+            }),
 
             // If there is a native function stored in the thunk, evaluate it
             // and replace this thunk's representation with the result.
@@ -206,7 +208,7 @@ impl Thunk {
     pub fn value(&self) -> Ref<Value> {
         Ref::map(self.0.borrow(), |thunk| match thunk {
             ThunkRepr::Evaluated(value) => value,
-            ThunkRepr::Blackhole => panic!("Thunk::value called on a black-holed thunk"),
+            ThunkRepr::Blackhole { .. } => panic!("Thunk::value called on a black-holed thunk"),
             ThunkRepr::Suspended { .. } | ThunkRepr::Native(_) => {
                 panic!("Thunk::value called on a suspended thunk")
             }
diff --git a/tvix/eval/src/vm/generators.rs b/tvix/eval/src/vm/generators.rs
index 853ab063465f..b7d8f7aeb5c5 100644
--- a/tvix/eval/src/vm/generators.rs
+++ b/tvix/eval/src/vm/generators.rs
@@ -293,7 +293,9 @@ impl<'o> VM<'o> {
                         // back to the outer VM loop.
                         VMRequest::ForceValue(value) => {
                             self.reenqueue_generator(name, span.clone(), generator);
-                            self.enqueue_generator("force", span, |co| value.force(co));
+                            self.enqueue_generator("force", span.clone(), |co| {
+                                value.force(co, span)
+                            });
                             return Ok(false);
                         }
 
@@ -313,7 +315,9 @@ impl<'o> VM<'o> {
                             self.reenqueue_generator(name, span.clone(), generator);
 
                             let value = self.stack[self.with_stack[idx]].clone();
-                            self.enqueue_generator("force", span, |co| value.force(co));
+                            self.enqueue_generator("force", span.clone(), |co| {
+                                value.force(co, span)
+                            });
 
                             return Ok(false);
                         }
@@ -328,7 +332,9 @@ impl<'o> VM<'o> {
                                 .expect("Tvix bug: generator requested captured with-value, but there is no call frame");
 
                             let value = call_frame.upvalues.with_stack().unwrap()[idx].clone();
-                            self.enqueue_generator("force", span, |co| value.force(co));
+                            self.enqueue_generator("force", span.clone(), |co| {
+                                value.force(co, span)
+                            });
 
                             return Ok(false);
                         }
@@ -441,7 +447,9 @@ impl<'o> VM<'o> {
                                 "generator should be reenqueued with the same frame ID"
                             );
 
-                            self.enqueue_generator("force", span, |co| value.force(co));
+                            self.enqueue_generator("force", span.clone(), |co| {
+                                value.force(co, span)
+                            });
                             return Ok(false);
                         }
 
diff --git a/tvix/eval/src/vm/mod.rs b/tvix/eval/src/vm/mod.rs
index 450d950b34a6..0fb1f9cd8cc9 100644
--- a/tvix/eval/src/vm/mod.rs
+++ b/tvix/eval/src/vm/mod.rs
@@ -483,7 +483,10 @@ impl<'o> VM<'o> {
                         let gen_span = frame.current_light_span();
 
                         self.push_call_frame(span, frame);
-                        self.enqueue_generator("force", gen_span, |co| thunk.force(co));
+                        self.enqueue_generator("force", gen_span.clone(), |co| {
+                            thunk.force(co, gen_span)
+                        });
+
                         return Ok(false);
                     }
                 }