Add TaskPersistence enum for task creation functions

bgw · bgw · commit 02ecb50a4989 · 2024-08-14T16:08:36.000-04:00
diff --git a/turbopack/crates/turbo-tasks-macros/src/func.rs b/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,
                         )
                     )
                 }
diff --git a/turbopack/crates/turbo-tasks-memory/src/task.rs b/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);
diff --git a/turbopack/crates/turbo-tasks-testing/src/lib.rs b/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!()
     }
diff --git a/turbopack/crates/turbo-tasks/src/backend.rs b/turbopack/crates/turbo-tasks/src/backend.rs
@@ -23,8 +23,8 @@ use crate::{
     task::shared_reference::TypedSharedReference,
     trait_helpers::{get_trait_method, has_trait, traits},
     triomphe_utils::unchecked_sidecast_triomphe_arc,
-    FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdSet, TraitRef, TraitTypeId,
-    ValueTypeId, VcRead, VcValueTrait, VcValueType,
+    FunctionId, RawVc, ReadRef, SharedReference, TaskId, TaskIdSet, TaskPersistence, TraitRef,
+    TraitTypeId, ValueTypeId, VcRead, VcValueTrait, VcValueType,
 };
 
 type TransientTaskRoot =
@@ -608,17 +608,17 @@ impl CachedTaskType {
         fn_id: FunctionId,
         mut this: Option<RawVc>,
         arg: &dyn MagicAny,
-        is_transient: bool,
+        persistence: TaskPersistence,
         turbo_tasks: Arc<dyn TurboTasksBackendApi<B>>,
     ) -> Result<RawVc> {
         if let Some(this) = this.as_mut() {
             *this = this.resolve().await?;
         }
         let arg = registry::get_function(fn_id).arg_meta.resolve(arg).await?;
         Ok(if let Some(this) = this {
-            turbo_tasks.this_call(fn_id, this, arg, is_transient)
+            turbo_tasks.this_call(fn_id, this, arg, persistence)
         } else {
-            turbo_tasks.native_call(fn_id, arg, is_transient)
+            turbo_tasks.native_call(fn_id, arg, persistence)
         })
     }
 
@@ -636,7 +636,7 @@ impl CachedTaskType {
         name: Cow<'static, str>,
         this: RawVc,
         arg: &dyn MagicAny,
-        is_transient: bool,
+        persistence: TaskPersistence,
         turbo_tasks: Arc<dyn TurboTasksBackendApi<B>>,
     ) -> Result<RawVc> {
         let this = this.resolve().await?;
@@ -647,7 +647,7 @@ impl CachedTaskType {
             .arg_meta
             .resolve(arg)
             .await?;
-        Ok(turbo_tasks.dynamic_this_call(native_fn, this, arg, is_transient))
+        Ok(turbo_tasks.dynamic_this_call(native_fn, this, arg, persistence))
     }
 
     /// Shared helper used by [`Self::resolve_trait_method`] and
diff --git a/turbopack/crates/turbo-tasks/src/id.rs b/turbopack/crates/turbo-tasks/src/id.rs
@@ -7,7 +7,7 @@ use std::{
 
 use serde::{de::Visitor, Deserialize, Serialize};
 
-use crate::registry;
+use crate::{registry, TaskPersistence};
 
 macro_rules! define_id {
     ($name:ident : $primitive:ty $(,derive($($derive:ty),*))?) => {
@@ -84,6 +84,14 @@ impl TaskId {
     pub fn is_transient(&self) -> bool {
         **self & TRANSIENT_TASK_BIT != 0
     }
+    pub fn persistence(&self) -> TaskPersistence {
+        // tasks with `TaskPersistence::LocalCells` have no `TaskId`, so we can ignore that case
+        if self.is_transient() {
+            TaskPersistence::Transient
+        } else {
+            TaskPersistence::Persistent
+        }
+    }
 }
 
 macro_rules! make_serializable {
diff --git a/turbopack/crates/turbo-tasks/src/lib.rs b/turbopack/crates/turbo-tasks/src/lib.rs
@@ -90,8 +90,8 @@ pub use magic_any::MagicAny;
 pub use manager::{
     dynamic_call, dynamic_this_call, emit, get_invalidator, mark_finished, mark_stateful,
     prevent_gc, run_once, run_once_with_reason, spawn_blocking, spawn_thread, trait_call,
-    turbo_tasks, CurrentCellRef, Invalidator, TurboTasks, TurboTasksApi, TurboTasksBackendApi,
-    TurboTasksCallApi, Unused, UpdateInfo,
+    turbo_tasks, CurrentCellRef, Invalidator, TaskPersistence, TurboTasks, TurboTasksApi,
+    TurboTasksBackendApi, TurboTasksCallApi, Unused, UpdateInfo,
 };
 pub use native_function::{FunctionMeta, NativeFunction};
 pub use raw_vc::{CellId, RawVc, ReadRawVcFuture, ResolveTypeError};
diff --git a/turbopack/crates/turbo-tasks/src/macro_helpers.rs b/turbopack/crates/turbo-tasks/src/macro_helpers.rs
@@ -8,7 +8,7 @@ pub use super::{
     magic_any::MagicAny,
     manager::{find_cell_by_type, notify_scheduled_tasks, spawn_detached_for_testing},
 };
-use crate::debug::ValueDebugFormatString;
+use crate::{debug::ValueDebugFormatString, TaskInput, TaskPersistence};
 
 #[inline(never)]
 pub async fn value_debug_format_field(value: ValueDebugFormatString<'_>) -> String {
@@ -21,6 +21,14 @@ pub async fn value_debug_format_field(value: ValueDebugFormatString<'_>) -> Stri
     }
 }
 
+pub fn get_non_local_persistence_from_inputs(inputs: &impl TaskInput) -> TaskPersistence {
+    if inputs.is_transient() {
+        TaskPersistence::Transient
+    } else {
+        TaskPersistence::Persistent
+    }
+}
+
 #[macro_export]
 macro_rules! stringify_path {
     ($path:path) => {
diff --git a/turbopack/crates/turbo-tasks/src/manager.rs b/turbopack/crates/turbo-tasks/src/manager.rs
