perf(turbo-tasks): Call .shrink_to_fit() on common collection types…

… when constructing a cell

turbopack/crates/turbo-tasks-backend/tests/shrink_to_fit.rs

@@ -0,0 +1 @@
+../../turbo-tasks-testing/tests/shrink_to_fit.rs

turbopack/crates/turbo-tasks-macros-shared/src/primitive_input.rs

@@ -1,16 +1,47 @@
+use proc_macro2::Span;
 use syn::{
     parse::{Parse, ParseStream},
-    Result, Type,
+    punctuated::Punctuated,
+    spanned::Spanned,
+    Meta, Result, Token, Type,
 };
 
 #[derive(Debug)]
 pub struct PrimitiveInput {
     pub ty: Type,
+    pub manual_shrink_to_fit: Option<Span>,
 }
 
 impl Parse for PrimitiveInput {
     fn parse(input: ParseStream) -> Result<Self> {
         let ty: Type = input.parse()?;
-        Ok(PrimitiveInput { ty })
+        let mut parsed_input = PrimitiveInput {
+            ty,
+            manual_shrink_to_fit: None,
+        };
+        if input.parse::<Option<Token![,]>>()?.is_some() {
+            let punctuated: Punctuated<Meta, Token![,]> = input.parse_terminated(Meta::parse)?;
+            for meta in punctuated {
+                match (
+                    meta.path()
+                        .get_ident()
+                        .map(ToString::to_string)
+                        .as_deref()
+                        .unwrap_or_default(),
+                    &meta,
+                ) {
+                    ("manual_shrink_to_fit", Meta::Path(_)) => {
+                        parsed_input.manual_shrink_to_fit = Some(meta.span())
+                    }
+                    (_, meta) => {
+                        return Err(syn::Error::new_spanned(
+                            meta,
+                            "unexpected token, expected: \"manual_shrink_to_fit\"",
+                        ))
+                    }
+                }
+            }
+        }
+        Ok(parsed_input)
     }
 }

turbopack/crates/turbo-tasks-macros/src/derive/mod.rs

@@ -1,6 +1,7 @@
 mod deterministic_hash_macro;
 mod key_value_pair_macro;
 mod resolved_value_macro;
+mod shrink_to_fit_macro;
 mod task_input_macro;
 mod trace_raw_vcs_macro;
 mod value_debug_format_macro;
@@ -9,6 +10,7 @@ mod value_debug_macro;
 pub use deterministic_hash_macro::derive_deterministic_hash;
 pub use key_value_pair_macro::derive_key_value_pair;
 pub use resolved_value_macro::derive_resolved_value;
+pub use shrink_to_fit_macro::derive_shrink_to_fit;
 use syn::{spanned::Spanned, Attribute, Meta, MetaList, NestedMeta};
 pub use task_input_macro::derive_task_input;
 pub use trace_raw_vcs_macro::derive_trace_raw_vcs;

turbopack/crates/turbo-tasks-macros/src/derive/shrink_to_fit_macro.rs

