feat(tvix/eval): implement unthunking in compiler r/6204

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>

1 files changed, 22 insertions, 1 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index 7f090978db99..bfb823b673c5 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