feat(tvix/eval): deduplicate overlap between Closure and Thunk r/5159

This commit deduplicates the Thunk-like functionality from Closure
and unifies it with Thunk.

Specifically, we now have one and only one way of breaking reference
cycles in the Value-graph: Thunk.  No other variant contains a
RefCell.  This should make it easier to reason about the behavior of
the VM.  InnerClosure and UpvaluesCarrier are no longer necessary.

This refactoring allowed an improvement in code generation:
`Rc<RefCell<>>`s are now created only for closures which do not have
self-references or deferred upvalues, instead of for all closures.
OpClosure has been split into two separate opcodes:

- OpClosure creates non-recursive closures with no deferred
  upvalues.  The VM will not create an `Rc<RefCell<>>` when executing
  this instruction.

- OpThunkClosure is used for closures with self-references or
  deferred upvalues.  The VM will create a Thunk when executing this
  opcode, but the Thunk will start out already in the
  `ThunkRepr::Evaluated` state, rather than in the
  `ThunkRepr::Suspeneded` state.

To avoid confusion, OpThunk has been renamed OpThunkSuspended.

Thanks to @sterni for suggesting that all this could be done without
adding an additional variant to ThunkRepr.  This does however mean
that there will be mutating accesses to `ThunkRepr::Evaluated`,
which was not previously the case.  The field `is_finalised:bool`
has been added to `Closure` to ensure that these mutating accesses
are performed only on finalised Closures.  Both the check and the
field are present only if `#[cfg(debug_assertions)]`.

Change-Id: I04131501029772f30e28da8281d864427685097f
Signed-off-by: Adam Joseph <adam@westernsemico.com>
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7019
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>

1 files changed, 76 insertions, 38 deletions

diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 887c297a53b9..818ec0f58aec 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -28,7 +28,8 @@ use codemap::Span;
 
 use crate::{
     errors::{Error, ErrorKind},
-    upvalues::{UpvalueCarrier, Upvalues},
+    upvalues::Upvalues,
+    value::Closure,
     vm::VM,
     Value,
 };
