refactor(tvix/eval): introduce Upvalues struct in closures & thunks r/4800

This struct will be responsible for tracking upvalues (and is a
convenient place to introduce optimisations for reducing value clones)
instead of a plain value vector.

The main motivation for this is that the upvalues will have to capture
the `with`-stack fully and I want to avoid duplicating the logic for
this between the two capturing types.

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

4 files changed, 65 insertions, 30 deletions

diff --git a/tvix/eval/src/upvalues.rs b/tvix/eval/src/upvalues.rs
index 9287e7abdf..b51ef500e0 100644
--- a/tvix/eval/src/upvalues.rs
+++ b/tvix/eval/src/upvalues.rs
@@ -3,34 +3,61 @@
 //! as well as closures (lambdas that capture variables from the
 //! surrounding scope).
 
-use std::cell::{Ref, RefMut};
+use std::{
+    cell::{Ref, RefMut},
+    ops::Index,
+};
 
 use crate::{opcode::UpvalueIdx, Value};
 
+/// Structure for carrying upvalues inside of thunks & closures. The
+/// implementation of this struct encapsulates the logic for capturing
+/// and accessing upvalues.
+#[derive(Clone, Debug)]
+pub struct Upvalues {
+    upvalues: Vec<Value>,
+}
+
+impl Upvalues {
+    pub fn with_capacity(count: usize) -> Self {
+        Upvalues {
+            upvalues: Vec::with_capacity(count),
+        }
+    }
+
+    /// Push an upvalue at the end of the upvalue list.
+    pub fn push(&mut self, value: Value) {
+        self.upvalues.push(value);
+    }
+}
+
+impl Index<UpvalueIdx> for Upvalues {
+    type Output = Value;
+
+    fn index(&self, index: UpvalueIdx) -> &Self::Output {
+        &self.upvalues[index.0]
+    }
+}
+
 /// `UpvalueCarrier` is implemented by all types that carry upvalues.
 pub trait UpvalueCarrier {
     fn upvalue_count(&self) -> usize;
 
     /// Read-only accessor for the stored upvalues.
-    fn upvalues(&self) -> Ref<'_, [Value]>;
+    fn upvalues(&self) -> Ref<'_, Upvalues>;
 
     /// Mutable accessor for stored upvalues.
-    fn upvalues_mut(&self) -> RefMut<'_, Vec<Value>>;
+    fn upvalues_mut(&self) -> RefMut<'_, Upvalues>;
 
     /// Read an upvalue at the given index.
     fn upvalue(&self, idx: UpvalueIdx) -> Ref<'_, Value> {
-        Ref::map(self.upvalues(), |v| &v[idx.0])
-    }
-
-    /// Push an upvalue at the end of the upvalue list.
-    fn push_upvalue(&self, value: Value) {
-        self.upvalues_mut().push(value);
+        Ref::map(self.upvalues(), |v| &v.upvalues[idx.0])
     }
 
     /// Resolve deferred upvalues from the provided stack slice,
     /// mutating them in the internal upvalue slots.
     fn resolve_deferred_upvalues(&self, stack: &[Value]) {
-        for upvalue in self.upvalues_mut().iter_mut() {
+        for upvalue in self.upvalues_mut().upvalues.iter_mut() {
             if let Value::DeferredUpvalue(idx) = upvalue {
                 *upvalue = stack[idx.0].clone();
             }
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index 4f5cc1154b..758994504e 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -4,7 +4,10 @@ use std::{
     rc::Rc,
 };
 
-use crate::{chunk::Chunk, upvalues::UpvalueCarrier, Value};
+use crate::{
+    chunk::Chunk,
+    upvalues::{UpvalueCarrier, Upvalues},
+};
 
 #[derive(Clone, Debug)]
 pub struct Lambda {
@@ -30,7 +33,7 @@ impl Lambda {
 #[derive(Clone, Debug)]
 pub struct InnerClosure {
     pub lambda: Rc<Lambda>,
-    pub upvalues: Vec<Value>,
+    upvalues: Upvalues,
 }
 
 #[repr(transparent)]
@@ -40,7 +43,7 @@ pub struct Closure(Rc<RefCell<InnerClosure>>);
 impl Closure {
     pub fn new(lambda: Rc<Lambda>) -> Self {
         Closure(Rc::new(RefCell::new(InnerClosure {
-            upvalues: Vec::with_capacity(lambda.upvalue_count),
+            upvalues: Upvalues::with_capacity(lambda.upvalue_count),
             lambda,
         })))
     }
@@ -59,11 +62,11 @@ impl UpvalueCarrier for Closure {
         self.0.borrow().lambda.upvalue_count
     }
 
-    fn upvalues(&self) -> Ref<'_, [Value]> {
-        Ref::map(self.0.borrow(), |c| c.upvalues.as_slice())
+    fn upvalues(&self) -> Ref<'_, Upvalues> {
+        Ref::map(self.0.borrow(), |c| &c.upvalues)
     }
 
-    fn upvalues_mut(&self) -> RefMut<'_, Vec<Value>> {
+    fn upvalues_mut(&self) -> RefMut<'_, Upvalues> {
         RefMut::map(self.0.borrow_mut(), |c| &mut c.upvalues)
     }
 }
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index ad56b5eaf6..5098be186f 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -24,7 +24,12 @@ use std::{
     rc::Rc,
 };
 
-use crate::{errors::ErrorKind, upvalues::UpvalueCarrier, vm::VM, Value};
+use crate::{
+    errors::ErrorKind,
+    upvalues::{UpvalueCarrier, Upvalues},
+    vm::VM,
+    Value,
+};
 
 use super::Lambda;
 
@@ -35,7 +40,7 @@ enum ThunkRepr {
     /// execution.
     Suspended {
         lambda: Rc<Lambda>,
-        upvalues: Vec<Value>,
+        upvalues: Upvalues,
     },
 
     /// Thunk currently under-evaluation; encountering a blackhole
@@ -52,7 +57,7 @@ pub struct Thunk(Rc<RefCell<ThunkRepr>>);
 impl Thunk {
     pub fn new(lambda: Rc<Lambda>) -> Self {
         Thunk(Rc::new(RefCell::new(ThunkRepr::Suspended {
-            upvalues: Vec::with_capacity(lambda.upvalue_count),
+            upvalues: Upvalues::with_capacity(lambda.upvalue_count),
             lambda: lambda.clone(),
         })))
     }
@@ -119,14 +124,14 @@ impl UpvalueCarrier for Thunk {
         panic!("upvalues() on non-suspended thunk");
     }
 
-    fn upvalues(&self) -> Ref<'_, [Value]> {
+    fn upvalues(&self) -> Ref<'_, Upvalues> {
         Ref::map(self.0.borrow(), |thunk| match thunk {
-            ThunkRepr::Suspended { upvalues, .. } => upvalues.as_slice(),
+            ThunkRepr::Suspended { upvalues, .. } => upvalues,
             _ => panic!("upvalues() on non-suspended thunk"),
         })
     }
 
-    fn upvalues_mut(&self) -> RefMut<'_, Vec<Value>> {
+    fn upvalues_mut(&self) -> RefMut<'_, Upvalues> {
         RefMut::map(self.0.borrow_mut(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
             _ => panic!("upvalues() on non-suspended thunk"),
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index c124a653a0..31964cf909 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -8,13 +8,13 @@ use crate::{
     errors::{Error, ErrorKind, EvalResult},
     observer::Observer,
     opcode::{CodeIdx, ConstantIdx, Count, JumpOffset, OpCode, StackIdx, UpvalueIdx},
-    upvalues::UpvalueCarrier,
+    upvalues::{UpvalueCarrier, Upvalues},
     value::{Builtin, Closure, Lambda, NixAttrs, NixList, Thunk, Value},
 };
 
 struct CallFrame {
     lambda: Rc<Lambda>,
-    upvalues: Vec<Value>,
+    upvalues: Upvalues,
     ip: usize,
     stack_offset: usize,
 }
@@ -22,7 +22,7 @@ struct CallFrame {
 impl CallFrame {
     /// Retrieve an upvalue from this frame at the given index.
     fn upvalue(&self, idx: UpvalueIdx) -> &Value {
-        &self.upvalues[idx.0]
+        &self.upvalues[idx]
     }
 }
 
@@ -176,7 +176,7 @@ impl<'o> VM<'o> {
     pub fn call(
         &mut self,
         lambda: Rc<Lambda>,
-        upvalues: Vec<Value>,
+        upvalues: Upvalues,
         arg_count: usize,
     ) -> EvalResult<Value> {
         self.observer
@@ -443,7 +443,7 @@ impl<'o> VM<'o> {
                     match callable {
                         Value::Closure(closure) => {
                             let result =
-                                self.call(closure.lambda(), closure.upvalues().to_vec(), 1)?;
+                                self.call(closure.lambda(), closure.upvalues().clone(), 1)?;
                             self.push(result)
                         }
 
@@ -468,7 +468,7 @@ impl<'o> VM<'o> {
                             // closure.
                             let mut frame = self.frame_mut();
                             frame.lambda = lambda;
-                            frame.upvalues = closure.upvalues().to_vec();
+                            frame.upvalues = closure.upvalues().clone();
                             frame.ip = 0; // reset instruction pointer to beginning
                         }
 
@@ -658,7 +658,7 @@ impl<'o> VM<'o> {
     fn populate_upvalues(
         &mut self,
         count: usize,
-        mut upvalues: RefMut<'_, Vec<Value>>,
+        mut upvalues: RefMut<'_, Upvalues>,
     ) -> EvalResult<()> {
         for _ in 0..count {
             match self.inc_ip() {
@@ -744,7 +744,7 @@ fn unwrap_or_clone_rc<T: Clone>(rc: Rc<T>) -> T {
 
 pub fn run_lambda(observer: &mut dyn Observer, lambda: Rc<Lambda>) -> EvalResult<Value> {
     let mut vm = VM::new(observer);
-    let value = vm.call(lambda, vec![], 0)?;
+    let value = vm.call(lambda, Upvalues::with_capacity(0), 0)?;
     vm.force_for_output(&value)?;
     Ok(value)
 }