Perform filtering in local resolution task

vercel · Jan 28, 2025 · 46bad4f · 46bad4f
1 parent e28304e
commit 46bad4f
Show file tree

Hide file tree

Showing 11 changed files with 206 additions and 73 deletions.
diff --git a/turbopack/crates/turbo-tasks-macros/src/func.rs b/turbopack/crates/turbo-tasks-macros/src/func.rs
@@ -480,19 +480,61 @@ impl TurboFn<'_> {
     }
 
     fn inline_input_idents(&self) -> impl Iterator<Item = &Ident> {
-        self.exposed_inputs
-            .iter()
-            .map(|Input { ident, .. }| ident)
+        self.exposed_input_idents()
             .filter(|id| inline_inputs_identifier_filter(id))
     }
 
-    pub fn input_types(&self) -> Vec<&Type> {
+    fn exposed_input_idents(&self) -> impl Iterator<Item = &Ident> {
+        self.exposed_inputs.iter().map(|Input { ident, .. }| ident)
+    }
+
+    pub fn exposed_input_types(&self) -> Vec<&Type> {
         self.exposed_inputs
             .iter()
             .map(|Input { ty, .. }| ty)
             .collect()
     }
 
+    pub fn filter_trait_call_args(&self) -> Option<FilterTraitCallArgsTokens> {
+        // we only need to do this on trait methods, but we're doing it on all methods because we
+        // don't know if we're a trait method or not (we could pass this information down)
+        if self.is_method() {
+            let inline_input_idents: Vec<_> = self.inline_input_idents().collect();
+            if inline_input_idents.len() != self.exposed_inputs.len() {
+                let exposed_input_idents: Vec<_> = self.exposed_input_idents().collect();
+                let exposed_input_types = self.exposed_input_types();
+                return Some(FilterTraitCallArgsTokens {
+                    filter_owned: quote! {
+                        |magic_any| {
+                            let (#(#exposed_input_idents,)*) =
+                                *turbo_tasks::macro_helpers
+                                    ::downcast_args_owned::<(#(#exposed_input_types,)*)>(magic_any);
+                            ::std::boxed::Box::new((#(#inline_input_idents,)*))
+                        }
+                    },
+                    filter_and_resolve: quote! {
+                        |magic_any| {
+                            Box::pin(async move {
+                                let (#(#exposed_input_idents,)*) = turbo_tasks::macro_helpers
+                                    ::downcast_args_ref::<(#(#exposed_input_types,)*)>(magic_any);
+                                let resolved = (#(
+                                    <_ as turbo_tasks::TaskInput>::resolve(
+                                        #inline_input_idents
+                                    ).await?,
+                                )*);
+                                Ok(
+                                    ::std::boxed::Box::new(resolved)
+                                    as ::std::boxed::Box<dyn turbo_tasks::MagicAny>
+                                )
+                            })
+                        }
+                    },
+                });
+            }
+        }
+        None
+    }
+
     pub fn persistence(&self) -> impl ToTokens {
         if self.local {
             quote! {
@@ -1046,46 +1088,79 @@ impl DefinitionContext {
     }
 }
 
+#[derive(Debug)]
+pub struct FilterTraitCallArgsTokens {
+    filter_owned: TokenStream,
+    filter_and_resolve: TokenStream,
+}
+
 #[derive(Debug)]
 pub struct NativeFn {
     pub function_path_string: String,
     pub function_path: ExprPath,
     pub is_method: bool,
+    pub filter_trait_call_args: Option<FilterTraitCallArgsTokens>,
     pub local: bool,
     pub local_cells: bool,
 }
 
 impl NativeFn {
     pub fn ty(&self) -> Type {
-        parse_quote! { turbo_tasks::NativeFunction }
+        parse_quote! { turbo_tasks::macro_helpers::NativeFunction }
     }
 
     pub fn definition(&self) -> TokenStream {
         let Self {
             function_path_string,
             function_path,
             is_method,
+            filter_trait_call_args,
             local,
             local_cells,
         } = self;
 
-        let constructor = if *is_method {
-            quote! { new_method }
+        if *is_method {
+            let arg_filter = if let Some(filter) = filter_trait_call_args {
+                let FilterTraitCallArgsTokens {
+                    filter_owned,
+                    filter_and_resolve,
+                } = filter;
+                quote! {
+                    ::std::option::Option::Some((
+                        #filter_owned,
+                        #filter_and_resolve,
+                    ))
+                }
+            } else {
+                quote! { ::std::option::Option::None }
+            };
+            quote! {
+                {
+                    #[allow(deprecated)]
+                    turbo_tasks::macro_helpers::NativeFunction::new_method(
+                        #function_path_string.to_owned(),
+                        turbo_tasks::macro_helpers::FunctionMeta {
+                            local: #local,
+                            local_cells: #local_cells,
+                        },
+                        #arg_filter,
+                        #function_path,
+                    )
+                }
+            }
         } else {
-            quote! { new_function }
-        };
-
-        quote! {
-            {
-                #[allow(deprecated)]
-                turbo_tasks::NativeFunction::#constructor(
-                    #function_path_string.to_owned(),
-                    turbo_tasks::FunctionMeta {
-                        local: #local,
-                        local_cells: #local_cells,
-                    },
-                    #function_path,
-                )
+            quote! {
+                {
+                    #[allow(deprecated)]
+                    turbo_tasks::macro_helpers::NativeFunction::new_function(
+                        #function_path_string.to_owned(),
+                        turbo_tasks::macro_helpers::FunctionMeta {
+                            local: #local,
+                            local_cells: #local_cells,
+                        },
+                        #function_path,
+                    )
+                }
             }
         }
     }

diff --git a/turbopack/crates/turbo-tasks-macros/src/function_macro.rs b/turbopack/crates/turbo-tasks-macros/src/function_macro.rs
@@ -59,6 +59,7 @@ pub fn function(args: TokenStream, input: TokenStream) -> TokenStream {
         function_path_string: ident.to_string(),
         function_path: parse_quote! { #inline_function_ident },
         is_method: turbo_fn.is_method(),
+        filter_trait_call_args: None, // not a trait method
         local,
         local_cells,
     };

diff --git a/turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs b/turbopack/crates/turbo-tasks-macros/src/value_impl_macro.rs
@@ -140,6 +140,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {
                     function_path_string: format!("{ty}::{ident}", ty = ty.to_token_stream()),
                     function_path: parse_quote! { <#ty>::#inline_function_ident },
                     is_method: turbo_fn.is_method(),
+                    filter_trait_call_args: None, // not a trait method
                     local,
                     local_cells,
                 };
@@ -252,6 +253,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {
                         <#ty as #inline_extension_trait_ident>::#inline_function_ident
                     },
                     is_method: turbo_fn.is_method(),
+                    filter_trait_call_args: turbo_fn.filter_trait_call_args(),
                     local,
                     local_cells,
                 };

diff --git a/turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs b/turbopack/crates/turbo-tasks-macros/src/value_trait_macro.rs
@@ -101,7 +101,7 @@ pub fn value_trait(args: TokenStream, input: TokenStream) -> TokenStream {
         };
 
         let turbo_signature = turbo_fn.signature();
-        let arg_types = turbo_fn.input_types();
+        let arg_types = turbo_fn.exposed_input_types();
         let dynamic_block = turbo_fn.dynamic_block(&trait_type_id_ident);
         dynamic_trait_fns.push(quote! {
             #turbo_signature #dynamic_block
@@ -120,6 +120,7 @@ pub fn value_trait(args: TokenStream, input: TokenStream) -> TokenStream {
                     <Box<dyn #trait_ident> as #inline_extension_trait_ident>::#inline_function_ident
                 },
                 is_method: turbo_fn.is_method(),
+                filter_trait_call_args: turbo_fn.filter_trait_call_args(),
                 // `local` and `local_cells` are currently unsupported here because:
                 // - The `#[turbo_tasks::function]` macro needs to be present for us to read this
                 //   argument. (This could be fixed)

diff --git a/turbopack/crates/turbo-tasks/src/lib.rs b/turbopack/crates/turbo-tasks/src/lib.rs
@@ -107,7 +107,6 @@ pub use manager::{
     CurrentCellRef, ReadConsistency, TaskPersistence, TurboTasks, TurboTasksApi,
     TurboTasksBackendApi, TurboTasksBackendApiExt, TurboTasksCallApi, Unused, UpdateInfo,
 };
-pub use native_function::{FunctionMeta, NativeFunction};
 pub use output::OutputContent;
 pub use raw_vc::{CellId, RawVc, ReadRawVcFuture, ResolveTypeError};
 pub use read_ref::ReadRef;

diff --git a/turbopack/crates/turbo-tasks/src/macro_helpers.rs b/turbopack/crates/turbo-tasks/src/macro_helpers.rs
@@ -6,14 +6,15 @@ pub use once_cell::sync::{Lazy, OnceCell};
 pub use serde;
 pub use tracing;
 
-pub use super::{
-    magic_any::MagicAny,
-    manager::{find_cell_by_type, notify_scheduled_tasks, spawn_detached_for_testing},
-};
 use crate::{
     debug::ValueDebugFormatString, shrink_to_fit::ShrinkToFit, task::TaskOutput, NonLocalValue,
     RawVc, TaskInput, TaskPersistence, Vc,
 };
+pub use crate::{
+    magic_any::MagicAny,
+    manager::{find_cell_by_type, notify_scheduled_tasks, spawn_detached_for_testing},
+    native_function::{downcast_args_owned, downcast_args_ref, FunctionMeta, NativeFunction},
+};
 
 #[inline(never)]
 pub async fn value_debug_format_field(value: ValueDebugFormatString<'_>) -> String {

diff --git a/turbopack/crates/turbo-tasks/src/manager.rs b/turbopack/crates/turbo-tasks/src/manager.rs
@@ -37,6 +37,7 @@ use crate::{
     },
     id_factory::{IdFactory, IdFactoryWithReuse},
     magic_any::MagicAny,
+    native_function::FunctionMeta,
     raw_vc::{CellId, RawVc},
     registry::{self, get_function},
     serialization_invalidation::SerializationInvalidator,
@@ -48,8 +49,8 @@ use crate::{
     trait_helpers::get_trait_method,
     util::StaticOrArc,
     vc::ReadVcFuture,
-    Completion, FunctionMeta, InvalidationReason, InvalidationReasonSet, ResolvedVc,
-    SharedReference, TaskId, TaskIdSet, ValueTypeId, Vc, VcRead, VcValueTrait, VcValueType,
+    Completion, InvalidationReason, InvalidationReasonSet, ResolvedVc, SharedReference, TaskId,
+    TaskIdSet, ValueTypeId, Vc, VcRead, VcValueTrait, VcValueType,
 };
 
 pub trait TurboTasksCallApi: Sync + Send {
@@ -638,6 +639,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
         if let RawVc::TaskCell(_, CellId { type_id, .. }) = this {
             match get_trait_method(trait_type, type_id, trait_fn_name) {
                 Ok(native_fn) => {
+                    let arg = registry::get_function(native_fn).arg_meta.filter_owned(arg);
                     return self.dynamic_call(native_fn, Some(this), arg, persistence);
                 }
                 Err(name) => {