2 files changed, 36 insertions, 4 deletions
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 9394dce48453..819a37ec71eb 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -927,7 +927,7 @@ impl Compiler<'_> {
         if lambda.upvalue_count == 0 {
             self.emit_constant(
                 if is_thunk {
-                    Value::Thunk(Thunk::new(lambda, span))
+                    Value::Thunk(Thunk::new_suspended(lambda, span))
                 } else {
                     Value::Closure(Closure::new(lambda))
                 },
@@ -942,11 +942,11 @@ impl Compiler<'_> {
         // which the result can be constructed.
         let blueprint_idx = self.chunk().push_constant(Value::Blueprint(lambda));
 
-        self.push_op(
+        let code_idx = self.push_op(
             if is_thunk {
-                OpCode::OpThunk(blueprint_idx)
+                OpCode::OpThunkSuspended(blueprint_idx)
             } else {
-                OpCode::OpClosure(blueprint_idx)
+                OpCode::OpThunkClosure(blueprint_idx)
             },
             node,
         );
@@ -957,6 +957,23 @@ impl Compiler<'_> {
             compiled.scope.upvalues,
             compiled.captures_with_stack,
         );
+
+        if !is_thunk && !self.scope()[outer_slot].needs_finaliser {
+            if !self.scope()[outer_slot].must_thunk {
+                // The closure has upvalues, but is not recursive.  Therefore no thunk is required,
+                // which saves us the overhead of Rc<RefCell<>>
+                self.chunk()[code_idx] = OpCode::OpClosure(blueprint_idx);
+            } else {
+                // This case occurs when a closure has upvalue-references to itself but does not need a
+                // finaliser.  Since no OpFinalise will be emitted later on we synthesize one here.
+                // It is needed here only to set [`Closure::is_finalised`] which is used for sanity checks.
+                #[cfg(debug_assertions)]
+                self.push_op(
+                    OpCode::OpFinalise(self.scope().stack_index(outer_slot)),
+                    &self.span_for(node),
+                );
+            }
+        }
     }
 
     fn compile_apply(&mut self, slot: LocalIdx, node: &ast::Apply) {
@@ -996,6 +1013,10 @@ impl Compiler<'_> {
                         self.push_op(OpCode::DataDeferredLocal(stack_idx), &upvalue.span);
                         self.scope_mut().mark_needs_finaliser(slot);
                     } else {
+                        // a self-reference
+                        if this_depth == target_depth && this_stack_slot == stack_idx {
+                            self.scope_mut().mark_must_thunk(slot);
+                        }
                         self.push_op(OpCode::DataLocalIdx(stack_idx), &upvalue.span);
                     }
                 }
diff --git a/tvix/eval/src/compiler/scope.rs b/tvix/eval/src/compiler/scope.rs
index 88df218ebdf7..9269a3c44c76 100644
--- a/tvix/eval/src/compiler/scope.rs
+++ b/tvix/eval/src/compiler/scope.rs
@@ -50,6 +50,9 @@ pub struct Local {
     /// Does this local need to be finalised after the enclosing scope
     /// is completely constructed?
     pub needs_finaliser: bool,
+
+    /// Does this local's upvalues contain a reference to itself?
+    pub must_thunk: bool,
 }
 
 impl Local {
@@ -228,6 +231,7 @@ impl Scope {
             name: LocalName::Phantom,
             depth: self.scope_depth,
             needs_finaliser: false,
+            must_thunk: false,
             used: true,
         });
 
@@ -243,6 +247,7 @@ impl Scope {
             depth: self.scope_depth,
             initialised: false,
             needs_finaliser: false,
+            must_thunk: false,
             used: false,
         });
 
@@ -259,6 +264,12 @@ impl Scope {
         self.locals[idx.0].needs_finaliser = true;
     }
 
+    /// Mark local as must be wrapped in a thunk.  This happens if
+    /// the local has a reference to itself in its upvalues.
+    pub fn mark_must_thunk(&mut self, idx: LocalIdx) {
+        self.locals[idx.0].must_thunk = true;
+    }
+
     /// Compute the runtime stack index for a given local by
     /// accounting for uninitialised variables at scopes below this
     /// one.