fix(tvix/eval): instantiate *new* closures from blueprints each time r/4659

The previous closure refactoring introduced a bug in which the same
closure object would get mutated constantly for each instance of a
closure, which is incorrect behaviour.

This commit instead introduces an explicit new Value variant for the
internal "blueprint" that the compiler generates (essentially just the
lambda) and uses this variant to construct the closure at runtime.

If the blueprint ever leaks out to a user somehow that is a critical
bug and tvix-eval will panic.

As a ~treat~ test for this, the fibonacci function is being used as it
is a self-recursive closure (i.e. different instantiations of the same
"blueprint") getting called with different values and it's good to
have it around.

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

6 files changed, 33 insertions, 17 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 78fe76ca01a9..0d2ed66f5bf5 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -805,16 +805,17 @@ impl Compiler {
         // If the function is not a closure, just emit it directly and
         // move on.
         if compiled.lambda.upvalue_count == 0 {
-            self.emit_constant(Value::Closure(Closure::new(compiled.lambda)));
+            self.emit_constant(Value::Closure(Closure::new(Rc::new(compiled.lambda))));
             return;
         }
 
         // If the function is a closure, we need to emit the variable
         // number of operands that allow the runtime to close over the
-        // upvalues.
+        // upvalues and leave a blueprint in the constant index from
+        // which the runtime closure can be constructed.
         let closure_idx = self
             .chunk()
-            .push_constant(Value::Closure(Closure::new(compiled.lambda)));
+            .push_constant(Value::Blueprint(Rc::new(compiled.lambda)));
 
         self.chunk().push_op(OpCode::OpClosure(closure_idx));
         for upvalue in compiled.scope.upvalues {
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-fib.exp b/tvix/eval/src/tests/tvix_tests/eval-okay-fib.exp
new file mode 100644
index 000000000000..8643cf6debac
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-fib.exp
@@ -0,0 +1 @@
+89
diff --git a/tvix/eval/src/tests/tvix_tests/eval-okay-fib.nix b/tvix/eval/src/tests/tvix_tests/eval-okay-fib.nix
new file mode 100644
index 000000000000..9a22d85ac5f1
--- /dev/null
+++ b/tvix/eval/src/tests/tvix_tests/eval-okay-fib.nix
@@ -0,0 +1,7 @@
+let
+  fib' = i: n: m: if i == 0
+    then n
+    else fib' (i - 1) m (n + m);
+
+  fib = n: fib' n 1 1;
+in fib 10
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index d0209cc50725..e5db43d58ace 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -29,7 +29,7 @@ impl Lambda {
 
 #[derive(Clone, Debug)]
 pub struct InnerClosure {
-    pub lambda: Lambda,
+    pub lambda: Rc<Lambda>,
     pub upvalues: Vec<Value>,
 }
 
@@ -38,7 +38,7 @@ pub struct InnerClosure {
 pub struct Closure(Rc<RefCell<InnerClosure>>);
 
 impl Closure {
-    pub fn new(lambda: Lambda) -> Self {
+    pub fn new(lambda: Rc<Lambda>) -> Self {
         Closure(Rc::new(RefCell::new(InnerClosure {
             upvalues: Vec::with_capacity(lambda.upvalue_count),
             lambda,
diff --git a/tvix/eval/src/value/mod.rs b/tvix/eval/src/value/mod.rs
index 54211e8ba313..911af9d6ae12 100644
--- a/tvix/eval/src/value/mod.rs
+++ b/tvix/eval/src/value/mod.rs
@@ -35,6 +35,7 @@ pub enum Value {
     AttrPath(Vec<NixString>),
     AttrNotFound,
     DynamicUpvalueMissing(NixString),
+    Blueprint(Rc<Lambda>),
 }
 
 impl Value {
@@ -55,9 +56,10 @@ impl Value {
             Value::Closure(_) | Value::Builtin(_) => "lambda",
 
             // Internal types
-            Value::AttrPath(_) | Value::AttrNotFound | Value::DynamicUpvalueMissing(_) => {
-                "internal"
-            }
+            Value::AttrPath(_)
+            | Value::AttrNotFound
+            | Value::DynamicUpvalueMissing(_)
+            | Value::Blueprint(_) => "internal",
         }
     }
 
@@ -166,6 +168,7 @@ impl Display for Value {
             // internal types
             Value::AttrPath(path) => write!(f, "internal[attrpath({})]", path.len()),
             Value::AttrNotFound => f.write_str("internal[not found]"),
+            Value::Blueprint(_) => f.write_str("internal[blueprint]"),
             Value::DynamicUpvalueMissing(name) => {
                 write!(f, "internal[no_dyn_upvalue({name})]")
             }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 5d5cb57dbabf..7b3de6406549 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -385,15 +385,19 @@ impl VM {
                 }
 
                 OpCode::OpClosure(idx) => {
-                    let value = self.chunk().constant(idx).clone();
-                    self.push(value.clone());
+                    let blueprint = match self.chunk().constant(idx) {
+                        Value::Blueprint(lambda) => lambda.clone(),
+                        _ => panic!("compiler bug: non-blueprint in blueprint slot"),
+                    };
+
+                    let closure = Closure::new(blueprint);
+                    self.push(Value::Closure(closure.clone()));
 
-                    // This refers to the same Rc, and from this point
-                    // on internally mutates the closure objects
-                    // upvalues. The closure is already in its stack
-                    // slot, which means that it can capture itself as
-                    // an upvalue for self-recursion.
-                    let closure = value.to_closure()?;
+                    // From this point on we internally mutate the
+                    // closure object's upvalues. The closure is
+                    // already in its stack slot, which means that it
+                    // can capture itself as an upvalue for
+                    // self-recursion.
 
                     debug_assert!(
                         closure.upvalue_count() > 0,
@@ -536,6 +540,6 @@ pub fn run_lambda(lambda: Lambda) -> EvalResult<Value> {
         with_stack: vec![],
     };
 
-    vm.call(Closure::new(lambda), 0);
+    vm.call(Closure::new(Rc::new(lambda)), 0);
     vm.run()
 }