chore(tvix/eval): thread `slot` value through all compiler methods r/4692

With this change any compilation of an expression is aware of its own
stack slot if it is leaving identifiers on the stack.

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

1 files changed, 58 insertions, 54 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 1c8d0182a7c1..d34564a75992 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -110,20 +110,20 @@ impl Compiler {
         match expr {
             ast::Expr::Literal(literal) => self.compile_literal(literal),
             ast::Expr::Path(path) => self.compile_path(path),
-            ast::Expr::Str(s) => self.compile_str(s),
-            ast::Expr::UnaryOp(op) => self.compile_unary_op(op),
-            ast::Expr::BinOp(op) => self.compile_binop(op),
+            ast::Expr::Str(s) => self.compile_str(slot, s),
+            ast::Expr::UnaryOp(op) => self.compile_unary_op(slot, op),
+            ast::Expr::BinOp(op) => self.compile_binop(slot, op),
             ast::Expr::HasAttr(has_attr) => self.compile_has_attr(slot, has_attr),
-            ast::Expr::List(list) => self.compile_list(list),
+            ast::Expr::List(list) => self.compile_list(slot, list),
             ast::Expr::AttrSet(attrs) => self.compile_attr_set(slot, attrs),
             ast::Expr::Select(select) => self.compile_select(slot, select),
-            ast::Expr::Assert(assert) => self.compile_assert(assert),
-            ast::Expr::IfElse(if_else) => self.compile_if_else(if_else),
-            ast::Expr::LetIn(let_in) => self.compile_let_in(let_in),
+            ast::Expr::Assert(assert) => self.compile_assert(slot, assert),
+            ast::Expr::IfElse(if_else) => self.compile_if_else(slot, if_else),
+            ast::Expr::LetIn(let_in) => self.compile_let_in(slot, let_in),
             ast::Expr::Ident(ident) => self.compile_ident(slot, ident),
-            ast::Expr::With(with) => self.compile_with(with),
+            ast::Expr::With(with) => self.compile_with(slot, with),
             ast::Expr::Lambda(lambda) => self.compile_lambda(slot, lambda),
-            ast::Expr::Apply(apply) => self.compile_apply(apply),
+            ast::Expr::Apply(apply) => self.compile_apply(slot, apply),
 
             // Parenthesized expressions are simply unwrapped, leaving
             // their value on the stack.
@@ -188,7 +188,7 @@ impl Compiler {
         self.emit_constant(value);
     }
 
-    fn compile_str(&mut self, node: ast::Str) {
+    fn compile_str(&mut self, slot: Option<LocalIdx>, node: ast::Str) {
         let mut count = 0;
 
         // The string parts are produced in literal order, however
@@ -202,7 +202,7 @@ impl Compiler {
                 // Interpolated expressions are compiled as normal and
                 // dealt with by the VM before being assembled into
                 // the final string.
-                ast::InterpolPart::Interpolation(node) => self.compile(None, node.expr().unwrap()),
+                ast::InterpolPart::Interpolation(node) => self.compile(slot, node.expr().unwrap()),
 
                 ast::InterpolPart::Literal(lit) => {
                     self.emit_constant(Value::String(lit.into()));
@@ -215,8 +215,8 @@ impl Compiler {
         }
     }
 
-    fn compile_unary_op(&mut self, op: ast::UnaryOp) {
-        self.compile(None, op.expr().unwrap());
+    fn compile_unary_op(&mut self, slot: Option<LocalIdx>, op: ast::UnaryOp) {
+        self.compile(slot, op.expr().unwrap());
 
         let opcode = match op.operator().unwrap() {
             ast::UnaryOpKind::Invert => OpCode::OpInvert,
@@ -226,7 +226,7 @@ impl Compiler {
         self.chunk().push_op(opcode);
     }
 
-    fn compile_binop(&mut self, op: ast::BinOp) {
+    fn compile_binop(&mut self, slot: Option<LocalIdx>, op: ast::BinOp) {
         use ast::BinOpKind;
 
         // Short-circuiting and other strange operators, which are
@@ -235,17 +235,17 @@ impl Compiler {
         // used for standard binary operators).
 
         match op.operator().unwrap() {
-            BinOpKind::And => return self.compile_and(op),
-            BinOpKind::Or => return self.compile_or(op),
-            BinOpKind::Implication => return self.compile_implication(op),
+            BinOpKind::And => return self.compile_and(slot, op),
+            BinOpKind::Or => return self.compile_or(slot, op),
+            BinOpKind::Implication => return self.compile_implication(slot, op),
             _ => {}
         };
 
         // For all other operators, the two values need to be left on
         // the stack in the correct order before pushing the
         // instruction for the operation itself.
-        self.compile(None, op.lhs().unwrap());
-        self.compile(None, op.rhs().unwrap());
+        self.compile(slot, op.lhs().unwrap());
+        self.compile(slot, op.rhs().unwrap());
 
         match op.operator().unwrap() {
             BinOpKind::Add => self.chunk().push_op(OpCode::OpAdd),
@@ -272,7 +272,7 @@ impl Compiler {
         };
     }
 
-    fn compile_and(&mut self, node: ast::BinOp) {
+    fn compile_and(&mut self, slot: Option<LocalIdx>, node: ast::BinOp) {
         debug_assert!(
             matches!(node.operator(), Some(ast::BinOpKind::And)),
             "compile_and called with wrong operator kind: {:?}",
@@ -280,7 +280,7 @@ impl Compiler {
         );
 
         // Leave left-hand side value on the stack.
-        self.compile(None, node.lhs().unwrap());
+        self.compile(slot, node.lhs().unwrap());
 
         // If this value is false, jump over the right-hand side - the
         // whole expression is false.
@@ -290,13 +290,13 @@ impl Compiler {
         // right-hand side on the stack. Its result is now the value
         // of the whole expression.
         self.chunk().push_op(OpCode::OpPop);
-        self.compile(None, node.rhs().unwrap());
+        self.compile(slot, node.rhs().unwrap());
 
         self.patch_jump(end_idx);
         self.chunk().push_op(OpCode::OpAssertBool);
     }
 
-    fn compile_or(&mut self, node: ast::BinOp) {
+    fn compile_or(&mut self, slot: Option<LocalIdx>, node: ast::BinOp) {
         debug_assert!(
             matches!(node.operator(), Some(ast::BinOpKind::Or)),
             "compile_or called with wrong operator kind: {:?}",
@@ -304,18 +304,18 @@ impl Compiler {
         );
 
         // Leave left-hand side value on the stack
-        self.compile(None, node.lhs().unwrap());
+        self.compile(slot, node.lhs().unwrap());
 
         // Opposite of above: If this value is **true**, we can
         // short-circuit the right-hand side.
         let end_idx = self.chunk().push_op(OpCode::OpJumpIfTrue(JumpOffset(0)));
         self.chunk().push_op(OpCode::OpPop);
-        self.compile(None, node.rhs().unwrap());
+        self.compile(slot, node.rhs().unwrap());
         self.patch_jump(end_idx);
         self.chunk().push_op(OpCode::OpAssertBool);
     }
 
-    fn compile_implication(&mut self, node: ast::BinOp) {
+    fn compile_implication(&mut self, slot: Option<LocalIdx>, node: ast::BinOp) {
         debug_assert!(
             matches!(node.operator(), Some(ast::BinOpKind::Implication)),
             "compile_implication called with wrong operator kind: {:?}",
@@ -323,20 +323,20 @@ impl Compiler {
         );
 
         // Leave left-hand side value on the stack and invert it.
-        self.compile(None, node.lhs().unwrap());
+        self.compile(slot, node.lhs().unwrap());
         self.chunk().push_op(OpCode::OpInvert);
 
         // Exactly as `||` (because `a -> b` = `!a || b`).
         let end_idx = self.chunk().push_op(OpCode::OpJumpIfTrue(JumpOffset(0)));
         self.chunk().push_op(OpCode::OpPop);
-        self.compile(None, node.rhs().unwrap());
+        self.compile(slot, node.rhs().unwrap());
         self.patch_jump(end_idx);
         self.chunk().push_op(OpCode::OpAssertBool);
     }
 
     fn compile_has_attr(&mut self, slot: Option<LocalIdx>, node: ast::HasAttr) {
         // Put the attribute set on the stack.
-        self.compile(None, node.expr().unwrap());
+        self.compile(slot, node.expr().unwrap());
 
         // Push all path fragments with an operation for fetching the
         // next nested element, for all fragments except the last one.
@@ -356,7 +356,7 @@ impl Compiler {
     fn compile_attr(&mut self, slot: Option<LocalIdx>, node: ast::Attr) {
         match node {
             ast::Attr::Dynamic(dynamic) => self.compile(slot, dynamic.expr().unwrap()),
-            ast::Attr::Str(s) => self.compile_str(s),
+            ast::Attr::Str(s) => self.compile_str(slot, s),
             ast::Attr::Ident(ident) => self.emit_literal_ident(&ident),
         }
     }
@@ -367,12 +367,12 @@ impl Compiler {
     //
     // The VM, after evaluating the code for each element, simply
     // constructs the list from the given number of elements.
-    fn compile_list(&mut self, node: ast::List) {
+    fn compile_list(&mut self, slot: Option<LocalIdx>, node: ast::List) {
         let mut count = 0;
 
         for item in node.items() {
             count += 1;
-            self.compile(None, item);
+            self.compile(slot, item);
         }
 
         self.chunk().push_op(OpCode::OpList(Count(count)));
@@ -490,7 +490,7 @@ impl Compiler {
         }
 
         // Push the set onto the stack
-        self.compile(None, set);
+        self.compile(slot, set);
 
         // Compile each key fragment and emit access instructions.
         //
@@ -538,7 +538,7 @@ impl Compiler {
         path: ast::Attrpath,
         default: ast::Expr,
     ) {
-        self.compile(None, set);
+        self.compile(slot, set);
         let mut jumps = vec![];
 
         for fragment in path.attrs() {
@@ -562,15 +562,15 @@ impl Compiler {
         self.patch_jump(final_jump);
     }
 
-    fn compile_assert(&mut self, node: ast::Assert) {
+    fn compile_assert(&mut self, slot: Option<LocalIdx>, node: ast::Assert) {
         // Compile the assertion condition to leave its value on the stack.
-        self.compile(None, node.condition().unwrap());
+        self.compile(slot, node.condition().unwrap());
         self.chunk().push_op(OpCode::OpAssert);
 
         // The runtime will abort evaluation at this point if the
         // assertion failed, if not the body simply continues on like
         // normal.
-        self.compile(None, node.body().unwrap());
+        self.compile(slot, node.body().unwrap());
     }
 
     // Compile conditional expressions using jumping instructions in the VM.
@@ -583,25 +583,29 @@ impl Compiler {
     //  Jump over else body  ││ 4  [  else body  ]←┼─┘
     //  if condition is true.└┼─5─→     ...        │
     //                        └────────────────────┘
-    fn compile_if_else(&mut self, node: ast::IfElse) {
-        self.compile(None, node.condition().unwrap());
+    fn compile_if_else(&mut self, slot: Option<LocalIdx>, node: ast::IfElse) {
+        self.compile(slot, node.condition().unwrap());
 
         let then_idx = self.chunk().push_op(OpCode::OpJumpIfFalse(JumpOffset(0)));
 
         self.chunk().push_op(OpCode::OpPop); // discard condition value
-        self.compile(None, node.body().unwrap());
+        self.compile(slot, node.body().unwrap());
 
         let else_idx = self.chunk().push_op(OpCode::OpJump(JumpOffset(0)));
 
         self.patch_jump(then_idx); // patch jump *to* else_body
         self.chunk().push_op(OpCode::OpPop); // discard condition value
-        self.compile(None, node.else_body().unwrap());
+        self.compile(slot, node.else_body().unwrap());
 
         self.patch_jump(else_idx); // patch jump *over* else body
     }
 
     // Compile an `inherit` node of a `let`-expression.
-    fn compile_let_inherit<I: Iterator<Item = ast::Inherit>>(&mut self, inherits: I) {
+    fn compile_let_inherit<I: Iterator<Item = ast::Inherit>>(
+        &mut self,
+        slot: Option<LocalIdx>,
+        inherits: I,
+    ) {
         for inherit in inherits {
             match inherit.from() {
                 // Within a `let` binding, inheriting from the outer
@@ -626,7 +630,7 @@ impl Compiler {
                             continue;
                         }
 
-                        self.compile_ident(None, ident.clone());
+                        self.compile_ident(slot, ident.clone());
                         let idx = self.declare_local(
                             ident.syntax().clone(),
                             ident.ident_token().unwrap().text(),
@@ -637,7 +641,7 @@ impl Compiler {
 
                 Some(from) => {
                     for ident in inherit.idents() {
-                        self.compile(None, from.expr().unwrap());
+                        self.compile(slot, from.expr().unwrap());
                         self.emit_literal_ident(&ident);
                         self.chunk().push_op(OpCode::OpAttrsSelect);
                         let idx = self.declare_local(
@@ -656,10 +660,10 @@ 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: ast::LetIn) {
+    fn compile_let_in(&mut self, slot: Option<LocalIdx>, node: ast::LetIn) {
         self.begin_scope();
 
-        self.compile_let_inherit(node.inherits());
+        self.compile_let_inherit(slot, node.inherits());
 
         // First pass to ensure that all identifiers are known;
         // required for resolving recursion.
@@ -704,7 +708,7 @@ impl Compiler {
         }
 
         // Deal with the body, then clean up the locals afterwards.
-        self.compile(None, node.body().unwrap());
+        self.compile(slot, node.body().unwrap());
         self.end_scope();
     }
 
@@ -775,12 +779,12 @@ impl Compiler {
     // 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: ast::With) {
+    fn compile_with(&mut self, slot: Option<LocalIdx>, node: ast::With) {
         self.begin_scope();
         // TODO: Detect if the namespace is just an identifier, and
         // resolve that directly (thus avoiding duplication on the
         // stack).
-        self.compile(None, node.namespace().unwrap());
+        self.compile(slot, node.namespace().unwrap());
         let local_idx = self.scope_mut().declare_phantom();
         let with_idx = self.scope().stack_index(local_idx);
 
@@ -788,7 +792,7 @@ impl Compiler {
 
         self.chunk().push_op(OpCode::OpPushWith(with_idx));
 
-        self.compile(None, node.body().unwrap());
+        self.compile(slot, node.body().unwrap());
 
         self.chunk().push_op(OpCode::OpPopWith);
         self.scope_mut().pop_with();
@@ -818,7 +822,7 @@ impl Compiler {
             }
         }
 
-        self.compile(None, node.body().unwrap());
+        self.compile(slot, node.body().unwrap());
         self.end_scope();
 
         // TODO: determine and insert enclosing name, if available.
@@ -851,13 +855,13 @@ impl Compiler {
         self.emit_upvalue_data(slot, compiled.scope.upvalues);
     }
 
-    fn compile_apply(&mut self, node: ast::Apply) {
+    fn compile_apply(&mut self, slot: Option<LocalIdx>, node: ast::Apply) {
         // To call a function, we leave its arguments on the stack,
         // followed by the function expression itself, and then emit a
         // call instruction. This way, the stack is perfectly laid out
         // to enter the function call straight away.
-        self.compile(None, node.argument().unwrap());
-        self.compile(None, node.lambda().unwrap());
+        self.compile(slot, node.argument().unwrap());
+        self.compile(slot, node.lambda().unwrap());
         self.chunk().push_op(OpCode::OpCall);
     }