From 2e018a50a74040d8db5c0dfeae5f829a5c7c0cf2 Mon Sep 17 00:00:00 2001
From: Vincent Ambo <mail@tazj.in>
Date: Sun, 4 Sep 2022 23:16:59 +0300
Subject: feat(tvix/eval): implement OpTailCall

If the last operation within a chunk is a function call, the call can
be executed in the same call frame without increasing the depth of the
call stack.

To enable this, a new OpTailCall instruction (similar to OpCall) is
introduced, but not yet emitted by the compiler.

Change-Id: I9ffbd7da6d2d6a8ec7a724646435dc6ee89712f2
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6457
Tested-by: BuildkiteCI
Reviewed-by: sterni <sternenseemann@systemli.org>
---
 tvix/eval/src/observer.rs | 12 ++++++++++++
 tvix/eval/src/opcode.rs   |  1 +
 tvix/eval/src/vm.rs       | 46 +++++++++++++++++++++++++++++++++++++---------
 3 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/tvix/eval/src/observer.rs b/tvix/eval/src/observer.rs
index 677c3f0811e7..8dc2a6a0ca96 100644
--- a/tvix/eval/src/observer.rs
+++ b/tvix/eval/src/observer.rs
@@ -41,6 +41,10 @@ pub trait Observer {
     /// Called when the runtime exits a call frame.
     fn observe_exit_frame(&mut self, _frame_at: usize) {}
 
+    /// Called when the runtime replaces the current call frame for a
+    /// tail call.
+    fn observe_tail_call(&mut self, _frame_at: usize, _: &Rc<Lambda>) {}
+
     /// Called when the runtime enters a builtin.
     fn observe_enter_builtin(&mut self, _name: &'static str) {}
 
@@ -150,6 +154,14 @@ impl<W: Write> Observer for TracingObserver<W> {
         let _ = writeln!(&mut self.writer, "=== exiting builtin {} ===", name);
     }
 
+    fn observe_tail_call(&mut self, frame_at: usize, lambda: &Rc<Lambda>) {
+        let _ = writeln!(
+            &mut self.writer,
+            "=== tail-calling {:p} in frame[{}] ===",
+            lambda, frame_at
+        );
+    }
+
     fn observe_execute_op(&mut self, ip: usize, op: &OpCode, stack: &[Value]) {
         let _ = write!(&mut self.writer, "{:04} {:?}\t[ ", ip, op);
 
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index 871111883ba4..565c44a572de 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -104,6 +104,7 @@ pub enum OpCode {
 
     // Lambdas & closures
     OpCall,
+    OpTailCall,
     OpGetUpvalue(UpvalueIdx),
     OpClosure(ConstantIdx),
 
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index b65a37582f43..a7b9bf3e9206 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -9,7 +9,7 @@ use crate::{
     observer::Observer,
     opcode::{CodeIdx, ConstantIdx, Count, JumpOffset, OpCode, StackIdx, UpvalueIdx},
     upvalues::UpvalueCarrier,
-    value::{Closure, Lambda, NixAttrs, NixList, Thunk, Value},
+    value::{Builtin, Closure, Lambda, NixAttrs, NixList, Thunk, Value},
 };
 
 struct CallFrame {
@@ -447,19 +447,33 @@ impl<'o> VM<'o> {
                             self.push(result)
                         }
 
-                        Value::Builtin(builtin) => {
-                            let builtin_name = builtin.name();
-                            self.observer.observe_enter_builtin(builtin_name);
+                        Value::Builtin(builtin) => self.call_builtin(builtin)?,
 
-                            let arg = self.pop();
-                            let result = fallible!(self, builtin.apply(self, arg));
+                        _ => return Err(self.error(ErrorKind::NotCallable)),
+                    };
+                }
 
-                            self.observer.observe_exit_builtin(builtin_name);
+                OpCode::OpTailCall => {
+                    let callable = self.pop();
 
-                            self.push(result);
+                    match callable {
+                        Value::Builtin(builtin) => self.call_builtin(builtin)?,
+
+                        Value::Closure(closure) => {
+                            let lambda = closure.lambda();
+                            self.observer.observe_tail_call(self.frames.len(), &lambda);
+
+                            // Replace the current call frames
+                            // internals with that of the tail-called
+                            // closure.
+                            let mut frame = self.frame_mut();
+                            frame.lambda = lambda;
+                            frame.upvalues = closure.upvalues().to_vec();
+                            frame.ip = 0; // reset instruction pointer to beginning
                         }
+
                         _ => return Err(self.error(ErrorKind::NotCallable)),
-                    };
+                    }
                 }
 
                 OpCode::OpGetUpvalue(upv_idx) => {
@@ -699,6 +713,20 @@ impl<'o> VM<'o> {
             _ => Ok(()),
         }
     }
+
+    fn call_builtin(&mut self, builtin: Builtin) -> EvalResult<()> {
+        let builtin_name = builtin.name();
+        self.observer.observe_enter_builtin(builtin_name);
+
+        let arg = self.pop();
+        let result = fallible!(self, builtin.apply(self, arg));
+
+        self.observer.observe_exit_builtin(builtin_name);
+
+        self.push(result);
+
+        Ok(())
+    }
 }
 
 // TODO: use Rc::unwrap_or_clone once it is stabilised.
-- 
cgit 1.4.1