5 files changed, 81 insertions, 18 deletions
diff --git a/tvix/eval/src/lib.rs b/tvix/eval/src/lib.rs
index eebccc873956..7156a03b6444 100644
--- a/tvix/eval/src/lib.rs
+++ b/tvix/eval/src/lib.rs
@@ -4,6 +4,7 @@ mod compiler;
 mod errors;
 mod eval;
 mod opcode;
+mod upvalues;
 mod value;
 mod vm;
 mod warnings;
diff --git a/tvix/eval/src/upvalues.rs b/tvix/eval/src/upvalues.rs
new file mode 100644
index 000000000000..9287e7abdf04
--- /dev/null
+++ b/tvix/eval/src/upvalues.rs
@@ -0,0 +1,39 @@
+//! This module encapsulates some logic for upvalue handling, which is
+//! relevant to both thunks (delayed computations for lazy-evaluation)
+//! as well as closures (lambdas that capture variables from the
+//! surrounding scope).
+
+use std::cell::{Ref, RefMut};
+
+use crate::{opcode::UpvalueIdx, Value};
+
+/// `UpvalueCarrier` is implemented by all types that carry upvalues.
+pub trait UpvalueCarrier {
+    fn upvalue_count(&self) -> usize;
+
+    /// Read-only accessor for the stored upvalues.
+    fn upvalues(&self) -> Ref<'_, [Value]>;
+
+    /// Mutable accessor for stored upvalues.
+    fn upvalues_mut(&self) -> RefMut<'_, Vec<Value>>;
+
+    /// Read an upvalue at the given index.
+    fn upvalue(&self, idx: UpvalueIdx) -> Ref<'_, Value> {
+        Ref::map(self.upvalues(), |v| &v[idx.0])
+    }
+
+    /// Push an upvalue at the end of the upvalue list.
+    fn push_upvalue(&self, value: Value) {
+        self.upvalues_mut().push(value);
+    }
+
+    /// Resolve deferred upvalues from the provided stack slice,
+    /// mutating them in the internal upvalue slots.
+    fn resolve_deferred_upvalues(&self, stack: &[Value]) {
+        for upvalue in self.upvalues_mut().iter_mut() {
+            if let Value::DeferredUpvalue(idx) = upvalue {
+                *upvalue = stack[idx.0].clone();
+            }
+        }
+    }
+}
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index e8301d82da11..4f5cc1154bd7 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -1,10 +1,10 @@
 //! This module implements the runtime representation of functions.
 use std::{
-    cell::{Ref, RefCell},
+    cell::{Ref, RefCell, RefMut},
     rc::Rc,
 };
 
-use crate::{chunk::Chunk, opcode::UpvalueIdx, Value};
+use crate::{chunk::Chunk, upvalues::UpvalueCarrier, Value};
 
 #[derive(Clone, Debug)]
 pub struct Lambda {
@@ -49,26 +49,21 @@ impl Closure {
         Ref::map(self.0.borrow(), |c| &c.lambda.chunk)
     }
 
-    pub fn upvalue(&self, idx: UpvalueIdx) -> Ref<'_, Value> {
-        Ref::map(self.0.borrow(), |c| &c.upvalues[idx.0])
+    pub fn lambda(&self) -> Rc<Lambda> {
+        self.0.borrow().lambda.clone()
     }
+}
 
-    pub fn upvalue_count(&self) -> usize {
+impl UpvalueCarrier for Closure {
+    fn upvalue_count(&self) -> usize {
         self.0.borrow().lambda.upvalue_count
     }
 
-    pub fn push_upvalue(&self, value: Value) {
-        self.0.borrow_mut().upvalues.push(value)
+    fn upvalues(&self) -> Ref<'_, [Value]> {
+        Ref::map(self.0.borrow(), |c| c.upvalues.as_slice())
     }
 
-    /// Resolve the deferred upvalues in the closure from a slice of
-    /// the current stack, using the indices stored in the deferred
-    /// values.
-    pub fn resolve_deferred_upvalues(&self, stack: &[Value]) {
-        for upvalue in self.0.borrow_mut().upvalues.iter_mut() {
-            if let Value::DeferredUpvalue(idx) = upvalue {
-                *upvalue = stack[idx.0].clone();
-            }
-        }
+    fn upvalues_mut(&self) -> RefMut<'_, Vec<Value>> {
+        RefMut::map(self.0.borrow_mut(), |c| &mut c.upvalues)
     }
 }
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 57c38e089460..3b7cb7f5f711 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -18,9 +18,12 @@
 //! object, but when forcing a thunk, the runtime *must* mutate the
 //! memoisable slot.
 
-use std::{cell::RefCell, rc::Rc};
+use std::{
+    cell::{Ref, RefCell, RefMut},
+    rc::Rc,
+};
 
-use crate::Value;
+use crate::{upvalues::UpvalueCarrier, Value};
 
 use super::Lambda;
 
@@ -53,3 +56,27 @@ impl Thunk {
         })))
     }
 }
+
+impl UpvalueCarrier for Thunk {
+    fn upvalue_count(&self) -> usize {
+        if let ThunkRepr::Suspended { lambda, .. } = &*self.0.borrow() {
+            return lambda.upvalue_count;
+        }
+
+        panic!("upvalues() on non-suspended thunk");
+    }
+
+    fn upvalues(&self) -> Ref<'_, [Value]> {
+        Ref::map(self.0.borrow(), |thunk| match thunk {
+            ThunkRepr::Suspended { upvalues, .. } => upvalues.as_slice(),
+            _ => panic!("upvalues() on non-suspended thunk"),
+        })
+    }
+
+    fn upvalues_mut(&self) -> RefMut<'_, Vec<Value>> {
+        RefMut::map(self.0.borrow_mut(), |thunk| match thunk {
+            ThunkRepr::Suspended { upvalues, .. } => upvalues,
+            _ => panic!("upvalues() on non-suspended thunk"),
+        })
+    }
+}
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index b363522994eb..aa67cb1ea817 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -7,6 +7,7 @@ use crate::{
     chunk::Chunk,
     errors::{Error, ErrorKind, EvalResult},
     opcode::{ConstantIdx, Count, JumpOffset, OpCode, StackIdx},
+    upvalues::UpvalueCarrier,
     value::{Closure, Lambda, NixAttrs, NixList, Value},
 };