@@ -0,0 +1,53 @@
+use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::quote;
+use syn::{parse_macro_input, DeriveInput, FieldsNamed, FieldsUnnamed};
+use turbo_tasks_macros_shared::{generate_exhaustive_destructuring, match_expansion};
+
+pub fn derive_shrink_to_fit(input: TokenStream) -> TokenStream {
+    let derive_input = parse_macro_input!(input as DeriveInput);
+    let ident = &derive_input.ident;
+    let (impl_generics, ty_generics, where_clause) = derive_input.generics.split_for_impl();
+
+    let shrink_items = match_expansion(&derive_input, &shrink_named, &shrink_unnamed, &shrink_unit);
+    quote! {
+        impl #impl_generics turbo_tasks::ShrinkToFit for #ident #ty_generics #where_clause {
+            fn shrink_to_fit(&mut self) {
+                #shrink_items
+            }
+        }
+    }
+    .into()
+}
+
+fn shrink_named(_ident: TokenStream2, fields: &FieldsNamed) -> (TokenStream2, TokenStream2) {
+    let (captures, fields_idents) = generate_exhaustive_destructuring(fields.named.iter());
+    (
+        captures,
+        quote! {
+            {#(
+                turbo_tasks::macro_helpers::ShrinkToFitDerefSpecialization::new(
+                    #fields_idents,
+                ).shrink_to_fit();
+            )*}
+        },
+    )
+}
+
+fn shrink_unnamed(_ident: TokenStream2, fields: &FieldsUnnamed) -> (TokenStream2, TokenStream2) {
+    let (captures, fields_idents) = generate_exhaustive_destructuring(fields.unnamed.iter());
+    (
+        captures,
+        quote! {
+            {#(
+                turbo_tasks::macro_helpers::ShrinkToFitDerefSpecialization::new(
+                    #fields_idents,
+                ).shrink_to_fit();
+            )*}
+        },
+    )
+}
+
+fn shrink_unit(_ident: TokenStream2) -> TokenStream2 {
+    quote! { { } }
+}

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

@@ -22,6 +22,11 @@ pub fn derive_trace_raw_vcs_attr(input: TokenStream) -> TokenStream {
     derive::derive_trace_raw_vcs(input)
 }
 
+#[proc_macro_derive(ShrinkToFit, attributes(turbo_tasks))]
+pub fn derive_shrink_to_fit(input: TokenStream) -> TokenStream {
+    derive::derive_shrink_to_fit(input)
+}
+
 #[proc_macro_derive(ResolvedValue, attributes(turbo_tasks))]
 pub fn derive_resolved_value_attr(input: TokenStream) -> TokenStream {
     derive::derive_resolved_value(input)

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

@@ -16,6 +16,16 @@ pub fn primitive(input: TokenStream) -> TokenStream {
         .into();
     };
 
+    let value_shrink_to_fit_impl = if input.manual_shrink_to_fit.is_none() {
+        Some(quote! {
+            impl turbo_tasks::ShrinkToFit for #ty {
+                fn shrink_to_fit(&mut self) {}
+            }
+        })
+    } else {
+        None
+    };
+
     let value_debug_impl = quote! {
         #[turbo_tasks::value_impl]
         impl turbo_tasks::debug::ValueDebug for #ty {
@@ -62,7 +72,7 @@ pub fn primitive(input: TokenStream) -> TokenStream {
         #value_type_and_register
 
         #value_debug_impl
-
+        #value_shrink_to_fit_impl
         #value_default_impl
     }
     .into()

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

@@ -349,49 +349,46 @@ pub fn value(args: TokenStream, input: TokenStream) -> TokenStream {
         quote! {}
     };
 
-    let derive = match serialization_mode {
+    let mut struct_attributes = vec![quote! {
+        #[derive(turbo_tasks::ShrinkToFit, turbo_tasks::trace::TraceRawVcs)]
+    }];
+    match serialization_mode {
+        SerializationMode::Auto | SerializationMode::AutoForInput => {
+            struct_attributes.push(quote! {
+                #[derive(
+                    turbo_tasks::macro_helpers::serde::Serialize,
+                    turbo_tasks::macro_helpers::serde::Deserialize,
+                )]
+                #[serde(crate = "turbo_tasks::macro_helpers::serde")]
+            })
+        }
         SerializationMode::None | SerializationMode::Custom | SerializationMode::CustomForInput => {
-            quote! {
-                #[derive(turbo_tasks::trace::TraceRawVcs)]
-            }
         }
-        SerializationMode::Auto | SerializationMode::AutoForInput => quote! {
-            #[derive(
-                turbo_tasks::trace::TraceRawVcs,
-                turbo_tasks::macro_helpers::serde::Serialize,
-                turbo_tasks::macro_helpers::serde::Deserialize,
-            )]
-            #[serde(crate = "turbo_tasks::macro_helpers::serde")]
-        },
     };
