feat(tvix/eval): optimise tail calls in emitted chunks r/4784

When the last instruction in a chunk is OpCall, make it an OpTailCall instead.

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

1 files changed, 20 insertions, 2 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 448999fb38..496f0aaf32 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -900,7 +900,11 @@ impl Compiler<'_, '_> {
 
         // Pop the lambda context back off, and emit the finished
         // lambda as a constant.
-        let compiled = self.contexts.pop().unwrap();
+        let mut compiled = self.contexts.pop().unwrap();
+
+        // Check if tail-call optimisation is possible and perform it.
+        optimise_tail_call(&mut compiled.lambda.chunk);
+
         let lambda = Rc::new(compiled.lambda);
         self.observer.observe_compiled_lambda(&lambda);
 
@@ -947,7 +951,8 @@ impl Compiler<'_, '_> {
         content(self, node, slot);
         self.end_scope(node);
 
-        let thunk = self.contexts.pop().unwrap();
+        let mut thunk = self.contexts.pop().unwrap();
+        optimise_tail_call(&mut thunk.lambda.chunk);
         let lambda = Rc::new(thunk.lambda);
         self.observer.observe_compiled_thunk(&lambda);
 
@@ -1304,6 +1309,19 @@ impl Compiler<'_, '_> {
     }
 }
 
+/// Perform tail-call optimisation if the last call within a
+/// compiled chunk is another call.
+fn optimise_tail_call(chunk: &mut Chunk) {
+    let last_op = chunk
+        .code
+        .last_mut()
+        .expect("compiler bug: chunk should never be empty");
+
+    if matches!(last_op, OpCode::OpCall) {
+        *last_op = OpCode::OpTailCall;
+    }
+}
+
 /// Prepare the full set of globals from additional globals supplied
 /// by the caller of the compiler, as well as the built-in globals
 /// that are always part of the language.