feat(tvix/eval): wrap Closure in Rc<> to match cppnix semantics r/5452

Change-Id: I595087eff943d38a9fc78a83d37e207bb2ab79bc
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7443
Reviewed-by: grfn <grfn@gws.fyi>
Tested-by: BuildkiteCI

6 files changed, 34 insertions, 61 deletions

diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index f75acd48ad..e4e9c14df6 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -323,7 +323,7 @@ mod pure_builtins {
 
     #[builtin("functionArgs")]
     fn builtin_function_args(_: &mut VM, f: Value) -> Result<Value, ErrorKind> {
-        let lambda = f.to_closure()?.lambda();
+        let lambda = &f.as_closure()?.lambda();
         let formals = if let Some(formals) = &lambda.formals {
             formals
         } else {
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 5be582fb28..87383c0ca9 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -948,7 +948,7 @@ impl Compiler<'_> {
                 if is_suspended_thunk {
                     Value::Thunk(Thunk::new_suspended(lambda, span))
                 } else {
-                    Value::Closure(Closure::new(lambda))
+                    Value::Closure(Rc::new(Closure::new(lambda)))
                 },
                 node,
             );
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index 7a21223e04..4c2f3514f0 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -66,9 +66,6 @@ impl Lambda {
 pub struct Closure {
     pub lambda: Rc<Lambda>,
     pub upvalues: Rc<Upvalues>,
-    /// true if all upvalues have been realised
-    #[cfg(debug_assertions)]
-    pub is_finalised: bool,
 }
 
 impl Closure {
@@ -79,19 +76,8 @@ impl Closure {
         )
     }
 
-    /// Do not call this function unless you have read
-    /// `tvix/docs/value-pointer-equality.md` carefully.
-    pub fn ptr_eq(&self, other: &Self) -> bool {
-        Rc::ptr_eq(&self.lambda, &other.lambda) && Rc::ptr_eq(&self.upvalues, &other.upvalues)
-    }
-
     pub fn new_with_upvalues(upvalues: Rc<Upvalues>, lambda: Rc<Lambda>) -> Self {
-        Closure {
-            upvalues,
-            lambda,
-            #[cfg(debug_assertions)]
-            is_finalised: true,
-        }
+        Closure { upvalues, lambda }
     }
 
     pub fn chunk(&self) -> &Chunk {
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index 2bca9e6d32..af75bc9a73 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -41,7 +41,7 @@ pub enum Value {
     Path(PathBuf),
     Attrs(Rc<NixAttrs>),
     List(NixList),
-    Closure(Closure),
+    Closure(Rc<Closure>), // must use Rc<Closure> here in order to get proper pointer equality
     Builtin(Builtin),
 
     // Internal values that, while they technically exist at runtime,
@@ -304,7 +304,13 @@ impl Value {
     gen_cast!(to_str, NixString, "string", Value::String(s), s.clone());
     gen_cast!(to_attrs, Rc<NixAttrs>, "set", Value::Attrs(a), a.clone());
     gen_cast!(to_list, NixList, "list", Value::List(l), l.clone());
-    gen_cast!(to_closure, Closure, "lambda", Value::Closure(c), c.clone());
+    gen_cast!(
+        as_closure,
+        Rc<Closure>,
+        "lambda",
+        Value::Closure(c),
+        c.clone()
+    );
 
     gen_cast_mut!(as_list_mut, NixList, "list", List);
 
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 0d4c26bab4..54d2b31a2b 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -69,12 +69,10 @@ pub struct Thunk(Rc<RefCell<ThunkRepr>>);
 impl Thunk {
     pub fn new_closure(lambda: Rc<Lambda>) -> Self {
         Thunk(Rc::new(RefCell::new(ThunkRepr::Evaluated(Value::Closure(
-            Closure {
+            Rc::new(Closure {
                 upvalues: Rc::new(Upvalues::with_capacity(lambda.upvalue_count)),
                 lambda: lambda.clone(),
-                #[cfg(debug_assertions)]
-                is_finalised: false,
-            },
+            }),
         )))))
     }
 
@@ -124,13 +122,6 @@ impl Thunk {
 
     pub fn finalise(&self, stack: &[Value]) {
         self.upvalues_mut().resolve_deferred_upvalues(stack);
-        #[cfg(debug_assertions)]
-        {
-            let inner: &mut ThunkRepr = &mut self.0.as_ref().borrow_mut();
-            if let ThunkRepr::Evaluated(Value::Closure(c)) = inner {
-                c.is_finalised = true;
-            }
-        }
     }
 
     pub fn is_evaluated(&self) -> bool {
@@ -145,19 +136,7 @@ impl Thunk {
     // API too much.
     pub fn value(&self) -> Ref<Value> {
         Ref::map(self.0.borrow(), |thunk| match thunk {
-            ThunkRepr::Evaluated(value) => {
-                #[cfg(debug_assertions)]
-                if matches!(
-                    value,
-                    Value::Closure(Closure {
-                        is_finalised: false,
-                        ..
-                    })
-                ) {
-                    panic!("Thunk::value called on an unfinalised closure");
-                }
-                return value;
-            }
+            ThunkRepr::Evaluated(value) => value,
             ThunkRepr::Blackhole => panic!("Thunk::value called on a black-holed thunk"),
             ThunkRepr::Suspended { .. } => panic!("Thunk::value called on a suspended thunk"),
         })
@@ -166,7 +145,7 @@ impl Thunk {
     pub fn upvalues(&self) -> Ref<'_, Upvalues> {
         Ref::map(self.0.borrow(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
-            ThunkRepr::Evaluated(Value::Closure(Closure { upvalues, .. })) => upvalues,
+            ThunkRepr::Evaluated(Value::Closure(c)) => &c.upvalues,
             _ => panic!("upvalues() on non-suspended thunk"),
         })
     }
@@ -174,19 +153,12 @@ impl Thunk {
     pub fn upvalues_mut(&self) -> RefMut<'_, Upvalues> {
         RefMut::map(self.0.borrow_mut(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
-            ThunkRepr::Evaluated(Value::Closure(Closure {
-                upvalues,
-                #[cfg(debug_assertions)]
-                is_finalised,
-                ..
-            })) => {
-                #[cfg(debug_assertions)]
-                if *is_finalised {
-                    panic!("Thunk::upvalues_mut() called on a finalised closure");
-                }
-                Rc::get_mut(upvalues)
-                    .expect("upvalues_mut() was called on a thunk which already had multiple references to it")
-            }
+            ThunkRepr::Evaluated(Value::Closure(c)) => Rc::get_mut(
+                &mut Rc::get_mut(c).unwrap().upvalues,
+            )
+            .expect(
+                "upvalues_mut() was called on a thunk which already had multiple references to it",
+            ),
             thunk => panic!("upvalues() on non-suspended thunk: {thunk:?}"),
         })
     }
@@ -194,7 +166,16 @@ impl Thunk {
     /// Do not use this without first reading and understanding
     /// `tvix/docs/value-pointer-equality.md`.
     pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
-        Rc::ptr_eq(&self.0, &other.0)
+        if Rc::ptr_eq(&self.0, &other.0) {
+            return true;
+        }
+        match &*self.0.borrow() {
+            ThunkRepr::Evaluated(Value::Closure(c1)) => match &*other.0.borrow() {
+                ThunkRepr::Evaluated(Value::Closure(c2)) => Rc::ptr_eq(c1, c2),
+                _ => false,
+            },
+            _ => false,
+        }
     }
 }
 
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index bfca5d5984..4ea155f5a3 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -533,7 +533,7 @@ impl<'o> VM<'o> {
                 (v1, v2) => {
                     if allow_pointer_equality_on_functions_and_thunks {
                         if let (Value::Closure(c1), Value::Closure(c2)) = (&v1, &v2) {
-                            if c1.ptr_eq(c2) {
+                            if Rc::ptr_eq(c1, c2) {
                                 continue;
                             }
                         }
@@ -864,10 +864,10 @@ 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(
+                self.push(Value::Closure(Rc::new(Closure::new_with_upvalues(
                     Rc::new(upvalues),
                     blueprint,
-                )));
+                ))));
             }
 
             OpCode::OpThunkSuspended(idx) | OpCode::OpThunkClosure(idx) => {