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

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>

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::Str(s) => self.compile_str(slot, s),
-            ast::Expr::UnaryOp(op) => self.compile_unary_op(op),
+            ast::Expr::UnaryOp(op) => self.compile_unary_op(slot, op),
-            ast::Expr::BinOp(op) => self.compile_binop(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::Assert(assert) => self.compile_assert(slot, assert),
-            ast::Expr::IfElse(if_else) => self.compile_if_else(if_else),
+            ast::Expr::IfElse(if_else) => self.compile_if_else(slot, if_else),
-            ast::Expr::LetIn(let_in) => self.compile_let_in(let_in),
+            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) {
+    fn compile_unary_op(&mut self, slot: Option<LocalIdx>, op: ast::UnaryOp) {
-        self.compile(None, op.expr().unwrap());
+        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::And => return self.compile_and(slot, op),
-            BinOpKind::Or => return self.compile_or(op),
+            BinOpKind::Or => return self.compile_or(slot, op),
-            BinOpKind::Implication => return self.compile_implication(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(slot, op.lhs().unwrap());
-        self.compile(None, op.rhs().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) {
+    fn compile_if_else(&mut self, slot: Option<LocalIdx>, node: ast::IfElse) {
-        self.compile(None, node.condition().unwrap());
+        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(slot, node.argument().unwrap());
-        self.compile(None, node.lambda().unwrap());
+        self.compile(slot, node.lambda().unwrap());
         self.chunk().push_op(OpCode::OpCall);
     }