refactor(turbo-tasks): Add TaskPersistence enum for task creation functions

turbopack/crates/turbo-tasks-macros/src/func.rs

@@ -1,7 +1,7 @@
 use std::collections::HashSet;
 
 use proc_macro2::{Ident, Span, TokenStream};
-use quote::{quote, quote_spanned};
+use quote::{quote, quote_spanned, ToTokens};
 use syn::{
     parenthesized,
     parse::{Parse, ParseStream},
@@ -24,6 +24,8 @@ pub struct TurboFn {
     inputs: Vec<Input>,
     /// Should we check that the return type contains a `ResolvedValue`?
     resolved: Option<Span>,
+    /// Should this function use `TaskPersistence::LocalCells`?
+    local_cells: bool,
 }
 
 #[derive(Debug)]
@@ -257,6 +259,7 @@ impl TurboFn {
             this,
             inputs,
             resolved: args.resolved,
+            local_cells: args.local_cells.is_some(),
         })
     }
 
@@ -301,17 +304,26 @@ impl TurboFn {
         }
     }
 
-    fn inputs(&self) -> Vec<&Ident> {
-        self.inputs
-            .iter()
-            .map(|Input { ident, .. }| ident)
-            .collect()
+    fn input_idents(&self) -> impl Iterator<Item = &Ident> {
+        self.inputs.iter().map(|Input { ident, .. }| ident)
     }
 
     pub fn input_types(&self) -> Vec<&Type> {
         self.inputs.iter().map(|Input { ty, .. }| ty).collect()
     }
 
+    pub fn persistence(&self) -> impl ToTokens {
+        if self.local_cells {
+            quote! {
+                turbo_tasks::TaskPersistence::LocalCells
+            }
+        } else {
+            quote! {
+                turbo_tasks::macro_helpers::get_non_local_persistence_from_inputs(&*inputs)
+            }
+        }
+    }
+
     fn converted_this(&self) -> Option<Expr> {
         self.this.as_ref().map(|Input { ty: _, ident }| {
             parse_quote! {
@@ -344,31 +356,33 @@ impl TurboFn {
     /// The block of the exposed function for a dynamic dispatch call to the
     /// given trait.
     pub fn dynamic_block(&self, trait_type_id_ident: &Ident) -> Block {
+        let Some(converted_this) = self.converted_this() else {
+            return parse_quote! {
+                {
+                    unimplemented!("trait methods without self are not yet supported")
+                }
+            };
+        };
+
         let ident = &self.ident;
         let output = &self.output;
         let assertions = self.get_assertions();
-        if let Some(converted_this) = self.converted_this() {
-            let inputs = self.inputs();
-            parse_quote! {
-                {
-                    #assertions
-                    let turbo_tasks_transient = #( turbo_tasks::TaskInput::is_transient(&#inputs) ||)* false;
-                    <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc(
-                        turbo_tasks::trait_call(
-                            *#trait_type_id_ident,
-                            std::borrow::Cow::Borrowed(stringify!(#ident)),
-                            #converted_this,
-                            Box::new((#(#inputs,)*)) as Box<dyn turbo_tasks::MagicAny>,
-                            turbo_tasks_transient,
-                        )
+        let inputs = self.input_idents();
+        let persistence = self.persistence();
+        parse_quote! {
+            {
+                #assertions
+                let inputs = std::boxed::Box::new((#(#inputs,)*));
+                let persistence = #persistence;
+                <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc(
+                    turbo_tasks::trait_call(
+                        *#trait_type_id_ident,
+                        std::borrow::Cow::Borrowed(stringify!(#ident)),
+                        #converted_this,
+                        inputs as std::boxed::Box<dyn turbo_tasks::MagicAny>,
+                        persistence,
                     )
-                }
-            }
-        } else {
-            parse_quote! {
-                {
-                    unimplemented!("trait methods without self are not yet supported")
-                }
+                )
             }
         }
     }
@@ -377,19 +391,21 @@ impl TurboFn {
     /// given native function.
     pub fn static_block(&self, native_function_id_ident: &Ident) -> Block {
         let output = &self.output;
-        let inputs = self.inputs();
+        let inputs = self.input_idents();
+        let persistence = self.persistence();
         let assertions = self.get_assertions();
         if let Some(converted_this) = self.converted_this() {
             parse_quote! {
                 {
                     #assertions
-                    let turbo_tasks_transient = #( turbo_tasks::TaskInput::is_transient(&#inputs) ||)* false;
+                    let inputs = std::boxed::Box::new((#(#inputs,)*));
+                    let persistence = #persistence;
                     <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc(
                         turbo_tasks::dynamic_this_call(
                             *#native_function_id_ident,
                             #converted_this,
-                            Box::new((#(#inputs,)*)) as Box<dyn turbo_tasks::MagicAny>,
-                            turbo_tasks_transient
+                            inputs as std::boxed::Box<dyn turbo_tasks::MagicAny>,
+                            persistence,
                         )
                     )
                 }
@@ -398,12 +414,13 @@ impl TurboFn {
             parse_quote! {
                 {
                     #assertions
-                    let turbo_tasks_transient = #( turbo_tasks::TaskInput::is_transient(&#inputs) ||)* false;
+                    let inputs = std::boxed::Box::new((#(#inputs,)*));
+                    let persistence = #persistence;
                     <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc(
                         turbo_tasks::dynamic_call(
                             *#native_function_id_ident,
-                            Box::new((#(#inputs,)*)) as Box<dyn turbo_tasks::MagicAny>,
-                            turbo_tasks_transient,
+                            inputs as std::boxed::Box<dyn turbo_tasks::MagicAny>,
+                            persistence,
                         )
                     )
                 }

turbopack/crates/turbo-tasks-memory/src/task.rs

@@ -810,7 +810,7 @@ impl Task {
                         native_fn_id,
                         *this,
                         &**arg,
-                        self.id.is_transient(),
+                        self.id.persistence(),
                         turbo_tasks,
                     ));
                     drop(entered);
@@ -833,7 +833,7 @@ impl Task {
                         name,
                         *this,
                         &**arg,
-                        self.id.is_transient(),
+                        self.id.persistence(),
                         turbo_tasks,
                     ));
                     drop(entered);

turbopack/crates/turbo-tasks-memory/tests/local_cell.rs

@@ -0,0 +1 @@
+../../turbo-tasks-testing/tests/local_cell.rs

turbopack/crates/turbo-tasks-testing/src/lib.rs

@@ -20,8 +20,8 @@ use turbo_tasks::{
     registry,
     test_helpers::with_turbo_tasks_for_testing,
     util::{SharedError, StaticOrArc},
-    CellId, ExecutionId, InvalidationReason, MagicAny, RawVc, TaskId, TraitTypeId, TurboTasksApi,
-    TurboTasksCallApi,
+    CellId, ExecutionId, InvalidationReason, MagicAny, RawVc, TaskId, TaskPersistence, TraitTypeId,
+    TurboTasksApi, TurboTasksCallApi,
 };
 
 pub use crate::run::{run, run_without_cache_check, Registration};
@@ -92,7 +92,7 @@ impl TurboTasksCallApi for VcStorage {
         &self,
         func: turbo_tasks::FunctionId,
         arg: Box<dyn MagicAny>,
-        _is_transient: bool,
+        _persistence: TaskPersistence,
     ) -> RawVc {
         self.dynamic_call(func, None, arg)
     }
@@ -102,7 +102,7 @@ impl TurboTasksCallApi for VcStorage {
         func: turbo_tasks::FunctionId,
         this_arg: RawVc,
         arg: Box<dyn MagicAny>,
-        _is_transient: bool,
+        _persistence: TaskPersistence,
     ) -> RawVc {
         self.dynamic_call(func, Some(this_arg), arg)
     }
@@ -111,7 +111,7 @@ impl TurboTasksCallApi for VcStorage {
         &self,
         _func: turbo_tasks::FunctionId,
         _arg: Box<dyn MagicAny>,
-        _is_transient: bool,
+        _persistence: TaskPersistence,
     ) -> RawVc {
         unreachable!()
     }
@@ -121,7 +121,7 @@ impl TurboTasksCallApi for VcStorage {
         _func: turbo_tasks::FunctionId,
         _this: RawVc,
         _arg: Box<dyn MagicAny>,
-        _is_transient: bool,
+        _persistence: TaskPersistence,
     ) -> RawVc {
         unreachable!()
     }
@@ -132,7 +132,7 @@ impl TurboTasksCallApi for VcStorage {
         _trait_fn_name: Cow<'static, str>,
         _this: RawVc,
         _arg: Box<dyn MagicAny>,
-        _is_transient: bool,
+        _persistence: TaskPersistence,
     ) -> RawVc {
         unreachable!()
     }

turbopack/crates/turbo-tasks-testing/tests/local_cell.rs

@@ -1,5 +1,6 @@
 #![feature(arbitrary_self_types)]
 
+use anyhow::Result;
 use turbo_tasks::{
     debug::ValueDebug, test_helpers::current_task_for_testing, ResolvedValue, ValueDefault, Vc,
 };
@@ -14,8 +15,8 @@ struct Wrapper(u32);
 struct TransparentWrapper(u32);
 
 #[tokio::test]
-async fn test_store_and_read() {
-    run(&REGISTRATION, async {
+async fn test_store_and_read() -> Result<()> {
+    run(&REGISTRATION, || async {
         let a: Vc<u32> = Vc::local_cell(42);
         assert_eq!(*a.await.unwrap(), 42);
 
@@ -24,13 +25,15 @@ async fn test_store_and_read() {
 
         let c = TransparentWrapper(42).local_cell();
         assert_eq!(*c.await.unwrap(), 42);
+
+        Ok(())
     })
     .await
 }
 
 #[tokio::test]
-async fn test_store_and_read_generic() {
-    run(&REGISTRATION, async {
+async fn test_store_and_read_generic() -> Result<()> {
+    run(&REGISTRATION, || async {
         // `Vc<Vec<Vc<T>>>` is stored as `Vc<Vec<Vc<()>>>` and requires special
         // transmute handling
         let cells: Vc<Vec<Vc<u32>>> =
@@ -42,6 +45,8 @@ async fn test_store_and_read_generic() {
         }
 
         assert_eq!(output, vec![1, 2, 3]);
+
+        Ok(())
     })
     .await
 }
@@ -53,10 +58,12 @@ async fn returns_resolved_local_vc() -> Vc<u32> {
     cell.resolve().await.unwrap()
 }
 
+#[ignore]
 #[tokio::test]
-async fn test_return_resolved() {
-    run(&REGISTRATION, async {
+async fn test_return_resolved() -> Result<()> {
+    run(&REGISTRATION, || async {
         assert_eq!(*returns_resolved_local_vc().await.unwrap(), 42);
+        Ok(())
     })
     .await
 }
@@ -65,8 +72,8 @@ async fn test_return_resolved() {
 trait UnimplementedTrait {}
 
 #[tokio::test]
-async fn test_try_resolve_sidecast() {
-    run(&REGISTRATION, async {
+async fn test_try_resolve_sidecast() -> Result<()> {
+    run(&REGISTRATION, || async {
         let trait_vc: Vc<Box<dyn ValueDebug>> = Vc::upcast(Vc::<u32>::local_cell(42));
 
         // `u32` is both a `ValueDebug` and a `ValueDefault`, so this sidecast is valid
@@ -80,13 +87,15 @@ async fn test_try_resolve_sidecast() {
             .await
             .unwrap();
         assert!(wrongly_sidecast_vc.is_none());
+
+        Ok(())
     })
     .await
 }
 
 #[tokio::test]
-async fn test_try_resolve_downcast_type() {
-    run(&REGISTRATION, async {
+async fn test_try_resolve_downcast_type() -> Result<()> {
+    run(&REGISTRATION, || async {
         let trait_vc: Vc<Box<dyn ValueDebug>> = Vc::upcast(Vc::<u32>::local_cell(42));
 
         let downcast_vc: Vc<u32> = Vc::try_resolve_downcast_type(trait_vc)
@@ -98,16 +107,19 @@ async fn test_try_resolve_downcast_type() {
         let wrongly_downcast_vc: Option<Vc<i64>> =
             Vc::try_resolve_downcast_type(trait_vc).await.unwrap();
         assert!(wrongly_downcast_vc.is_none());
+
+        Ok(())
     })
     .await
 }
 
 #[tokio::test]
-async fn test_get_task_id() {
-    run(&REGISTRATION, async {
+async fn test_get_task_id() -> Result<()> {
+    run(&REGISTRATION, || async {
         // the task id as reported by the RawVc
         let vc_task_id = Vc::into_raw(Vc::<()>::local_cell(())).get_task_id();
         assert_eq!(vc_task_id, current_task_for_testing());
+        Ok(())
     })
     .await
 }
@@ -139,25 +151,31 @@ async fn get_untracked_local_cell() -> Vc<Untracked> {
         .unwrap()
 }
 
+#[ignore]
 #[tokio::test]
 #[should_panic(expected = "Local Vcs must only be accessed within their own task")]
 async fn test_panics_on_local_cell_escape_read() {
-    run(&REGISTRATION, async {
+    run(&REGISTRATION, || async {
         get_untracked_local_cell()
             .await
             .unwrap()
             .cell
             .await
             .unwrap();
+        Ok(())
     })
     .await
+    .unwrap()
 }
 
+#[ignore]
 #[tokio::test]
 #[should_panic(expected = "Local Vcs must only be accessed within their own task")]
 async fn test_panics_on_local_cell_escape_get_task_id() {
-    run(&REGISTRATION, async {
+    run(&REGISTRATION, || async {
         Vc::into_raw(get_untracked_local_cell().await.unwrap().cell).get_task_id();
+        Ok(())
     })
     .await
+    .unwrap()
 }