refactor(tvix/eval): non-recursive thunk forcing r/5485

Introduces continuation-passing-based trampolining of thunk forcing to
avoid recursing when forcing deeply nested expressions.

This is required for evaluating large expressions.

This change was extracted out of cl/7362.

Co-authored-by: Vincent Ambo <tazjin@tvl.su>
Co-authored-by: Griffin Smith <grfn@gws.fyi>
Change-Id: Ifc1747e712663684b2fff53095de62b8459a47f3
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7551
Reviewed-by: grfn <grfn@gws.fyi>
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>

1 files changed, 115 insertions, 29 deletions

diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 680cc9b1fb..c7cdfa2441 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -30,7 +30,7 @@ use crate::{
     spans::LightSpan,
     upvalues::Upvalues,
     value::{Builtin, Closure},
-    vm::VM,
+    vm::{Trampoline, TrampolineAction, VM},
     Value,
 };
 
@@ -84,10 +84,6 @@ impl Thunk {
         })))
     }
 
-    /// Create a new thunk from suspended Rust code.
-    ///
-    /// The suspended code will be executed and expected to return a
-    /// value whenever the thunk is forced like any other thunk.
     pub fn new_suspended_native(
         native: Rc<Box<dyn Fn(&mut VM) -> Result<Value, ErrorKind>>>,
     ) -> Self {
@@ -103,7 +99,7 @@ impl Thunk {
             None,
             move |v: Vec<Value>, vm: &mut VM| {
                 // sanity check that only the dummy argument was popped
-                assert_eq!(v.len(), 1);
+                assert!(v.len() == 1);
                 assert!(matches!(v[0], Value::Null));
                 native(vm)
             },
@@ -127,41 +123,103 @@ impl Thunk {
         })))
     }
 
+    /// Force a thunk from a context that can't handle trampoline
+    /// continuations, eg outside the VM's normal execution loop.  Calling
+    /// `force_trampoline()` instead should be preferred whenever possible.
+    pub fn force(&self, vm: &mut VM) -> Result<(), ErrorKind> {
+        if self.is_forced() {
+            return Ok(());
+        }
+        vm.push(Value::Thunk(self.clone()));
+        let mut trampoline = Self::force_trampoline(vm)?;
+        loop {
+            match trampoline.action {
+                None => (),
+                Some(TrampolineAction::EnterFrame {
+                    lambda,
+                    upvalues,
+                    arg_count,
+                    light_span: _,
+                }) => vm.enter_frame(lambda, upvalues, arg_count)?,
+            }
+            match trampoline.continuation {
+                None => break (),
+                Some(cont) => {
+                    trampoline = cont(vm)?;
+                    continue;
+                }
+            }
+        }
+        vm.pop();
+        Ok(())
+    }
+
     /// 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.
-    pub fn force(&self, vm: &mut VM) -> Result<(), ErrorKind> {
+    ///
+    /// 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() {
+            Value::Thunk(thunk) => thunk.force_trampoline_self(vm),
+            v => {
+                vm.push(v);
+                Ok(Trampoline::default())
+            }
+        }
+    }
+
+    fn force_trampoline_self(&self, vm: &mut VM) -> Result<Trampoline, ErrorKind> {
         loop {
-            let mut thunk_mut = self.0.borrow_mut();
+            if !self.is_suspended() {
+                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);
+                    }
 
-            match *thunk_mut {
-                ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => {
-                    let inner_repr = inner_thunk.0.borrow().clone();
-                    *thunk_mut = inner_repr;
+                    ThunkRepr::Evaluated(ref v) => {
+                        vm.push(v.clone());
+                        return Ok(Trampoline::default());
+                    }
+                    ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion),
+                    _ => panic!("impossible"),
                 }
-
-                ThunkRepr::Evaluated(_) => return Ok(()),
-                ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion),
-
-                ThunkRepr::Suspended { .. } => {
-                    if let ThunkRepr::Suspended {
+            } else {
+                match self.0.replace(ThunkRepr::Blackhole) {
+                    ThunkRepr::Suspended {
                         lambda,
                         upvalues,
                         light_span,
-                    } = std::mem::replace(&mut *thunk_mut, ThunkRepr::Blackhole)
-                    {
-                        drop(thunk_mut);
-                        vm.enter_frame(lambda, upvalues, 0).map_err(|e| {
-                            ErrorKind::ThunkForce(Box::new(Error {
-                                span: light_span.span(),
-                                ..e
-                            }))
-                        })?;
-                        (*self.0.borrow_mut()) = ThunkRepr::Evaluated(vm.pop())
+                    } => {
+                        let self_clone = self.clone();
+                        return Ok(Trampoline {
+                            action: Some(TrampolineAction::EnterFrame {
+                                lambda,
+                                upvalues: upvalues.clone(),
+                                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));
+                                return Self::force_trampoline(vm).map_err(|kind| Error {
+                                    kind,
+                                    span: light_span.span(),
+                                });
+                            })),
+                        });
                     }
+                    _ => panic!("impossible"),
                 }
             }
         }
@@ -175,6 +233,20 @@ impl Thunk {
         matches!(*self.0.borrow(), ThunkRepr::Evaluated(_))
     }
 
+    pub fn is_suspended(&self) -> bool {
+        matches!(*self.0.borrow(), ThunkRepr::Suspended { .. })
+    }
+
+    /// Returns true if forcing this thunk will not change it.
+    pub fn is_forced(&self) -> bool {
+        match *self.0.borrow() {
+            ThunkRepr::Blackhole => panic!("is_forced() called on a blackholed thunk"),
+            ThunkRepr::Evaluated(Value::Thunk(_)) => false,
+            ThunkRepr::Evaluated(_) => true,
+            _ => false,
+        }
+    }
+
     /// Returns a reference to the inner evaluated value of a thunk.
     /// It is an error to call this on a thunk that has not been
     /// forced, or is not otherwise known to be fully evaluated.
@@ -183,7 +255,21 @@ impl Thunk {
     // API too much.
     pub fn value(&self) -> Ref<Value> {
         Ref::map(self.0.borrow(), |thunk| match thunk {
-            ThunkRepr::Evaluated(value) => value,
+            ThunkRepr::Evaluated(value) => {
+                /*
+                #[cfg(debug_assertions)]
+                if matches!(
+                    value,
+                    Value::Closure(Closure {
+                        is_finalised: false,
+                        ..
+                    })
+                ) {
+                    panic!("Thunk::value called on an unfinalised closure");
+                }
+                */
+                return value;
+            }
             ThunkRepr::Blackhole => panic!("Thunk::value called on a black-holed thunk"),
             ThunkRepr::Suspended { .. } => panic!("Thunk::value called on a suspended thunk"),
         })