refactor(tvix/eval): introduce UpvalueCarrier trait r/4677

This trait abstracts over the commonalities of upvalue handling
between closures and thunks.

It allows the VM to simplify the code used for setting up upvalues,
without duplicating between the two different types.

Note that this does not yet refactor the VM code to optimally make use
of this.

Change-Id: If8de5181f26ae1fa00d554f1ae6ea473ee4b6070
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6347
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>

2 files changed, 40 insertions, 18 deletions

diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index e8301d82da..4f5cc1154b 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 57c38e0894..3b7cb7f5f7 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"),
+        })
+    }
+}