refactor(tvix/eval): encapsulate internal mutability within Closure r/4651

This is required to efficiently construct the upvalue array at
runtime, as there are situations where during Closure construction
multiple things already have a reference to the closure (e.g. a
self-reference).

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

3 files changed, 47 insertions, 25 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 2fac7560a6f6..520eb3bd6404 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -87,8 +87,7 @@ impl Compiler {
     }
 
     fn chunk(&mut self) -> &mut Chunk {
-        Rc::<Chunk>::get_mut(self.context_mut().lambda.chunk())
-            .expect("compiler flaw: long-lived chunk reference")
+        &mut self.context_mut().lambda.chunk
     }
 
     fn scope(&self) -> &Scope {
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index 2b5fcf6c9819..d0209cc50725 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -1,12 +1,15 @@
 //! This module implements the runtime representation of functions.
-use std::rc::Rc;
+use std::{
+    cell::{Ref, RefCell},
+    rc::Rc,
+};
 
-use crate::{chunk::Chunk, Value};
+use crate::{chunk::Chunk, opcode::UpvalueIdx, Value};
 
 #[derive(Clone, Debug)]
 pub struct Lambda {
     // name: Option<NixString>,
-    pub(crate) chunk: Rc<Chunk>,
+    pub(crate) chunk: Chunk,
     pub(crate) upvalue_count: usize,
 }
 
@@ -19,22 +22,42 @@ impl Lambda {
         }
     }
 
-    pub fn chunk(&mut self) -> &mut Rc<Chunk> {
+    pub fn chunk(&mut self) -> &mut Chunk {
         &mut self.chunk
     }
 }
 
 #[derive(Clone, Debug)]
-pub struct Closure {
+pub struct InnerClosure {
     pub lambda: Lambda,
     pub upvalues: Vec<Value>,
 }
 
+#[repr(transparent)]
+#[derive(Clone, Debug)]
+pub struct Closure(Rc<RefCell<InnerClosure>>);
+
 impl Closure {
     pub fn new(lambda: Lambda) -> Self {
-        Closure {
+        Closure(Rc::new(RefCell::new(InnerClosure {
             upvalues: Vec::with_capacity(lambda.upvalue_count),
             lambda,
-        }
+        })))
+    }
+
+    pub fn chunk(&self) -> Ref<'_, Chunk> {
+        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 upvalue_count(&self) -> usize {
+        self.0.borrow().lambda.upvalue_count
+    }
+
+    pub fn push_upvalue(&self, value: Value) {
+        self.0.borrow_mut().upvalues.push(value)
     }
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 707d2f4b26b3..527f58ddc8d6 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -1,12 +1,12 @@
 //! This module implements the virtual (or abstract) machine that runs
 //! Tvix bytecode.
 
-use std::rc::Rc;
+use std::{cell::Ref, rc::Rc};
 
 use crate::{
     chunk::Chunk,
     errors::{ErrorKind, EvalResult},
-    opcode::{Count, JumpOffset, OpCode, StackIdx, UpvalueIdx},
+    opcode::{Count, JumpOffset, OpCode, StackIdx},
     value::{Closure, Lambda, NixAttrs, NixList, Value},
 };
 
@@ -85,8 +85,8 @@ impl VM {
         &self.frames[self.frames.len() - 1]
     }
 
-    fn chunk(&self) -> &Chunk {
-        &self.frame().closure.lambda.chunk
+    fn chunk(&self) -> Ref<'_, Chunk> {
+        self.frame().closure.chunk()
     }
 
     fn frame_mut(&mut self) -> &mut CallFrame {
@@ -353,35 +353,35 @@ impl VM {
                     };
                 }
 
-                OpCode::OpGetUpvalue(UpvalueIdx(upv_idx)) => {
-                    let value = self.frame().closure.upvalues[upv_idx].clone();
+                OpCode::OpGetUpvalue(upv_idx) => {
+                    let value = self.frame().closure.upvalue(upv_idx).clone();
                     self.push(value);
                 }
 
                 OpCode::OpClosure(idx) => {
-                    let mut closure = self.chunk().constant(idx).clone().to_closure()?;
+                    let closure = self.chunk().constant(idx).clone().to_closure()?;
 
                     debug_assert!(
-                        closure.lambda.upvalue_count > 0,
+                        closure.upvalue_count() > 0,
                         "OpClosure should not be called for plain lambdas"
                     );
 
-                    for _ in 0..closure.lambda.upvalue_count {
+                    for _ in 0..closure.upvalue_count() {
                         match self.inc_ip() {
                             OpCode::DataLocalIdx(StackIdx(local_idx)) => {
                                 let idx = self.frame().stack_offset + local_idx;
-                                closure.upvalues.push(self.stack[idx].clone());
+                                closure.push_upvalue(self.stack[idx].clone());
                             }
 
-                            OpCode::DataUpvalueIdx(UpvalueIdx(upv_idx)) => {
-                                closure
-                                    .upvalues
-                                    .push(self.frame().closure.upvalues[upv_idx].clone());
+                            OpCode::DataUpvalueIdx(upv_idx) => {
+                                closure.push_upvalue(self.frame().closure.upvalue(upv_idx).clone());
                             }
 
                             OpCode::DataDynamicIdx(ident_idx) => {
-                                let ident = self.chunk().constant(ident_idx).as_str()?;
-                                closure.upvalues.push(self.resolve_with(ident)?);
+                                let chunk = self.chunk();
+                                let ident = chunk.constant(ident_idx).as_str()?.to_string();
+                                drop(chunk); // some lifetime trickery due to cell::Ref
+                                closure.push_upvalue(self.resolve_with(&ident)?);
                             }
 
                             _ => panic!("compiler error: missing closure operand"),