feat(tvix/eval): compile `with` expression r/4551

Adds an additional structure to the compiler's scope to track the
runtime "with stack", i.e. the stack of values through which
identifiers should be dynamically resolved within a with-scope.

When encountering a `with` expression, the value from which the
bindings should be resolved is pushed onto the stack and tracked by
the compiler in the "with stack", as well as with a "phantom value"
which indicates that the stack contains an additional slot which is
not available to users via identifiers.

Runtime handling of this is not yet implemented.

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

3 files changed, 65 insertions, 11 deletions

diff --git a/tvix/eval/src/compiler.rs b/tvix/eval/src/compiler.rs
index c5516d199add..5edf83535b67 100644
--- a/tvix/eval/src/compiler.rs
+++ b/tvix/eval/src/compiler.rs
@@ -40,6 +40,17 @@ struct Local {
 
     // Scope depth of this local.
     depth: usize,
+
+    // Phantom locals are not actually accessible by users (e.g.
+    // intermediate values used for `with`).
+    phantom: bool,
+}
+
+/// Represents a stack offset containing keys which are currently
+/// in-scope through a with expression.
+#[derive(Debug)]
+struct With {
+    depth: usize,
 }
 
 /// Represents a scope known during compilation, which can be resolved
@@ -54,6 +65,10 @@ struct Scope {
 
     // How many scopes "deep" are these locals?
     scope_depth: usize,
+
+    // Stack indices of attribute sets currently in scope through
+    // `with`.
+    with_stack: Vec<With>,
 }
 
 struct Compiler {
@@ -140,6 +155,11 @@ impl Compiler {
                 self.compile_let_in(node)
             }
 
+            rnix::SyntaxKind::NODE_WITH => {
+                let node = rnix::types::With::cast(node).unwrap();
+                self.compile_with(node)
+            }
+
             kind => panic!("visiting unsupported node: {:?}", kind),
         }
     }
@@ -689,7 +709,7 @@ impl Compiler {
     // Unless in a non-standard scope, the encountered values are
     // simply pushed on the stack and their indices noted in the
     // entries vector.
-    fn compile_let_in(&mut self, node: rnix::types::LetIn) -> Result<(), Error> {
+    fn compile_let_in(&mut self, node: rnix::types::LetIn) -> EvalResult<()> {
         self.begin_scope();
         let mut entries = vec![];
         let mut from_inherits = vec![];
@@ -709,10 +729,7 @@ impl Compiler {
 
                 Some(_) => {
                     for ident in inherit.idents() {
-                        self.scope.locals.push(Local {
-                            name: ident.as_str().to_string(),
-                            depth: self.scope.scope_depth,
-                        });
+                        self.push_local(ident.as_str());
                     }
                     from_inherits.push(inherit);
                 }
@@ -732,11 +749,7 @@ impl Compiler {
             }
 
             entries.push(entry.value().unwrap());
-
-            self.scope.locals.push(Local {
-                name: path.pop().unwrap(),
-                depth: self.scope.scope_depth,
-            });
+            self.push_local(path.pop().unwrap());
         }
 
         // Now we can add instructions to look up each inherited value
@@ -766,6 +779,26 @@ impl Compiler {
         Ok(())
     }
 
+    // Compile `with` expressions by emitting instructions that
+    // pop/remove the indices of attribute sets that are implicitly in
+    // scope through `with` on the "with-stack".
+    fn compile_with(&mut self, node: rnix::types::With) -> EvalResult<()> {
+        // TODO: Detect if the namespace is just an identifier, and
+        // resolve that directly (thus avoiding duplication on the
+        // stack).
+        self.compile(node.namespace().unwrap())?;
+
+        self.push_phantom();
+        self.scope.with_stack.push(With {
+            depth: self.scope.scope_depth,
+        });
+
+        self.chunk
+            .push_op(OpCode::OpPushWith(self.scope.locals.len() - 1));
+
+        self.compile(node.body().unwrap())
+    }
+
     // Emit the literal string value of an identifier. Required for
     // several operations related to attribute sets, where identifiers
     // are used as string keys.
@@ -819,11 +852,27 @@ impl Compiler {
         }
     }
 
+    fn push_local<S: Into<String>>(&mut self, name: S) {
+        self.scope.locals.push(Local {
+            name: name.into(),
+            depth: self.scope.scope_depth,
+            phantom: false,
+        });
+    }
+
+    fn push_phantom(&mut self) {
+        self.scope.locals.push(Local {
+            name: "".into(),
+            depth: self.scope.scope_depth,
+            phantom: true,
+        });
+    }
+
     fn resolve_local(&mut self, name: &str) -> Option<usize> {
         let scope = &self.scope;
 
         for (idx, local) in scope.locals.iter().enumerate().rev() {
-            if local.name == name {
+            if !local.phantom && local.name == name {
                 return Some(idx);
             }
         }
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index ce17dab63124..4a4d47045fbd 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -52,6 +52,9 @@ pub enum OpCode {
     OpAttrOrNotFound,
     OpAttrsIsSet,
 
+    // `with`-handling
+    OpPushWith(usize),
+
     // Lists
     OpList(usize),
     OpConcat,
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index d18cd7977bbe..881d79ba9853 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -276,6 +276,8 @@ impl VM {
                     let value = self.stack[local_idx].clone();
                     self.push(value)
                 }
+
+                OpCode::OpPushWith(_idx) => todo!("with handling not implemented"),
             }
 
             #[cfg(feature = "disassembler")]