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>

1 files changed, 30 insertions, 7 deletions

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)
     }
 }