feat(tvix/eval): compile simple `let ... in ...` expressions r/4522

These expressions now leave the binding values on the stack, and clean
up the scope after the body of the expression.

While variable access is not yet implemented (as the identifier node
remains unhandled), this already gives us the correct stack behaviour.

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

3 files changed, 98 insertions, 0 deletions

diff --git a/tvix/eval/src/compiler.rs b/tvix/eval/src/compiler.rs
index a0c62b112869..cd98531c7831 100644
--- a/tvix/eval/src/compiler.rs
+++ b/tvix/eval/src/compiler.rs
@@ -49,6 +49,7 @@ struct Local {
 /// TODO(tazjin): `with`-stack
 /// TODO(tazjin): flag "specials" (e.g. note depth if builtins are
 /// overridden)
+#[derive(Default)]
 struct Locals {
     locals: Vec<Local>,
 
@@ -58,6 +59,8 @@ struct Locals {
 
 struct Compiler {
     chunk: Chunk,
+    locals: Locals,
+
     warnings: Vec<EvalWarning>,
     root_dir: PathBuf,
 }
@@ -133,6 +136,11 @@ impl Compiler {
                 self.compile_if_else(node)
             }
 
+            rnix::SyntaxKind::NODE_LET_IN => {
+                let node = rnix::types::LetIn::cast(node).unwrap();
+                self.compile_let_in(node)
+            }
+
             kind => panic!("visiting unsupported node: {:?}", kind),
         }
     }
@@ -633,6 +641,50 @@ impl Compiler {
         Ok(())
     }
 
+    // Compile a standard `let ...; in ...` statement.
+    //
+    // Unless in a non-standard scope, the encountered values are
+    // simply pushed on the stack and their indices noted in the
+    fn compile_let_in(&mut self, node: rnix::types::LetIn) -> Result<(), Error> {
+        self.begin_scope();
+        let mut entries = vec![];
+
+        // Before compiling the values of a let expression, all keys
+        // need to already be added to the known locals. This is
+        // because in Nix these bindings are always recursive (they
+        // can even refer to themselves).
+        for entry in node.entries() {
+            let key = entry.key().unwrap();
+            let path = key.path().collect::<Vec<_>>();
+
+            if path.len() != 1 {
+                todo!("nested bindings in let expressions :(")
+            }
+
+            entries.push(entry.value().unwrap());
+
+            self.locals.locals.push(Local {
+                name: key.node().clone(), // TODO(tazjin): Just an Rc?
+                depth: self.locals.scope_depth,
+            });
+        }
+
+        for _ in node.inherits() {
+            todo!("inherit in let not yet implemented")
+        }
+
+        // Now we can compile each expression, leaving the values on
+        // the stack in the right order.
+        for value in entries {
+            self.compile(value)?;
+        }
+
+        // Deal with the body, then clean up the locals afterwards.
+        self.compile(node.body().unwrap())?;
+        self.end_scope();
+        Ok(())
+    }
+
     fn patch_jump(&mut self, idx: CodeIdx) {
         let offset = self.chunk.code.len() - 1 - idx.0;
 
@@ -647,6 +699,34 @@ impl Compiler {
             op => panic!("attempted to patch unsupported op: {:?}", op),
         }
     }
+
+    fn begin_scope(&mut self) {
+        self.locals.scope_depth += 1;
+    }
+
+    fn end_scope(&mut self) {
+        let mut scope = &mut self.locals;
+        debug_assert!(scope.scope_depth != 0, "can not end top scope");
+        scope.scope_depth -= 1;
+
+        // When ending a scope, all corresponding locals need to be
+        // removed, but the value of the body needs to remain on the
+        // stack. This is implemented by a separate instruction.
+        let mut pops = 0;
+
+        // TL;DR - iterate from the back while things belonging to the
+        // ended scope still exist.
+        while scope.locals.len() > 0
+            && scope.locals[scope.locals.len() - 1].depth > scope.scope_depth
+        {
+            pops += 1;
+            scope.locals.pop();
+        }
+
+        if pops > 0 {
+            self.chunk.push_op(OpCode::OpCloseScope(pops));
+        }
+    }
 }
 
 pub fn compile(ast: rnix::AST, location: Option<PathBuf>) -> EvalResult<CompilationResult> {
@@ -668,6 +748,7 @@ pub fn compile(ast: rnix::AST, location: Option<PathBuf>) -> EvalResult<Compilat
         root_dir,
         chunk: Chunk::default(),
         warnings: vec![],
+        locals: Default::default(),
     };
 
     c.compile(ast.node())?;
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index 4cf4f695b0c7..ebd91dd43924 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -61,4 +61,7 @@ pub enum OpCode {
 
     // Type assertion operators
     OpAssertBool,
+
+    // Close scopes while leaving their expression value around.
+    OpCloseScope(usize), // number of locals to pop
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 232e27aabbd3..eb860eae007e 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -240,6 +240,20 @@ impl VM {
                         });
                     }
                 }
+
+                // Remove the given number of elements from the stack,
+                // but retain the top value.
+                OpCode::OpCloseScope(count) => {
+                    // Immediately move the top value into the right
+                    // position.
+                    let target_idx = self.stack.len() - 1 - count;
+                    self.stack[target_idx] = self.pop();
+
+                    // Then drop the remaining values.
+                    for _ in 0..(count - 1) {
+                        self.pop();
+                    }
+                }
             }
 
             if self.ip == self.chunk.code.len() {