feat(tvix/eval): wrap Closure::upvalues in Rc r/5324

See cl/7372; Nix equality semantics require the ability to track
pointer equality of upvalue-sets.

Signed-off-by: Adam Joseph <adam@westernsemico.com>
Change-Id: I82ba517499cf370189a80355e4e46a5caaab7153
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7373
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

3 files changed, 11 insertions, 6 deletions

diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index c9664112f1cd..9ac35d43d837 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -65,7 +65,7 @@ impl Lambda {
 #[derive(Clone, Debug)]
 pub struct Closure {
     pub lambda: Rc<Lambda>,
-    pub upvalues: Upvalues,
+    pub upvalues: Rc<Upvalues>,
     /// true if all upvalues have been realised
     #[cfg(debug_assertions)]
     pub is_finalised: bool,
@@ -73,10 +73,13 @@ pub struct Closure {
 
 impl Closure {
     pub fn new(lambda: Rc<Lambda>) -> Self {
-        Self::new_with_upvalues(Upvalues::with_capacity(lambda.upvalue_count), lambda)
+        Self::new_with_upvalues(
+            Rc::new(Upvalues::with_capacity(lambda.upvalue_count)),
+            lambda,
+        )
     }
 
-    pub fn new_with_upvalues(upvalues: Upvalues, lambda: Rc<Lambda>) -> Self {
+    pub fn new_with_upvalues(upvalues: Rc<Upvalues>, lambda: Rc<Lambda>) -> Self {
         Closure {
             upvalues,
             lambda,
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 5107c6328da4..1be13bfe893f 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -70,7 +70,7 @@ impl Thunk {
     pub fn new_closure(lambda: Rc<Lambda>) -> Self {
         Thunk(Rc::new(RefCell::new(ThunkRepr::Evaluated(Value::Closure(
             Closure {
-                upvalues: Upvalues::with_capacity(lambda.upvalue_count),
+                upvalues: Rc::new(Upvalues::with_capacity(lambda.upvalue_count)),
                 lambda: lambda.clone(),
                 #[cfg(debug_assertions)]
                 is_finalised: false,
@@ -184,7 +184,8 @@ impl Thunk {
                 if *is_finalised {
                     panic!("Thunk::upvalues_mut() called on a finalised closure");
                 }
-                upvalues
+                Rc::get_mut(upvalues)
+                    .expect("upvalues_mut() was called on a thunk which already had multiple references to it")
             }
             thunk => panic!("upvalues() on non-suspended thunk: {thunk:?}"),
         })
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 400d85adac30..1ae97908e559 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -694,7 +694,8 @@ impl<'o> VM<'o> {
                 let mut upvalues = Upvalues::with_capacity(blueprint.upvalue_count);
                 self.populate_upvalues(upvalue_count, &mut upvalues)?;
                 self.push(Value::Closure(Closure::new_with_upvalues(
-                    upvalues, blueprint,
+                    Rc::new(upvalues),
+                    blueprint,
                 )));
             }