From 5d3fb33baa8da6bb8cfbd078e2b295d773a721ba Mon Sep 17 00:00:00 2001
From: Vincent Ambo <mail@tazj.in>
Date: Mon, 30 Jan 2023 00:37:47 +0300
Subject: feat(tvix/eval): implement unthunking in compiler

This feature allows the compiler to detect situations where the
created thunk is useless and can be merged into the parent chunk
instead.

The only case where the compiler does this initially is when
statically optimising a select expression.

For example, previously the expression `builtins.length` compiled into
two thunks:

1. An "inner" thunk which contained an `OpConstant` that had the
   optimised `length` builtin in it.

2. An "outer" thunk which contained an `OpConstant` to access the
   inner thunk, and the trailing OpForce of the top-level program.

With this change, the inner thunk is skipped completely and the outer
chunk directly contains the `length` builtin access.

This can be applied in several situations, some easier than others,
and we will add them in as we go along.

Change-Id: Ie44521445fce1199f99b5b17712833faea9bc357
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7959
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
---
 tvix/eval/src/compiler/mod.rs | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

(limited to 'tvix/eval/src')

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 7f090978db..bfb823b673 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -57,6 +57,7 @@ struct LambdaCtx {
     lambda: Lambda,
     scope: Scope,
     captures_with_stack: bool,
+    unthunk: bool,
 }
 
 impl LambdaCtx {
@@ -65,6 +66,7 @@ impl LambdaCtx {
             lambda: Lambda::default(),
             scope: Default::default(),
             captures_with_stack: false,
+            unthunk: false,
         }
     }
 
@@ -73,6 +75,7 @@ impl LambdaCtx {
             lambda: Lambda::default(),
             scope: self.scope.inherit(),
             captures_with_stack: false,
+            unthunk: false,
         }
     }
 }
@@ -664,6 +667,10 @@ impl Compiler<'_> {
                     if let Some(ident) = expr_static_attr_str(&attr) {
                         if let Some(selected_value) = attrs.select(ident.as_str()) {
                             *constant = selected_value.clone();
+
+                            // If this worked, we can unthunk the current thunk.
+                            self.unthunk();
+
                             return true;
                         }
                     }
@@ -1003,7 +1010,13 @@ impl Compiler<'_> {
         self.compile_lambda_or_thunk(true, outer_slot, node, content)
     }
 
-    /// Compile an expression into a runtime cloure or thunk
+    /// Mark the current thunk as redundant, i.e. possible to merge directly
+    /// into its parent lambda context without affecting runtime behaviour.
+    fn unthunk(&mut self) {
+        self.context_mut().unthunk = true;
+    }
+
+    /// Compile an expression into a runtime closure or thunk
     fn compile_lambda_or_thunk<N, F>(
         &mut self,
         is_suspended_thunk: bool,
@@ -1034,6 +1047,14 @@ impl Compiler<'_> {
         // lambda as a constant.
         let mut compiled = self.contexts.pop().unwrap();
 
+        // The compiler might have decided to unthunk, i.e. raise the compiled
+        // code to the parent context. In that case we do so and return right
+        // away.
+        if compiled.unthunk && is_suspended_thunk {
+            self.chunk().extend(compiled.lambda.chunk);
+            return;
+        }
+
         // Emit an instruction to inform the VM that the chunk has ended.
         compiled
             .lambda
-- 
cgit 1.4.1