[turbopack] Devirtualize some trait calls

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<Option<&'static turbo_tasks::macro_helpers::NativeFunction>> =
+                 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<dyn turbo_tasks::MagicAny>,
+                        persistence,
+                    )
+                }
+                None => {
                     turbo_tasks::trait_call(
                         *TRAIT_METHOD,
                         this,
                         inputs as std::boxed::Box<dyn turbo_tasks::MagicAny>,
                         persistence,
                     )
-                )
-            }
+                }
+            })
+        }
         }
     }
 

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<Registry<TraitType>> = Lazy::new(|| {
     Registry::new_from_items(all_traits)
 });
 
+static DEVIRTUALIZED_TRAIT_METHODS: Lazy<FxHashMap<&'static TraitMethod, &'static NativeFunction>> =
+    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)
 }

turbopack/crates/turbo-tasks/src/value_type.rs

@@ -34,7 +34,7 @@ pub struct ValueType {
     /// Set of traits available
     traits: AutoSet<TraitTypeId>,
     /// 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)>,