feat(tvix/eval): implement compilation of upvalue access r/4623

This adds a new upvalue tracking structure in the compiler to resolve
upvalues and track their positions within a function when compiling a
closure.

The compiler will emit runtime upvalue access instructions after this
commit, but the creation of the runtime closure object etc. is not yet
wired up.

Change-Id: Ib0c2c25f686bfd45f797c528753068858e3a770d
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6289
Tested-by: BuildkiteCI
Reviewed-by: grfn <grfn@gws.fyi>

4 files changed, 66 insertions, 4 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 2d888a6e3bed..3ade6831c286 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -22,7 +22,7 @@ use std::rc::Rc;
 
 use crate::chunk::Chunk;
 use crate::errors::{Error, ErrorKind, EvalResult};
-use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, StackIdx};
+use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, StackIdx, UpvalueIdx};
 use crate::value::{Closure, Lambda, Value};
 use crate::warnings::{EvalWarning, WarningKind};
 
@@ -63,6 +63,15 @@ struct With {
     depth: usize,
 }
 
+#[derive(Debug, PartialEq)]
+enum Upvalue {
+    /// This upvalue captures a local from the stack.
+    Stack(StackIdx),
+
+    /// This upvalue captures an enclosing upvalue.
+    Upvalue(UpvalueIdx),
+}
+
 /// Represents a scope known during compilation, which can be resolved
 /// directly to stack indices.
 ///
@@ -72,6 +81,7 @@ struct With {
 #[derive(Default)]
 struct Scope {
     locals: Vec<Local>,
+    upvalues: Vec<Upvalue>,
 
     // How many scopes "deep" are these locals?
     scope_depth: usize,
@@ -772,13 +782,19 @@ impl Compiler {
         match self.scope_mut().resolve_local(ident.text()) {
             Some(idx) => self.chunk().push_op(OpCode::OpGetLocal(idx)),
             None => {
+                // Are we possibly dealing with an upvalue?
+                if let Some(idx) = self.resolve_upvalue(self.contexts.len() - 1, ident.text()) {
+                    self.chunk().push_op(OpCode::OpGetUpvalue(idx));
+                    return;
+                }
+
                 if self.scope().with_stack.is_empty() {
                     self.emit_error(node.syntax().clone(), ErrorKind::UnknownStaticVariable);
                     return;
                 }
 
-                // Variable needs to be dynamically resolved
-                // at runtime.
+                // Variable needs to be dynamically resolved at
+                // runtime.
                 self.emit_constant(Value::String(ident.text().into()));
                 self.chunk().push_op(OpCode::OpResolveWith)
             }
@@ -976,6 +992,42 @@ impl Compiler {
         });
     }
 
+    fn resolve_upvalue(&mut self, ctx_idx: usize, name: &str) -> Option<UpvalueIdx> {
+        if ctx_idx == 0 {
+            // There can not be any upvalue at the outermost context.
+            return None;
+        }
+
+        if let Some(idx) = self.contexts[ctx_idx - 1].scope.resolve_local(name) {
+            return Some(self.add_upvalue(ctx_idx, Upvalue::Stack(idx)));
+        }
+
+        // If the upvalue comes from an enclosing context, we need to
+        // recurse to make sure that the upvalues are created at each
+        // level.
+        if let Some(idx) = self.resolve_upvalue(ctx_idx - 1, name) {
+            return Some(self.add_upvalue(ctx_idx, Upvalue::Upvalue(idx)));
+        }
+
+        None
+    }
+
+    fn add_upvalue(&mut self, ctx_idx: usize, upvalue: Upvalue) -> UpvalueIdx {
+        // If there is already an upvalue closing over the specified
+        // index, retrieve that instead.
+        for (idx, existing) in self.contexts[ctx_idx].scope.upvalues.iter().enumerate() {
+            if *existing == upvalue {
+                return UpvalueIdx(idx);
+            }
+        }
+
+        self.contexts[ctx_idx].scope.upvalues.push(upvalue);
+
+        let idx = UpvalueIdx(self.contexts[ctx_idx].lambda.upvalue_count);
+        self.contexts[ctx_idx].lambda.upvalue_count += 1;
+        idx
+    }
+
     fn emit_warning(&mut self, node: rnix::SyntaxNode, kind: WarningKind) {
         self.warnings.push(EvalWarning { node, kind })
     }
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index f2daf9bd8e52..6cce7ea9b176 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -13,9 +13,14 @@ pub struct CodeIdx(pub usize);
 
 /// Index of a value in the runtime stack.
 #[repr(transparent)]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub struct StackIdx(pub usize);
 
+/// Index of an upvalue within a closure's upvalue list.
+#[repr(transparent)]
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct UpvalueIdx(pub usize);
+
 /// Offset by which an instruction pointer should change in a jump.
 #[repr(transparent)]
 #[derive(Clone, Copy, Debug)]
@@ -99,4 +104,5 @@ pub enum OpCode {
 
     // Lambdas
     OpCall,
+    OpGetUpvalue(UpvalueIdx),
 }
diff --git a/tvix/eval/src/value/function.rs b/tvix/eval/src/value/function.rs
index 45efc24f09f8..5d7247416ee0 100644
--- a/tvix/eval/src/value/function.rs
+++ b/tvix/eval/src/value/function.rs
@@ -7,6 +7,7 @@ use crate::chunk::Chunk;
 pub struct Lambda {
     // name: Option<NixString>,
     pub(crate) chunk: Rc<Chunk>,
+    pub(crate) upvalue_count: usize,
 }
 
 impl Lambda {
@@ -14,6 +15,7 @@ impl Lambda {
         Lambda {
             // name: None,
             chunk: Default::default(),
+            upvalue_count: 0,
         }
     }
 
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 85f81bac6efa..1c7772f2a7e0 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -366,6 +366,8 @@ impl VM {
                         _ => return Err(ErrorKind::NotCallable.into()),
                     };
                 }
+
+                OpCode::OpGetUpvalue(_) => todo!("getting upvalues"),
             }
 
             #[cfg(feature = "disassembler")]