refactor(tvix/eval): return call frame result from VM::call r/4748

Previously, "calling" (setting up the VM run loop for executing a call
frame) and "running" (running this loop to completion) were separate
operations.

This was basically an attempt to avoid nesting `VM::run` invocations.
However, doing things this way introduced some tricky bugs for exiting
out of the call frames of thunks vs. builtins & closures.

For now, we unify the two operations and always return the value to
the caller directly. For now this makes calls a little less effective,
but it gives us a chance to nail down some other strange behaviours
and then re-optimise this afterwards.

To make sure we tackle this again further down I've added it to the
list of known possible optimisations.

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

1 files changed, 20 insertions, 12 deletions

diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 8d616b8d73b8..fecaae37aa5b 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -161,7 +161,14 @@ impl VM {
         }
     }
 
-    pub fn call(&mut self, lambda: Rc<Lambda>, upvalues: Vec<Value>, arg_count: usize) {
+    /// Execute the given lambda in this VM's context, returning its
+    /// value after its stack frame completes.
+    pub fn call(
+        &mut self,
+        lambda: Rc<Lambda>,
+        upvalues: Vec<Value>,
+        arg_count: usize,
+    ) -> EvalResult<Value> {
         let frame = CallFrame {
             lambda,
             upvalues,
@@ -170,22 +177,22 @@ impl VM {
         };
 
         self.frames.push(frame);
+        self.run()
     }
 
-    pub fn run(&mut self) -> EvalResult<Value> {
+    /// Run the VM's current stack frame to completion and return the
+    /// value.
+    fn run(&mut self) -> EvalResult<Value> {
         #[cfg(feature = "disassembler")]
         let mut tracer = Tracer::new();
 
         loop {
+            // Break the loop if this call frame has already run to
+            // completion, pop it off, and return the value to the
+            // caller.
             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;
+                return Ok(self.pop());
             }
 
             let op = self.inc_ip();
@@ -413,7 +420,9 @@ impl VM {
                     let callable = self.pop();
                     match callable {
                         Value::Closure(closure) => {
-                            self.call(closure.lambda(), closure.upvalues().to_vec(), 1)
+                            let result =
+                                self.call(closure.lambda(), closure.upvalues().to_vec(), 1)?;
+                            self.push(result)
                         }
 
                         Value::Builtin(builtin) => {
@@ -684,8 +693,7 @@ pub fn run_lambda(lambda: Lambda) -> EvalResult<Value> {
         with_stack: vec![],
     };
 
-    vm.call(Rc::new(lambda), vec![], 0);
-    let value = vm.run()?;
+    let value = vm.call(Rc::new(lambda), vec![], 0)?;
     vm.force_for_output(&value)?;
     Ok(value)
 }