diff --git a/turbopack/crates/turbo-tasks-macros/src/func.rs b/turbopack/crates/turbo-tasks-macros/src/func.rs index 054f13e133651..26369f28b6659 100644 --- a/turbopack/crates/turbo-tasks-macros/src/func.rs +++ b/turbopack/crates/turbo-tasks-macros/src/func.rs @@ -561,22 +561,35 @@ impl TurboFn<'_> { let inputs = self.exposed_input_idents(); let persistence = self.persistence_with_this(); parse_quote! { - { - #assertions - let inputs = std::boxed::Box::new((#(#inputs,)*)); - let this = #converted_this; - let persistence = #persistence; - static TRAIT_METHOD: turbo_tasks::macro_helpers::Lazy<&'static turbo_tasks::TraitMethod> = - turbo_tasks::macro_helpers::Lazy::new(|| #trait_type_ident.get(stringify!(#ident))); - <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc( + { + #assertions + let inputs = std::boxed::Box::new((#(#inputs,)*)); + let this = #converted_this; + let persistence = #persistence; + static TRAIT_METHOD: turbo_tasks::macro_helpers::Lazy<&'static turbo_tasks::TraitMethod> = + turbo_tasks::macro_helpers::Lazy::new(|| #trait_type_ident.get(stringify!(#ident))); + static DEVIRTUALIZED: turbo_tasks::macro_helpers::Lazy> = + turbo_tasks::macro_helpers::Lazy::new(|| turbo_tasks::registry::get_devirtualized_trait_impl(*TRAIT_METHOD)); + + <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc(match *DEVIRTUALIZED { + Some(f) => { + turbo_tasks::dynamic_call( + f, + Some(this), + inputs as std::boxed::Box, + persistence, + ) + } + None => { turbo_tasks::trait_call( *TRAIT_METHOD, this, inputs as std::boxed::Box, persistence, ) - ) - } + } + }) + } } } diff --git a/turbopack/crates/turbo-tasks/src/registry.rs b/turbopack/crates/turbo-tasks/src/registry.rs index 7f96208ad37a9..a9a312966b941 100644 --- a/turbopack/crates/turbo-tasks/src/registry.rs +++ b/turbopack/crates/turbo-tasks/src/registry.rs @@ -5,7 +5,7 @@ use once_cell::sync::Lazy; use rustc_hash::{FxHashMap, FxHashSet}; use crate::{ - TraitType, ValueType, + TraitMethod, TraitType, ValueType, id::{FunctionId, TraitTypeId, ValueTypeId}, macro_helpers::CollectableFunction, native_function::NativeFunction, @@ -171,6 +171,56 @@ static TRAITS: Lazy> = Lazy::new(|| { Registry::new_from_items(all_traits) }); +static DEVIRTUALIZED_TRAIT_METHODS: Lazy> = + Lazy::new(|| { + #[derive(Default)] + enum Impls { + #[default] + None, + One(&'static NativeFunction), + MoreThanOne, + } + impl Impls { + fn push(&mut self, func: &'static NativeFunction) { + let new = match self { + Impls::None => Impls::One(func), + Impls::One(_) => Impls::MoreThanOne, + Impls::MoreThanOne => Impls::MoreThanOne, + }; + *self = new; + } + } + let mut trait_method_to_impls = FxHashMap::default(); + for tt in &TRAITS.id_to_item { + for (_, tm) in &tt.methods { + trait_method_to_impls + .insert(tm, tm.default_method.map(Impls::One).unwrap_or(Impls::None)); + } + } + for value in &VALUES.id_to_item { + for (tm, nf) in &value.trait_methods { + trait_method_to_impls.entry(tm).or_default().push(nf); + } + } + + trait_method_to_impls + .into_iter() + .filter_map(|(k, v)| { + if let Impls::One(f) = v { + Some((k, f)) + } else { + None + } + }) + .collect() + }); + +pub fn get_devirtualized_trait_impl( + trait_method: &'static TraitMethod, +) -> Option<&'static NativeFunction> { + DEVIRTUALIZED_TRAIT_METHODS.get(trait_method).copied() +} + pub fn get_trait_type_id(trait_type: &'static TraitType) -> TraitTypeId { TRAITS.get_id(trait_type) } diff --git a/turbopack/crates/turbo-tasks/src/value_type.rs b/turbopack/crates/turbo-tasks/src/value_type.rs index c97188ebb1c03..b5d0e07ea4a32 100644 --- a/turbopack/crates/turbo-tasks/src/value_type.rs +++ b/turbopack/crates/turbo-tasks/src/value_type.rs @@ -34,7 +34,7 @@ pub struct ValueType { /// Set of traits available traits: AutoSet, /// List of trait methods available - trait_methods: AutoMap<&'static TraitMethod, &'static NativeFunction>, + pub(crate) trait_methods: AutoMap<&'static TraitMethod, &'static NativeFunction>, /// Functors for serialization any_serialization: Option<(AnySerializationFn, AnyDeserializeSeed)>,