@@ -36,6 +37,10 @@ use crate::{
 use super::Lambda;
 
 /// Internal representation of the different states of a thunk.
+///
+/// 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, PartialEq)]
 enum ThunkRepr {
     /// Thunk is closed over some values, suspended and awaiting
@@ -43,6 +48,7 @@ enum ThunkRepr {
     Suspended {
         lambda: Rc<Lambda>,
         upvalues: Upvalues,
+        span: Span,
     },
 
     /// Thunk currently under-evaluation; encountering a blackhole
@@ -53,21 +59,31 @@ enum ThunkRepr {
     Evaluated(Value),
 }
 
+/// A thunk is created for any value which requires non-strict
+/// evaluation due to self-reference or lazy semantics (or both).
+/// Every reference cycle involving `Value`s will contain at least
+/// one `Thunk`.
 #[derive(Clone, Debug, PartialEq)]
-pub struct Thunk {
-    inner: Rc<RefCell<ThunkRepr>>,
-    span: Span,
-}
+pub struct Thunk(Rc<RefCell<ThunkRepr>>);
 
 impl Thunk {
-    pub fn new(lambda: Rc<Lambda>, span: Span) -> Self {
-        Thunk {
-            inner: Rc::new(RefCell::new(ThunkRepr::Suspended {
+    pub fn new_closure(lambda: Rc<Lambda>) -> Self {
+        Thunk(Rc::new(RefCell::new(ThunkRepr::Evaluated(Value::Closure(
+            Closure {
                 upvalues: Upvalues::with_capacity(lambda.upvalue_count),
                 lambda: lambda.clone(),
-            })),
+                #[cfg(debug_assertions)]
+                is_finalised: false,
+            },
+        )))))
+    }
+
+    pub fn new_suspended(lambda: Rc<Lambda>, span: Span) -> Self {
+        Thunk(Rc::new(RefCell::new(ThunkRepr::Suspended {
+            upvalues: Upvalues::with_capacity(lambda.upvalue_count),
+            lambda: lambda.clone(),
             span,
-        }
+        })))
     }
 
     /// Evaluate the content of a thunk, potentially repeatedly, until
@@ -79,11 +95,11 @@ impl Thunk {
     /// are replaced.
     pub fn force(&self, vm: &mut VM) -> Result<(), ErrorKind> {
         loop {
-            let mut thunk_mut = self.inner.borrow_mut();
+            let mut thunk_mut = self.0.borrow_mut();
 
             match *thunk_mut {
                 ThunkRepr::Evaluated(Value::Thunk(ref inner_thunk)) => {
-                    let inner_repr = inner_thunk.inner.borrow().clone();
+                    let inner_repr = inner_thunk.0.borrow().clone();
                     *thunk_mut = inner_repr;
                 }
 
@@ -91,24 +107,33 @@ impl Thunk {
                 ThunkRepr::Blackhole => return Err(ErrorKind::InfiniteRecursion),
 
                 ThunkRepr::Suspended { .. } => {
-                    if let ThunkRepr::Suspended { lambda, upvalues } =
-                        std::mem::replace(&mut *thunk_mut, ThunkRepr::Blackhole)
+                    if let ThunkRepr::Suspended {
+                        lambda,
+                        upvalues,
+                        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: self.span,
-                                ..e
-                            }))
-                        })?;
-                        let evaluated = ThunkRepr::Evaluated(vm.pop());
-                        (*self.inner.borrow_mut()) = evaluated;
+                        vm.enter_frame(lambda, upvalues, 0)
+                            .map_err(|e| ErrorKind::ThunkForce(Box::new(Error { span, ..e })))?;
+                        (*self.0.borrow_mut()) = ThunkRepr::Evaluated(vm.pop())
                     }
                 }
             }
         }
     }
 
+    pub fn finalise(&self, stack: &[Value]) {
+        self.upvalues_mut().resolve_deferred_upvalues(stack);
+        #[cfg(debug_assertions)]
+        {
+            let inner: &mut ThunkRepr = &mut self.0.as_ref().borrow_mut();
+            if let ThunkRepr::Evaluated(Value::Closure(c)) = inner {
+                c.is_finalised = true;
+            }
+        }
+    }
+
     /// 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.
@@ -116,35 +141,48 @@ impl Thunk {
     // difficult to represent in the type system without impacting the
     // API too much.
     pub fn value(&self) -> Ref<Value> {
-        Ref::map(self.inner.borrow(), |thunk| {
+        Ref::map(self.0.borrow(), |thunk| {
             if let ThunkRepr::Evaluated(value) = thunk {
+                #[cfg(debug_assertions)]
+                if matches!(
+                    value,
+                    Value::Closure(Closure {
+                        is_finalised: false,
+                        ..
+                    })
+                ) {
+                    panic!("Thunk::value called on an unfinalised closure");
+                }
                 return value;
             }
 
             panic!("Thunk::value called on non-evaluated thunk");
         })
     }
-}
-
-impl UpvalueCarrier for Thunk {
-    fn upvalue_count(&self) -> usize {
-        if let ThunkRepr::Suspended { lambda, .. } = &*self.inner.borrow() {
-            return lambda.upvalue_count;
-        }
-
-        panic!("upvalues() on non-suspended thunk");
-    }
 
-    fn upvalues(&self) -> Ref<'_, Upvalues> {
-        Ref::map(self.inner.borrow(), |thunk| match thunk {
+    pub fn upvalues(&self) -> Ref<'_, Upvalues> {
+        Ref::map(self.0.borrow(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
+            ThunkRepr::Evaluated(Value::Closure(Closure { upvalues, .. })) => upvalues,
             _ => panic!("upvalues() on non-suspended thunk"),
         })
     }
 
-    fn upvalues_mut(&self) -> RefMut<'_, Upvalues> {
-        RefMut::map(self.inner.borrow_mut(), |thunk| match thunk {
+    pub fn upvalues_mut(&self) -> RefMut<'_, Upvalues> {
+        RefMut::map(self.0.borrow_mut(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues,
+            ThunkRepr::Evaluated(Value::Closure(Closure {
+                upvalues,
+                #[cfg(debug_assertions)]
+                is_finalised,
+                ..
+            })) => {
+                #[cfg(debug_assertions)]
+                if *is_finalised {
+                    panic!("Thunk::upvalues_mut() called on a finalised closure");
+                }
+                upvalues
+            }
             thunk => panic!("upvalues() on non-suspended thunk: {thunk:?}"),
         })
     }
@@ -152,7 +190,7 @@ impl UpvalueCarrier for Thunk {
 
 impl Display for Thunk {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self.inner.try_borrow() {
+        match self.0.try_borrow() {
             Ok(repr) => match &*repr {
                 ThunkRepr::Evaluated(v) => v.fmt(f),
                 _ => f.write_str("internal[thunk]"),