fix(tvix/eval): ensure all evaluated thunks are correctly memoized r/5828

This fixes a very complicated bug (b/246). Evaluation
progresses *much* further after this, leading to several less
complicated bugs likely being uncovered by this

What was the problem?
=====================

Previously, when evaluating a thunk, we had a code path that looked
like this:

    match *thunk {
        ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => {
            let inner_repr = inner_thunk.0.borrow().clone();
            drop(thunk);
            self.0.replace(inner_repr);
        }
        /* ... */
    }

This code path created a copy of the inner `ThunkRepr` of a nested
thunk, and moved that copy into the `ThunkRepr` of the parent.

The effect of this was that the original `ThunkRepr` (unforced!) lived
on in the original thunk, without the memoization of the subsequent
forcing applying to it.

This had the result that Tvix would repeatedly evaluate these thunks
without ever memoizing them, if they occured repeatedly as shared
inner thunks. Most notably, this would *always* occur when
builtins.import was used.

What's the solution?
====================

I have completely rewritten `Thunk::force_trampoline_self` to make all
flows that can occur in it explicit. I have also removed the outer
loop inside of that function, and resorted to more use of trampolining
instead.

The function is now well-commented and it should be possible to read
it from top-to-bottom and get a general sense of what is going on,
though the trampolining itself (which is implemented in the VM) needs
to be at least partially understood for this.

What's the new problem(s)?
==========================

One new (known) problem is that we have to construct `Error` instances
in all error types here, but we do not have spans available in some
thunk-related situations. Due to b/238 we cannot ask the VM for an
arbitrary span from the callsite leading to the force. This means that
there are now code paths where, under certain conditions, causing an
evaluation error during thunk forcing will panic.

To fix this we will need to investigate and fix b/238, and/or add a
span tracking mechanism to thunks themselves.

What other impacts does this have?
==================================

With this commit, eval of nixpkgs mostly succeeds (things like stdenv
evaluate to the same hashes for us and C++ Nix, meaning we now
construct identical derivations without eval breaking).

Due to this we progress much further into nixpkgs, which lets us
uncover more additional bugs. For example, after this commit we can
quickly see that cl/7949 introduces some kind of behavioural issue and
should not be merged as-is (this was not apparent before).

Additionally, tvix-eval is now seemingly very fast. When doing
performance analysis of a nixpkgs eval, we now mostly see the code
path for shelling out to C++ Nix to add things to the store in there.
We still need those code paths, so we can not (yet) do a performance
analysis beyond that.

Change-Id: I738525bad8bc5ede5d8c737f023b14b8f4160612
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8012
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

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);
                 }
             }