Devirtualize some trait calls

Author: lukesandberg

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

@@ -561,22 +561,37 @@ 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));
+
+            // Perform direct calls to devirtualized trait methods.
+            // ideally we would call directly to the outer implementation function on the only implementation
+            <#output as turbo_tasks::task::TaskOutput>::try_from_raw_vc(match *DEVIRTUALIZED {
+                Some(f) => {
+                    turbo_tasks::devirtualized_trait_call(
+                        f,
+                        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-testing/src/lib.rs

@@ -119,7 +119,15 @@ impl TurboTasksCallApi for VcStorage {
     ) -> RawVc {
         unreachable!()
     }
-
+    fn devirtualized_trait_call(
+        &self,
+        __trait_impl: &'static turbo_tasks::macro_helpers::NativeFunction,
+        _this: RawVc,
+        _arg: Box<dyn MagicAny>,
+        _persistence: TaskPersistence,
+    ) -> RawVc {
+        unreachable!()
+    }
     fn run(
         &self,
         _future: Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>,

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

@@ -107,9 +107,9 @@ pub use key_value_pair::KeyValuePair;
 pub use magic_any::MagicAny;
 pub use manager::{
     CurrentCellRef, ReadConsistency, ReadTracking, TaskPersistence, TurboTasks, TurboTasksApi,
-    TurboTasksBackendApi, TurboTasksCallApi, Unused, UpdateInfo, dynamic_call, emit, mark_finished,
-    mark_root, mark_session_dependent, mark_stateful, prevent_gc, run, run_once,
-    run_once_with_reason, trait_call, turbo_tasks, turbo_tasks_scope,
+    TurboTasksBackendApi, TurboTasksCallApi, Unused, UpdateInfo, devirtualized_trait_call,
+    dynamic_call, emit, mark_finished, mark_root, mark_session_dependent, mark_stateful,
+    prevent_gc, run, run_once, run_once_with_reason, trait_call, turbo_tasks, turbo_tasks_scope,
 };
 pub use output::OutputContent;
 pub use raw_vc::{CellId, RawVc, ReadRawVcFuture, ResolveTypeError};

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

@@ -73,7 +73,13 @@ pub trait TurboTasksCallApi: Sync + Send {
         arg: Box<dyn MagicAny>,
         persistence: TaskPersistence,
     ) -> RawVc;
-
+    fn devirtualized_trait_call(
+        &self,
+        trait_impl: &'static NativeFunction,
+        this: RawVc,
+        arg: Box<dyn MagicAny>,
+        persistence: TaskPersistence,
+    ) -> RawVc;
     fn run(
         &self,
         future: Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>,
@@ -656,6 +662,17 @@ impl<B: Backend + 'static> TurboTasks<B> {
         self.schedule_local_task(task_type, persistence)
     }
 
+    pub fn devirtualized_trait_call(
+        &self,
+        trait_impl: &'static NativeFunction,
+        this: RawVc,
+        arg: Box<dyn MagicAny>,
+        persistence: TaskPersistence,
+    ) -> RawVc {
+        let arg = trait_impl.arg_meta.filter_owned(arg);
+        return self.dynamic_call(trait_impl, Some(this), arg, persistence);
+    }
+
     pub fn trait_call(
         &self,
         trait_method: &'static TraitMethod,
@@ -1183,7 +1200,15 @@ impl<B: Backend + 'static> TurboTasksCallApi for TurboTasks<B> {
     ) -> RawVc {
         self.trait_call(trait_method, this, arg, persistence)
     }
-
+    fn devirtualized_trait_call(
+        &self,
+        trait_impl: &'static NativeFunction,
+        this: RawVc,
+        arg: Box<dyn MagicAny>,
+        persistence: TaskPersistence,
+    ) -> RawVc {
+        self.devirtualized_trait_call(trait_impl, this, arg, persistence)
+    }
     #[track_caller]
     fn run(
         &self,
@@ -1579,6 +1604,15 @@ pub fn trait_call(
 ) -> RawVc {
     with_turbo_tasks(|tt| tt.trait_call(trait_method, this, arg, persistence))
 }
+/// Calls [`TurboTasks::devirtualized_trait_call`] for the current turbo tasks instance.
+pub fn devirtualized_trait_call(
+    trait_impl: &'static NativeFunction,
+    this: RawVc,
+    arg: Box<dyn MagicAny>,
+    persistence: TaskPersistence,
+) -> RawVc {
+    with_turbo_tasks(|tt| tt.devirtualized_trait_call(trait_impl, this, arg, persistence))
+}
 
 pub fn turbo_tasks() -> Arc<dyn TurboTasksApi> {
     TURBO_TASKS.with(|arc| arc.clone())

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)>,

turbopack/crates/turbopack-core/src/ident.rs

@@ -4,7 +4,7 @@ use anyhow::Result;
 use once_cell::sync::Lazy;
 use regex::Regex;
 use serde::{Deserialize, Serialize};
-use turbo_rcstr::RcStr;
+use turbo_rcstr::{RcStr, rcstr};
 use turbo_tasks::{
     NonLocalValue, ReadRef, ResolvedVc, TaskInput, ValueToString, Vc, trace::TraceRawVcs,
 };
@@ -104,6 +104,17 @@ impl AssetIdent {
     }
 }
 
+#[turbo_tasks::value]
+struct TestStruct {}
+
+#[turbo_tasks::value_impl]
+impl ValueToString for TestStruct {
+    #[turbo_tasks::function]
+    fn to_string(self: Vc<Self>) -> Vc<RcStr> {
+        Vc::cell(rcstr!("huh"))
+    }
+}
+
 #[turbo_tasks::value_impl]
 impl AssetIdent {
     #[turbo_tasks::function]