feat(tvix/eval): Add Compiler::thunk method for emitting thunks r/4674

The logic in this method is *very* similar to `compile_lambda`. It is
intended to be called around any expression that should be
thunked (such as function applications, attribute set values, etc.).

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

3 files changed, 41 insertions, 1 deletions

diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index c9cd66aca78f..7c20307b843b 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -26,7 +26,7 @@ use std::rc::Rc;
 use crate::chunk::Chunk;
 use crate::errors::{Error, ErrorKind, EvalResult};
 use crate::opcode::{CodeIdx, Count, JumpOffset, OpCode, UpvalueIdx};
-use crate::value::{Closure, Lambda, Value};
+use crate::value::{Closure, Lambda, Thunk, Value};
 use crate::warnings::{EvalWarning, WarningKind};
 
 use self::scope::{Local, LocalIdx, LocalPosition, Scope, Upvalue};
@@ -848,6 +848,41 @@ impl Compiler {
         self.chunk().push_op(OpCode::OpCall);
     }
 
+    /// Compile an expression into a runtime thunk which should be
+    /// lazily evaluated when accessed.
+    // TODO: almost the same as Compiler::compile_lambda; unify?
+    fn thunk<F>(&mut self, slot: Option<LocalIdx>, content: F)
+    where
+        F: FnOnce(&mut Compiler, Option<LocalIdx>),
+    {
+        self.contexts.push(LambdaCtx::new());
+        self.begin_scope();
+        content(self, slot);
+        self.end_scope();
+
+        let thunk = self.contexts.pop().unwrap();
+
+        #[cfg(feature = "disassembler")]
+        {
+            crate::disassembler::disassemble_chunk(&thunk.lambda.chunk);
+        }
+
+        // Emit the thunk directly if it does not close over the
+        // environment.
+        if thunk.lambda.upvalue_count == 0 {
+            self.emit_constant(Value::Thunk(Thunk::new(thunk.lambda)));
+            return;
+        }
+
+        // Otherwise prepare for runtime construction of the thunk.
+        let blueprint_idx = self
+            .chunk()
+            .push_constant(Value::Blueprint(Rc::new(thunk.lambda)));
+
+        self.chunk().push_op(OpCode::OpThunk(blueprint_idx));
+        self.emit_upvalue_data(slot, thunk.scope.upvalues);
+    }
+
     /// Emit the data instructions that the runtime needs to correctly
     /// assemble the provided upvalues array.
     fn emit_upvalue_data(&mut self, slot: Option<LocalIdx>, upvalues: Vec<Upvalue>) {
diff --git a/tvix/eval/src/opcode.rs b/tvix/eval/src/opcode.rs
index 45a6337229d2..a8802e8c0a2b 100644
--- a/tvix/eval/src/opcode.rs
+++ b/tvix/eval/src/opcode.rs
@@ -107,6 +107,9 @@ pub enum OpCode {
     OpGetUpvalue(UpvalueIdx),
     OpClosure(ConstantIdx),
 
+    // Thunks
+    OpThunk(ConstantIdx),
+
     /// Finalise initialisation of the upvalues of the value in the
     /// given stack index after the scope is fully bound.
     OpFinalise(StackIdx),
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 1aaba9b0573e..b363522994eb 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -435,6 +435,8 @@ impl VM {
                     }
                 }
 
+                OpCode::OpThunk(_idx) => todo!("runtime thunk construction"),
+
                 OpCode::OpFinalise(StackIdx(idx)) => {
                     match &self.stack[self.frame().stack_offset + idx] {
                         Value::Closure(closure) => closure