feat(tvix/eval): implement opcode for function calls in VM r/4583

Nix functions always have a single argument and we do not yet make
efforts to optimise this in Tvix for known multi-argument functions
being directly applied.

For this reason, the call instruction is fairly simple and just calls
out to construct a new call frame.

Note that the logic for terminating the run loop has moved to the top
of the dispatch; this is because the loop run needs to be skipped if
the call frame for the current lambda has just been dropped.

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

3 files changed, 30 insertions, 8 deletions

diff --git a/tvix/eval/src/errors.rs b/tvix/eval/src/errors.rs
index fd3fddc326ad..a2fbdf02d186 100644
--- a/tvix/eval/src/errors.rs
+++ b/tvix/eval/src/errors.rs
@@ -32,6 +32,9 @@ pub enum ErrorKind {
     // Unknown variable in dynamic scope (with, rec, ...).
     UnknownDynamicVariable(String),
 
+    // Attempt to call something that is not callable.
+    NotCallable,
+
     ParseErrors(Vec<rnix::parser::ParseError>),
 
     AssertionFailed,
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index cda3fdf26be5..39a32d3bad2b 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -75,4 +75,7 @@ pub enum OpCode {
 
     // Asserts stack top is a boolean, and true.
     OpAssert,
+
+    // Lambdas
+    OpCall,
 }
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 1b9419ce2232..ff81349a1531 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -112,11 +112,11 @@ impl VM {
         &self.stack[self.stack.len() - 1 - offset]
     }
 
-    fn call(&mut self, lambda: Lambda) {
+    fn call(&mut self, lambda: Lambda, arg_count: usize) {
         let frame = CallFrame {
             lambda,
             ip: 0,
-            stack_offset: self.stack.len(),
+            stack_offset: self.stack.len() - arg_count,
         };
 
         self.frames.push(frame);
@@ -127,6 +127,17 @@ impl VM {
         let mut tracer = Tracer::new();
 
         'dispatch: loop {
+            if self.frame().ip == self.chunk().code.len() {
+                // If this is the end of the top-level function,
+                // return, otherwise pop the call frame.
+                if self.frames.len() == 1 {
+                    return Ok(self.pop());
+                }
+
+                self.frames.pop();
+                continue;
+            }
+
             let op = self.inc_ip();
             match op {
                 OpCode::OpConstant(idx) => {
@@ -308,7 +319,8 @@ impl VM {
                 }
 
                 OpCode::OpGetLocal(local_idx) => {
-                    let value = self.stack[local_idx].clone();
+                    let idx = self.frame().stack_offset + local_idx;
+                    let value = self.stack[idx].clone();
                     self.push(value)
                 }
 
@@ -341,16 +353,20 @@ impl VM {
                         return Err(ErrorKind::AssertionFailed.into());
                     }
                 }
+
+                OpCode::OpCall => {
+                    let callable = self.pop();
+                    match callable {
+                        Value::Lambda(lambda) => self.call(lambda, 1),
+                        _ => return Err(ErrorKind::NotCallable.into()),
+                    };
+                }
             }
 
             #[cfg(feature = "disassembler")]
             {
                 tracer.trace(&op, self.frame().ip, &self.stack);
             }
-
-            if self.frame().ip == self.chunk().code.len() {
-                return Ok(self.pop());
-            }
         }
     }
 
@@ -400,6 +416,6 @@ pub fn run_lambda(lambda: Lambda) -> EvalResult<Value> {
         with_stack: vec![],
     };
 
-    vm.call(lambda);
+    vm.call(lambda, 0);
     vm.run()
 }