-    let debug_derive = if inner_type.is_some() {
+    if inner_type.is_some() {
         // Transparent structs have their own manual `ValueDebug` implementation.
-        quote! {
+        struct_attributes.push(quote! {
             #[repr(transparent)]
-        }
+        });
     } else {
-        quote! {
+        struct_attributes.push(quote! {
             #[derive(
                 turbo_tasks::debug::ValueDebugFormat,
                 turbo_tasks::debug::internal::ValueDebug,
             )]
-        }
-    };
-    let eq_derive = if manual_eq {
-        quote!()
-    } else {
-        quote!(
+        });
+    }
+    if !manual_eq {
+        struct_attributes.push(quote! {
             #[derive(PartialEq, Eq)]
-        )
-    };
-    let resolved_derive = if let Some(span) = resolved {
-        quote_spanned!(
+        });
+    }
+    if let Some(span) = resolved {
+        struct_attributes.push(quote_spanned! {
             span =>
             #[derive(turbo_tasks::ResolvedValue)]
-        )
-    } else {
-        quote!()
-    };
+        });
+    }
 
     let new_value_type = match serialization_mode {
         SerializationMode::None => quote! {
@@ -449,10 +446,7 @@ pub fn value(args: TokenStream, input: TokenStream) -> TokenStream {
     );
 
     let expanded = quote! {
-        #derive
-        #debug_derive
-        #eq_derive
-        #resolved_derive
+        #(#struct_attributes)*
         #item
 
         impl #ident {

turbopack/crates/turbo-tasks-memory/tests/shrink_to_fit.rs

@@ -0,0 +1 @@
+../../turbo-tasks-testing/tests/shrink_to_fit.rs

turbopack/crates/turbo-tasks-testing/tests/shrink_to_fit.rs

@@ -0,0 +1,26 @@
+#![feature(arbitrary_self_types)]
+#![feature(arbitrary_self_types_pointers)]
+
+use anyhow::Result;
+use turbo_tasks::Vc;
+use turbo_tasks_testing::{register, run, Registration};
+
+static REGISTRATION: Registration = register!();
+
+#[turbo_tasks::value(transparent)]
+struct Wrapper(Vec<u32>);
+
+#[tokio::test]
+async fn test_shrink_to_fit() -> Result<()> {
+    run(&REGISTRATION, || async {
+        // `Vec::shrink_to_fit` is implicitly called when a cell is constructed.
+        let a: Vc<Wrapper> = Vc::cell(Vec::with_capacity(100));
+        assert_eq!(a.await?.capacity(), 0);
+
+        let b: Vc<Wrapper> = Vc::local_cell(Vec::with_capacity(100));
+        assert_eq!(b.await?.capacity(), 0);
+
+        Ok(())
+    })
+    .await
+}

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

@@ -66,6 +66,7 @@ mod read_ref;
 pub mod registry;
 mod scope;
 mod serialization_invalidation;
+mod shrink_to_fit;
 pub mod small_duration;
 mod state;
 pub mod task;
@@ -111,6 +112,7 @@ pub use read_ref::ReadRef;
 use rustc_hash::FxHasher;
 pub use scope::scope;
 pub use serialization_invalidation::SerializationInvalidator;
+pub use shrink_to_fit::ShrinkToFit;
 pub use state::{State, TransientState};
 pub use task::{task_input::TaskInput, SharedReference, TypedSharedReference};
 pub use trait_ref::{IntoTraitRef, TraitRef};

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

@@ -1,4 +1,6 @@
 //! Runtime helpers for [turbo-tasks-macro].
+use std::ops::{Deref, DerefMut};
+
 pub use async_trait::async_trait;
 pub use once_cell::sync::{Lazy, OnceCell};
 pub use serde;
@@ -9,8 +11,8 @@ pub use super::{
     manager::{find_cell_by_type, notify_scheduled_tasks, spawn_detached_for_testing},
 };
 use crate::{
-    debug::ValueDebugFormatString, task::TaskOutput, RawVc, ResolvedValue, TaskInput,
-    TaskPersistence, Vc,
+    debug::ValueDebugFormatString, shrink_to_fit::ShrinkToFit, task::TaskOutput, RawVc,
+    ResolvedValue, TaskInput, TaskPersistence, Vc,
 };
 
 #[inline(never)]
@@ -56,3 +58,60 @@ macro_rules! stringify_path {
         stringify!($path)
     };
 }
+
+/// A wrapper type that uses the [autoderef specialization hack][autoderef] to call
+/// [`ShrinkToFit::shrink_to_fit`] on types that implement [`ShrinkToFit`].
+///
+/// This uses a a no-op method [`ShrinkToFitFallbackNoop::shrink_to_fit`] on types that do not
+/// implement [`ShrinkToFit`].
+///
+/// This is used by the derive macro for [`ShrinkToFit`], which is called by the
+/// [turbo_tasks::value][crate::value] macro.
+///
+/// [autoderef]: http://lukaskalbertodt.github.io/2019/12/05/generalized-autoref-based-specialization.html
+pub struct ShrinkToFitDerefSpecialization<'a, T> {
+    inner: ShrinkToFitFallbackNoop<'a, T>,
+}
+
+impl<'a, T> ShrinkToFitDerefSpecialization<'a, T> {
+    pub fn new(real: &'a mut T) -> Self {
+        Self {
+            inner: ShrinkToFitFallbackNoop { real },
+        }
+    }
+}
+
+impl<T> ShrinkToFitDerefSpecialization<'_, T>
+where
+    T: ShrinkToFit,
+{
+    pub fn shrink_to_fit(&mut self) {
+        // call the real `ShrinkToFit::shrink_to_fit` method
+        self.inner.real.shrink_to_fit()
+    }
+}
+
+impl<'a, T> Deref for ShrinkToFitDerefSpecialization<'a, T> {
+    type Target = ShrinkToFitFallbackNoop<'a, T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<T> DerefMut for ShrinkToFitDerefSpecialization<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+// Implements `ShrinkToFit` using a no-op `ShrinkToFit::shrink_to_fit` method.
+pub struct ShrinkToFitFallbackNoop<'a, T> {
+    real: &'a mut T,
+}
+
+impl<T> ShrinkToFitFallbackNoop<'_, T> {
+    /// A no-op function called as part of [`ShrinkToFitDerefSpecialization`] when `T` does not
+    /// implement [`ShrinkToFit`].
+    pub fn shrink_to_fit(&mut self) {}
+}

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

