1 files changed, 37 insertions, 10 deletions

diff --git a/tvix/eval/src/upvalues.rs b/tvix/eval/src/upvalues.rs
index 9287e7abdf04..b51ef500e031 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();
             }