feat(tvix/eval): add Thunk::unwrap_or_clone() r/7054

This commit adds Thunk::unwrap_or_clone(), which uses
Rc::try_unwrap() to avoid cloning the Value out of a an Rc which has
only one strong reference.

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

1 files changed, 34 insertions, 1 deletions

diff --git a/tvix/eval/src/value/thunk.rs b/tvix/eval/src/value/thunk.rs
index 9c048d40b4..1ff5beb0d4 100644
--- a/tvix/eval/src/value/thunk.rs
+++ b/tvix/eval/src/value/thunk.rs
@@ -92,6 +92,28 @@ impl ThunkRepr {
             ThunkRepr::Suspended { lambda, .. } => format!("thunk({:p})", *lambda),
         }
     }
+
+    /// Return the Value within a fully-evaluated ThunkRepr; panics
+    /// if the thunk is not fully-evaluated.
+    fn expect(self) -> Value {
+        match self {
+            ThunkRepr::Evaluated(value) => value,
+            ThunkRepr::Blackhole { .. } => panic!("Thunk::expect() called on a black-holed thunk"),
+            ThunkRepr::Suspended { .. } | ThunkRepr::Native(_) => {
+                panic!("Thunk::expect() called on a suspended thunk")
+            }
+        }
+    }
+
+    fn expect_ref(&self) -> &Value {
+        match self {
+            ThunkRepr::Evaluated(value) => value,
+            ThunkRepr::Blackhole { .. } => panic!("Thunk::expect() called on a black-holed thunk"),
+            ThunkRepr::Suspended { .. } | ThunkRepr::Native(_) => {
+                panic!("Thunk::expect() called on a suspended thunk")
+            }
+        }
+    }
 }
 
 /// A thunk is created for any value which requires non-strict
@@ -178,7 +200,7 @@ impl Thunk {
         // If the current thunk is already fully evaluated, return its evaluated
         // value. The VM will continue running the code that landed us here.
         if self.is_forced() {
-            return Ok(self.value().clone());
+            return Ok(self.unwrap_or_clone());
         }
 
         // Begin evaluation of this thunk by marking it as a blackhole, meaning
@@ -280,6 +302,17 @@ impl Thunk {
         })
     }
 
+    /// Returns the inner evaluated value of a thunk, cloning it if
+    /// the Rc has more than one strong reference.  It is an error
+    /// to call this on a thunk that has not been forced, or is not
+    /// otherwise known to be fully evaluated.
+    fn unwrap_or_clone(self) -> Value {
+        match Rc::try_unwrap(self.0) {
+            Ok(refcell) => refcell.into_inner().expect(),
+            Err(rc) => Ref::map(rc.borrow(), |thunkrepr| thunkrepr.expect_ref()).clone(),
+        }
+    }
+
     pub fn upvalues(&self) -> Ref<'_, Upvalues> {
         Ref::map(self.0.borrow(), |thunk| match thunk {
             ThunkRepr::Suspended { upvalues, .. } => upvalues.as_ref(),