@@ -8,11 +8,11 @@ use crate::{
 };
 
 __turbo_tasks_internal_primitive!(());
-__turbo_tasks_internal_primitive!(String);
+__turbo_tasks_internal_primitive!(String, manual_shrink_to_fit);
 __turbo_tasks_internal_primitive!(RcStr);
 __turbo_tasks_internal_primitive!(Option<String>);
 __turbo_tasks_internal_primitive!(Option<RcStr>);
-__turbo_tasks_internal_primitive!(Vec<RcStr>);
+__turbo_tasks_internal_primitive!(Vec<RcStr>, manual_shrink_to_fit);
 __turbo_tasks_internal_primitive!(Option<u16>);
 __turbo_tasks_internal_primitive!(Option<u64>);
 __turbo_tasks_internal_primitive!(bool);
@@ -30,8 +30,8 @@ __turbo_tasks_internal_primitive!(usize);
 __turbo_tasks_internal_primitive!(isize);
 __turbo_tasks_internal_primitive!(serde_json::Value);
 __turbo_tasks_internal_primitive!(Duration);
-__turbo_tasks_internal_primitive!(Vec<u8>);
-__turbo_tasks_internal_primitive!(Vec<bool>);
+__turbo_tasks_internal_primitive!(Vec<u8>, manual_shrink_to_fit);
+__turbo_tasks_internal_primitive!(Vec<bool>, manual_shrink_to_fit);
 
 #[turbo_tasks::value(transparent, eq = "manual")]
 #[derive(Debug, Clone)]