about summary refs log tree commit diff
diff options
context:
space:
mode:
-rw-r--r--tvix/eval/src/builtins/impure.rs3
-rw-r--r--tvix/eval/src/compiler/mod.rs4
-rw-r--r--tvix/eval/src/tests/one_offs.rs7
-rw-r--r--tvix/eval/src/value/thunk.rs230
-rw-r--r--tvix/eval/src/vm.rs3
5 files changed, 196 insertions, 51 deletions
diff --git a/tvix/eval/src/builtins/impure.rs b/tvix/eval/src/builtins/impure.rs
index 3eaebf101e..afc85d4c1b 100644
--- a/tvix/eval/src/builtins/impure.rs
+++ b/tvix/eval/src/builtins/impure.rs
@@ -3,7 +3,6 @@ use smol_str::SmolStr;
 
 use std::{
     env,
-    rc::Rc,
     time::{SystemTime, UNIX_EPOCH},
 };
 
@@ -69,7 +68,7 @@ pub fn impure_builtins() -> Vec<(&'static str, Value)> {
 
     result.push((
         "storeDir",
-        Value::Thunk(Thunk::new_suspended_native(Rc::new(
+        Value::Thunk(Thunk::new_suspended_native(Box::new(
             |vm: &mut VM| match vm.io().store_dir() {
                 None => Ok(Value::Null),
                 Some(dir) => Ok(Value::String(dir.into())),
diff --git a/tvix/eval/src/compiler/mod.rs b/tvix/eval/src/compiler/mod.rs
index ed0daa0264..9e648fb998 100644
--- a/tvix/eval/src/compiler/mod.rs
+++ b/tvix/eval/src/compiler/mod.rs
@@ -1252,7 +1252,7 @@ fn compile_src_builtin(
     let file = source.add_file(format!("<src-builtins/{}.nix>", name), code.to_string());
     let weak = weak.clone();
 
-    Value::Thunk(Thunk::new_suspended_native(Rc::new(move |_| {
+    Value::Thunk(Thunk::new_suspended_native(Box::new(move |_| {
         let result = compile(
             &parsed.tree().expr().unwrap(),
             None,
@@ -1313,7 +1313,7 @@ pub fn prepare_globals(
         let weak_globals = weak.clone();
         builtins.insert(
             "builtins",
-            Value::Thunk(Thunk::new_suspended_native(Rc::new(move |_| {
+            Value::Thunk(Thunk::new_suspended_native(Box::new(move |_| {
                 Ok(weak_globals
                     .upgrade()
                     .unwrap()
diff --git a/tvix/eval/src/tests/one_offs.rs b/tvix/eval/src/tests/one_offs.rs
index 63bb8f7af3..13b87267a9 100644
--- a/tvix/eval/src/tests/one_offs.rs
+++ b/tvix/eval/src/tests/one_offs.rs
@@ -15,5 +15,10 @@ fn test_source_builtin() {
         result.errors
     );
 
-    assert!(matches!(result.value.unwrap(), Value::Integer(42)));
+    let value = result.value.unwrap();
+    assert!(
+        matches!(value, Value::Integer(42)),
+        "expected the integer 42, but got {}",
+        value,
+    );
 }
diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 2d48550dad..c9479fde37 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -39,8 +39,7 @@ use crate::{
 use super::{Lambda, TotalDisplay};
 
 /// Internal representation of a suspended native thunk.
-#[derive(Clone)]
-struct SuspendedNative(Rc<dyn Fn(&mut VM) -> Result<Value, ErrorKind>>);
+struct SuspendedNative(Box<dyn Fn(&mut VM) -> Result<Value, ErrorKind>>);
 
 impl Debug for SuspendedNative {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -53,7 +52,7 @@ impl Debug for SuspendedNative {
 /// Upvalues must be finalised before leaving the initial state
 /// (Suspended or RecursiveClosure).  The [`value()`] function may
 /// not be called until the thunk is in the final state (Evaluated).
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 enum ThunkRepr {
     /// Thunk is closed over some values, suspended and awaiting
     /// execution.
@@ -99,7 +98,7 @@ impl Thunk {
         })))
     }
 
-    pub fn new_suspended_native(native: Rc<dyn Fn(&mut VM) -> Result<Value, ErrorKind>>) -> Self {
+    pub fn new_suspended_native(native: Box<dyn Fn(&mut VM) -> Result<Value, ErrorKind>>) -> Self {
         Thunk(Rc::new(RefCell::new(ThunkRepr::Native(SuspendedNative(
             native,
         )))))
@@ -112,8 +111,8 @@ impl Thunk {
         if self.is_forced() {
             return Ok(());
         }
-        vm.push(Value::Thunk(self.clone()));
-        let mut trampoline = Self::force_trampoline(vm)?;
+
+        let mut trampoline = Self::force_trampoline(vm, Value::Thunk(self.clone()))?;
         loop {
             match trampoline.action {
                 None => (),
@@ -139,15 +138,11 @@ impl Thunk {
     /// Evaluate the content of a thunk, potentially repeatedly, until a
     /// non-thunk value is returned.
     ///
-    /// This will change the existing thunk (and thus all references to it,
-    /// providing memoization) through interior mutability. In case of nested
-    /// thunks, the intermediate thunk representations are replaced.
-    ///
-    /// The thunk to be forced should be at the top of the VM stack,
-    /// and will be left there (but possibly partially forced) when
-    /// this function returns.
-    pub fn force_trampoline(vm: &mut VM) -> Result<Trampoline, ErrorKind> {
-        match vm.pop() {
+    /// When this function returns, the result of one "round" of forcing is left
+    /// at the top of the stack. This may still be a partially evaluated thunk
+    /// which must be further run through the trampoline.
+    pub fn force_trampoline(vm: &mut VM, outer: Value) -> Result<Trampoline, ErrorKind> {
+        match outer {
             Value::Thunk(thunk) => thunk.force_trampoline_self(vm),
             v => {
                 vm.push(v);
@@ -156,59 +151,204 @@ impl Thunk {
         }
     }
 
+    /// Analyses `self` and, upon finding a suspended thunk, requests evaluation
+    /// of the contained code from the VM. Control flow may pass back and forth
+    /// between this function and the VM multiple times through continuations
+    /// that call `force_trampoline` again if nested thunks are encountered.
+    ///
+    /// This function is entered again by returning a continuation that calls
+    /// [force_trampoline].
+    // When working on this function, care should be taken to ensure that each
+    // evaluated thunk's *own copy* of its inner representation is replaced by
+    // evaluated results and blackholes, as appropriate. It is a critical error
+    // to move the representation of one thunk into another and can lead to
+    // hard-to-debug performance issues.
+    // TODO: check Rc count when replacing inner repr, to skip it optionally
     fn force_trampoline_self(&self, vm: &mut VM) -> Result<Trampoline, ErrorKind> {
-        loop {
-            if !self.is_suspended() {
-                // thunk is already evaluated (or infinitely
-                // recursive), inspect the situation and replace the
-                // inner representation if this is a nested thunk
-                let thunk = self.0.borrow();
-                match *thunk {
-                    ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => {
-                        let inner_repr = inner_thunk.0.borrow().clone();
-                        drop(thunk);
-                        self.0.replace(inner_repr);
-                    }
+        // If the current thunk is already fully evaluated, leave its evaluated
+        // value on the stack and return an empty trampoline. The VM will
+        // continue running the code that landed us here.
+        if self.is_forced() {
+            vm.push(self.value().clone());
+            return Ok(Trampoline::default());
+        }
 
-                    ThunkRepr::Evaluated(ref v) => {
-                        vm.push(v.clone());
-                        return Ok(Trampoline::default());
-                    }
+        // Begin evaluation of this thunk by marking it as a blackhole, meaning
+        // that any other trampoline loop round encountering this thunk before
+        // its evaluation is completed detected an evaluation cycle.
+        let inner = self.0.replace(ThunkRepr::Blackhole);
+
+        match inner {
+            // If there was already a blackhole in the thunk, this is an
+            // evaluation cycle.
+            ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion),
+
+            // If there is a native function stored in the thunk, evaluate it
+            // and replace this thunk's representation with it. Then bounces off
+            // the trampoline, to handle the case of the native function
+            // returning another thunk.
+            ThunkRepr::Native(native) => {
+                let value = native.0(vm)?;
+                self.0.replace(ThunkRepr::Evaluated(value.clone()));
+                let self_clone = self.clone();
+
+                return Ok(Trampoline {
+                    action: None,
+                    continuation: Some(Box::new(move |vm| {
+                        Thunk::force_trampoline(vm, Value::Thunk(self_clone))
+                            .map_err(|kind| Error::new(kind, todo!("BUG: b/238")))
+                    })),
+                });
+            }
+
+            // When encountering a suspended thunk, construct a trampoline that
+            // enters the thunk's code in the VM and replaces the thunks
+            // representation with the evaluated one upon return.
+            //
+            // Thunks may be nested, so this case initiates another round of
+            // trampolining to ensure that the returned value is forced.
+            ThunkRepr::Suspended {
+                lambda,
+                upvalues,
+                light_span,
+            } => {
+                // Clone self to move an Rc pointing to *this* thunk instance
+                // into the continuation closure.
+                let self_clone = self.clone();
+
+                return Ok(Trampoline {
+                    // Ask VM to enter frame of this thunk ...
+                    action: Some(TrampolineAction::EnterFrame {
+                        lambda,
+                        upvalues,
+                        arg_count: 0,
+                        light_span: light_span.clone(),
+                    }),
+
+                    // ... and replace the inner representation once that is done,
+                    // looping back around to here.
+                    continuation: Some(Box::new(move |vm: &mut VM| {
+                        let should_be_blackhole =
+                            self_clone.0.replace(ThunkRepr::Evaluated(vm.pop()));
+                        debug_assert!(matches!(should_be_blackhole, ThunkRepr::Blackhole));
+
+                        Thunk::force_trampoline(vm, Value::Thunk(self_clone))
+                            .map_err(|kind| Error::new(kind, light_span.span()))
+                    })),
+                });
+            }
+
+            // Note by tazjin: I have decided at this point to fully unroll the inner thunk handling
+            // here, leaving no room for confusion about how inner thunks are handled. This *could*
+            // be written in a shorter way (for example by using a helper function that handles all
+            // cases in which inner thunks can trivially be turned into a value), but given that we
+            // have been bitten by this logic repeatedly, I think it is better to let it be slightly
+            // verbose for now.
+
+            // If an inner thunk is found and already fully-forced, we can
+            // short-circuit and replace the representation of self with it.
+            ThunkRepr::Evaluated(Value::Thunk(ref inner)) if inner.is_forced() => {
+                self.0.replace(ThunkRepr::Evaluated(inner.value().clone()));
+                vm.push(inner.value().clone());
+                return Ok(Trampoline::default());
+            }
+
+            // Otherwise we handle inner thunks mostly as above, with the
+            // primary difference that we set the representations of *both*
+            // thunks in this case.
+            ThunkRepr::Evaluated(Value::Thunk(ref inner)) => {
+                // The inner thunk is now under evaluation, mark it as such.
+                let inner_repr = inner.0.replace(ThunkRepr::Blackhole);
+
+                match inner_repr {
                     ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion),
-                    _ => unreachable!(),
-                }
-            } else {
-                match self.0.replace(ThunkRepr::Blackhole) {
+
+                    // Same as for the native case above, but results are placed
+                    // in *both* thunks.
+                    ThunkRepr::Native(native) => {
+                        let value = native.0(vm)?;
+                        self.0.replace(ThunkRepr::Evaluated(value.clone()));
+                        inner.0.replace(ThunkRepr::Evaluated(value.clone()));
+                        let self_clone = self.clone();
+
+                        return Ok(Trampoline {
+                            action: None,
+                            continuation: Some(Box::new(move |vm| {
+                                Thunk::force_trampoline(vm, Value::Thunk(self_clone))
+                                    .map_err(|kind| Error::new(kind, todo!("BUG: b/238")))
+                            })),
+                        });
+                    }
+
+                    // Inner suspended thunks are trampolined to the VM, and
+                    // their results written to both thunks in the continuation.
                     ThunkRepr::Suspended {
                         lambda,
                         upvalues,
                         light_span,
                     } => {
                         let self_clone = self.clone();
+                        let inner_clone = inner.clone();
+
                         return Ok(Trampoline {
+                            // Ask VM to enter frame of this thunk ...
                             action: Some(TrampolineAction::EnterFrame {
                                 lambda,
                                 upvalues,
                                 arg_count: 0,
                                 light_span: light_span.clone(),
                             }),
-                            continuation: Some(Box::new(move |vm| {
-                                let should_be_blackhole =
-                                    self_clone.0.replace(ThunkRepr::Evaluated(vm.pop()));
-                                assert!(matches!(should_be_blackhole, ThunkRepr::Blackhole));
-                                vm.push(Value::Thunk(self_clone));
-                                Self::force_trampoline(vm)
+
+                            // ... and replace the inner representations.
+                            continuation: Some(Box::new(move |vm: &mut VM| {
+                                let result = vm.pop();
+
+                                let self_blackhole =
+                                    self_clone.0.replace(ThunkRepr::Evaluated(result.clone()));
+                                debug_assert!(matches!(self_blackhole, ThunkRepr::Blackhole));
+
+                                let inner_blackhole =
+                                    inner_clone.0.replace(ThunkRepr::Evaluated(result.clone()));
+                                debug_assert!(matches!(inner_blackhole, ThunkRepr::Blackhole));
+
+                                Thunk::force_trampoline(vm, Value::Thunk(self_clone))
                                     .map_err(|kind| Error::new(kind, light_span.span()))
                             })),
                         });
                     }
-                    ThunkRepr::Native(native) => {
-                        let value = native.0(vm)?;
-                        self.0.replace(ThunkRepr::Evaluated(value.clone()));
+
+                    // If the inner thunk is some arbitrary other value (this is
+                    // almost guaranteed to be another thunk), change our
+                    // representation to the same inner thunk and bounce off the
+                    // trampoline. The inner thunk is changed *back* to the same
+                    // state.
+                    //
+                    // This is safe because we are not cloning the innermost
+                    // thunk's representation, so while the inner thunk will not
+                    // eventually have its representation replaced by _this_
+                    // trampoline run, we will return the correct representation
+                    // out of here and memoize the innermost thunk.
+                    ThunkRepr::Evaluated(v) => {
+                        self.0.replace(ThunkRepr::Evaluated(v.clone()));
+                        inner.0.replace(ThunkRepr::Evaluated(v.clone()));
+                        let self_clone = self.clone();
+
+                        return Ok(Trampoline {
+                            action: None,
+                            continuation: Some(Box::new(move |vm: &mut VM| {
+                                // TODO(tazjin): not sure about this span ...
+                                // let span = vm.current_span();
+                                Thunk::force_trampoline(vm, Value::Thunk(self_clone))
+                                    .map_err(|kind| Error::new(kind, todo!("BUG: b/238")))
+                            })),
+                        });
                     }
-                    _ => unreachable!(),
                 }
             }
+
+            // This branch can not occur here, it would have been caught by our
+            // `self.is_forced()` check above.
+            ThunkRepr::Evaluated(_) => unreachable!("BUG: definition of Thunk::is_forced changed"),
         }
     }
 
diff --git a/tvix/eval/src/vm.rs b/tvix/eval/src/vm.rs
index 46c9f79b2e..cd6fc211b4 100644
--- a/tvix/eval/src/vm.rs
+++ b/tvix/eval/src/vm.rs
@@ -1017,7 +1017,8 @@ impl<'o> VM<'o> {
 
             OpCode::OpForce => {
                 if let Some(Value::Thunk(_)) = self.stack.last() {
-                    let trampoline = fallible!(self, Thunk::force_trampoline(self));
+                    let value = self.pop();
+                    let trampoline = fallible!(self, Thunk::force_trampoline(self, value));
                     return Ok(trampoline);
                 }
             }