feat: Enable tree shaking of turbopack (#66689)

### What?

Enable tree shaking of turbopack.

### Why?

Verify the implementation and improve the bundle size.

### How?

Closes PACK-3111

---------

Co-authored-by: Tobias Koppers <tobias.koppers@googlemail.com>

crates/next-api/src/pages.rs

@@ -37,7 +37,7 @@ use turbopack_binding::{
             asset::AssetContent,
             chunk::{
                 availability_info::AvailabilityInfo, ChunkingContext, ChunkingContextExt,
-                EntryChunkGroupResult, EvaluatableAssets,
+                EntryChunkGroupResult, EvaluatableAsset, EvaluatableAssets,
             },
             context::AssetContext,
             file_source::FileSource,
@@ -51,7 +51,7 @@ use turbopack_binding::{
             source::Source,
             virtual_output::VirtualOutputAsset,
         },
-        ecmascript::{resolve::esm_resolve, EcmascriptModuleAsset},
+        ecmascript::resolve::esm_resolve,
         nodejs::NodeJsChunkingContext,
         turbopack::{
             module_options::ModuleOptionsContext,
@@ -633,7 +633,7 @@ impl PageEndpoint {
             );
 
             let Some(client_module) =
-                Vc::try_resolve_downcast_type::<EcmascriptModuleAsset>(client_module).await?
+                Vc::try_resolve_sidecast::<Box<dyn EvaluatableAsset>>(client_module).await?
             else {
                 bail!("expected an ECMAScript module asset");
             };
@@ -659,7 +659,7 @@ impl PageEndpoint {
             .context("expected Next.js client runtime to resolve to a module")?;
 
             let Some(client_main_module) =
-                Vc::try_resolve_downcast_type::<EcmascriptModuleAsset>(client_main_module).await?
+                Vc::try_resolve_sidecast::<Box<dyn EvaluatableAsset>>(client_main_module).await?
             else {
                 bail!("expected an ECMAScript module asset");
             };
@@ -670,8 +670,8 @@ impl PageEndpoint {
                 client_module.ident(),
                 this.pages_project
                     .client_runtime_entries()
-                    .with_entry(Vc::upcast(client_main_module))
-                    .with_entry(Vc::upcast(client_module)),
+                    .with_entry(client_main_module)
+                    .with_entry(client_module),
                 Value::new(AvailabilityInfo::Root),
             );
 

crates/next-api/src/server_actions.rs

@@ -29,7 +29,8 @@ use turbopack_binding::{
             virtual_source::VirtualSource,
         },
         ecmascript::{
-            chunk::EcmascriptChunkPlaceable, parse::ParseResult, EcmascriptModuleAsset,
+            chunk::EcmascriptChunkPlaceable, parse::ParseResult,
+            tree_shake::asset::EcmascriptModulePartAsset, EcmascriptModuleAsset,
             EcmascriptModuleAssetType,
         },
     },
@@ -246,6 +247,14 @@ async fn to_rsc_context(
         } else {
             ReferenceType::TypeScript(TypeScriptReferenceSubType::Undefined)
         }
+    } else if let Some(module) =
+        Vc::try_resolve_downcast_type::<EcmascriptModulePartAsset>(module).await?
+    {
+        if module.await?.full_module.await?.ty == EcmascriptModuleAssetType::Ecmascript {
+            ReferenceType::EcmaScriptModules(EcmaScriptModulesReferenceSubType::Undefined)
+        } else {
+            ReferenceType::TypeScript(TypeScriptReferenceSubType::Undefined)
+        }
     } else {
         ReferenceType::TypeScript(TypeScriptReferenceSubType::Undefined)
     };
@@ -294,14 +303,21 @@ pub fn parse_server_actions<C: Comments>(
 /// the exported action function. If not, we return a None.
 #[turbo_tasks::function]
 async fn parse_actions(module: Vc<Box<dyn Module>>) -> Result<Vc<OptionActionMap>> {
-    let Some(ecmascript_asset) =
+    let parsed = if let Some(ecmascript_asset) =
         Vc::try_resolve_downcast_type::<EcmascriptModuleAsset>(module).await?
-    else {
+    {
+        ecmascript_asset.failsafe_parse()
+    } else if let Some(ecmascript_asset) =
+        Vc::try_resolve_downcast_type::<EcmascriptModulePartAsset>(module).await?
+    {
+        ecmascript_asset.await?.full_module.failsafe_parse()
+    } else {
         return Ok(OptionActionMap::none());
     };
+
     let ParseResult::Ok {
         comments, program, ..
-    } = &*ecmascript_asset.failsafe_parse().await?
+    } = &*parsed.await?
     else {
         // The file might be be parse-able, but this is reported separately.
         return Ok(OptionActionMap::none());

crates/next-core/src/next_client/context.rs

@@ -18,7 +18,6 @@ use turbopack_binding::{
             free_var_references,
             resolve::{parse::Request, pattern::Pattern},
         },
-        ecmascript::TreeShakingMode,
         node::{
             execution_context::ExecutionContext,
             transforms::postcss::{PostCssConfigLocation, PostCssTransformOptions},
@@ -47,6 +46,7 @@ use crate::{
         get_next_client_resolved_map,
     },
     next_shared::{
+        next_js_special_exports,
         resolve::{
             get_invalid_server_only_resolve_plugin, ModuleFeatureReportResolvePlugin,
             NextSharedRuntimeResolvePlugin,
@@ -203,6 +203,8 @@ pub async fn get_client_module_options_context(
     mode: Vc<NextMode>,
     next_config: Vc<NextConfig>,
 ) -> Result<Vc<ModuleOptionsContext>> {
+    let next_mode = mode.await?;
+
     let resolve_options_context =
         get_client_resolve_options_context(project_path, ty, mode, next_config, execution_context);
 
@@ -238,6 +240,12 @@ pub async fn get_client_module_options_context(
     let enable_webpack_loaders =
         webpack_loader_options(project_path, next_config, false, conditions).await?;
 
+    let tree_shaking_mode_for_user_code = *next_config
+        .tree_shaking_mode_for_user_code(next_mode.is_development())
+        .await?;
+    let tree_shaking_mode_for_foreign_code = *next_config
+        .tree_shaking_mode_for_foreign_code(next_mode.is_development())
+        .await?;
     let use_swc_css = *next_config.use_swc_css().await?;
     let target_browsers = env.runtime_versions();
 
@@ -278,7 +286,8 @@ pub async fn get_client_module_options_context(
         enable_typeof_window_inlining: Some(TypeofWindow::Object),
         preset_env_versions: Some(env),
         execution_context: Some(execution_context),
-        tree_shaking_mode: Some(TreeShakingMode::ReexportsOnly),
+        tree_shaking_mode: tree_shaking_mode_for_user_code,
+        special_exports: Some(next_js_special_exports()),
         enable_postcss_transform,
         side_effect_free_packages: next_config.optimize_package_imports().await?.clone_value(),
         ..Default::default()
@@ -290,6 +299,7 @@ pub async fn get_client_module_options_context(
         enable_webpack_loaders: foreign_enable_webpack_loaders,
         enable_postcss_transform: enable_foreign_postcss_transform,
         custom_rules: foreign_next_client_rules,
+        tree_shaking_mode: tree_shaking_mode_for_foreign_code,
         // NOTE(WEB-1016) PostCSS transforms should also apply to foreign code.
         ..module_options_context.clone()
     };

crates/next-core/src/next_client_reference/ecmascript_client_reference/ecmascript_client_reference_proxy_module.rs

@@ -22,7 +22,6 @@ use turbopack_binding::turbopack::{
             EcmascriptChunkType, EcmascriptExports,
         },
         utils::StringifyJs,
-        EcmascriptModuleAsset,
     },
 };
 
@@ -65,7 +64,7 @@ impl EcmascriptClientReferenceProxyModule {
     }
 
     #[turbo_tasks::function]
-    async fn proxy_module(&self) -> Result<Vc<EcmascriptModuleAsset>> {
+    async fn proxy_module(&self) -> Result<Vc<Box<dyn EcmascriptChunkPlaceable>>> {
         let mut code = CodeBuilder::default();
 
         let server_module_path = &*self.server_module_ident.path().to_string().await?;
@@ -156,7 +155,7 @@ impl EcmascriptClientReferenceProxyModule {
             .module();
 
         let Some(proxy_module) =
-            Vc::try_resolve_downcast_type::<EcmascriptModuleAsset>(proxy_module).await?
+            Vc::try_resolve_sidecast::<Box<dyn EcmascriptChunkPlaceable>>(proxy_module).await?
         else {
             bail!("proxy asset is not an ecmascript module");
         };

crates/next-core/src/next_config.rs

@@ -12,6 +12,7 @@ use turbopack_binding::{
             issue::{Issue, IssueSeverity, IssueStage, OptionStyledString, StyledString},
             resolve::ResolveAliasMap,
         },
+        ecmascript::{OptionTreeShaking, TreeShakingMode},
         ecmascript_plugin::transform::{
             emotion::EmotionTransformConfig, relay::RelayConfig,
             styled_components::StyledComponentsTransformConfig,
@@ -580,6 +581,8 @@ pub struct ExperimentalConfig {
     /// (doesn't apply to Turbopack).
     webpack_build_worker: Option<bool>,
     worker_threads: Option<bool>,
+
+    tree_shaking: Option<bool>,
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, TraceRawVcs)]
@@ -1116,6 +1119,35 @@ impl NextConfig {
                 .unwrap_or_default(),
         ))
     }
+
+    #[turbo_tasks::function]
+    pub async fn tree_shaking_mode_for_foreign_code(
+        self: Vc<Self>,
+        is_development: bool,
+    ) -> Result<Vc<OptionTreeShaking>> {
+        let tree_shaking = self.await?.experimental.tree_shaking;
+
+        Ok(OptionTreeShaking(match tree_shaking {
+            Some(false) => Some(TreeShakingMode::ReexportsOnly),
+            Some(true) => Some(TreeShakingMode::ModuleFragments),
+            None => {
+                if is_development {
+                    Some(TreeShakingMode::ReexportsOnly)
+                } else {
+                    Some(TreeShakingMode::ModuleFragments)
+                }
+            }
+        })
+        .cell())
+    }
+
+    #[turbo_tasks::function]
+    pub async fn tree_shaking_mode_for_user_code(
+        self: Vc<Self>,
+        _is_development: bool,
+    ) -> Result<Vc<OptionTreeShaking>> {
+        Ok(Vc::cell(Some(TreeShakingMode::ReexportsOnly)))
+    }
 }
 
 /// A subset of ts/jsconfig that next.js implicitly

crates/next-core/src/next_server/context.rs

@@ -53,6 +53,7 @@ use crate::{
     next_import_map::get_next_server_import_map,
     next_server::resolve::ExternalPredicate,
     next_shared::{
+        next_js_special_exports,
         resolve::{
             get_invalid_client_only_resolve_plugin, get_invalid_styled_jsx_resolve_plugin,
             ModuleFeatureReportResolvePlugin, NextExternalResolvePlugin,
@@ -386,6 +387,7 @@ pub async fn get_server_module_options_context(
     next_config: Vc<NextConfig>,
     next_runtime: NextRuntime,
 ) -> Result<Vc<ModuleOptionsContext>> {
+    let next_mode = mode.await?;
     let mut next_server_rules =
         get_next_server_transforms_rules(next_config, ty.into_value(), mode, false, next_runtime)
             .await?;
@@ -444,6 +446,12 @@ pub async fn get_server_module_options_context(
     let enable_webpack_loaders =
         webpack_loader_options(project_path, next_config, false, conditions).await?;
 
+    let tree_shaking_mode_for_user_code = *next_config
+        .tree_shaking_mode_for_user_code(next_mode.is_development())
+        .await?;
+    let tree_shaking_mode_for_foreign_code = *next_config
+        .tree_shaking_mode_for_foreign_code(next_mode.is_development())
+        .await?;
     let use_swc_css = *next_config.use_swc_css().await?;
     let versions = RuntimeVersions(Default::default()).cell();
 
@@ -486,8 +494,9 @@ pub async fn get_server_module_options_context(
         enable_typeof_window_inlining: Some(TypeofWindow::Undefined),
         execution_context: Some(execution_context),
         use_swc_css,
-        tree_shaking_mode: Some(TreeShakingMode::ReexportsOnly),
+        tree_shaking_mode: tree_shaking_mode_for_user_code,
         import_externals: *next_config.import_externals().await?,
+        special_exports: Some(next_js_special_exports()),
         ignore_dynamic_requests: true,
         side_effect_free_packages: next_config.optimize_package_imports().await?.clone_value(),
         ..Default::default()
@@ -537,6 +546,7 @@ pub async fn get_server_module_options_context(
                 enable_webpack_loaders: foreign_enable_webpack_loaders,
                 // NOTE(WEB-1016) PostCSS transforms should also apply to foreign code.
                 enable_postcss_transform: enable_foreign_postcss_transform,
+                tree_shaking_mode: tree_shaking_mode_for_foreign_code,
                 ..module_options_context.clone()
             };
 
@@ -592,6 +602,7 @@ pub async fn get_server_module_options_context(
                 enable_webpack_loaders: foreign_enable_webpack_loaders,
                 // NOTE(WEB-1016) PostCSS transforms should also apply to foreign code.
                 enable_postcss_transform: enable_foreign_postcss_transform,
+                tree_shaking_mode: tree_shaking_mode_for_foreign_code,
                 ..module_options_context.clone()
             };
             let internal_module_options_context = ModuleOptionsContext {
@@ -660,6 +671,7 @@ pub async fn get_server_module_options_context(
                 enable_webpack_loaders: foreign_enable_webpack_loaders,
                 // NOTE(WEB-1016) PostCSS transforms should also apply to foreign code.
                 enable_postcss_transform: enable_foreign_postcss_transform,
+                tree_shaking_mode: tree_shaking_mode_for_foreign_code,
                 ..module_options_context.clone()
             };
             let internal_module_options_context = ModuleOptionsContext {
@@ -723,6 +735,7 @@ pub async fn get_server_module_options_context(
                 enable_webpack_loaders: foreign_enable_webpack_loaders,
                 // NOTE(WEB-1016) PostCSS transforms should also apply to foreign code.
                 enable_postcss_transform: enable_foreign_postcss_transform,
+                tree_shaking_mode: tree_shaking_mode_for_foreign_code,
                 ..module_options_context.clone()
             };
             let internal_module_options_context = ModuleOptionsContext {
@@ -803,6 +816,7 @@ pub async fn get_server_module_options_context(
                 enable_webpack_loaders: foreign_enable_webpack_loaders,
                 // NOTE(WEB-1016) PostCSS transforms should also apply to foreign code.
                 enable_postcss_transform: enable_foreign_postcss_transform,
+                tree_shaking_mode: tree_shaking_mode_for_foreign_code,
                 ..module_options_context.clone()
             };
             let internal_module_options_context = ModuleOptionsContext {
@@ -841,7 +855,7 @@ pub async fn get_server_module_options_context(
 pub fn get_build_module_options_context() -> Vc<ModuleOptionsContext> {
     ModuleOptionsContext {
         enable_typescript_transform: Some(Default::default()),
-        tree_shaking_mode: Some(TreeShakingMode::ReexportsOnly),
+        tree_shaking_mode: Some(TreeShakingMode::ModuleFragments),
         esm_url_rewrite_behavior: Some(UrlRewriteBehavior::Full),
         ..Default::default()
     }

crates/next-core/src/next_shared/mod.rs

@@ -1,3 +1,31 @@
+use turbo_tasks::{RcStr, Vc};
+
 pub(crate) mod resolve;
 pub(crate) mod transforms;
 pub(crate) mod webpack_rules;
+
+#[turbo_tasks::function]
+pub fn next_js_special_exports() -> Vc<Vec<RcStr>> {
+    Vc::cell(
+        [
+            "config",
+            "middleware",
+            "runtime",
+            "revalidate",
+            "dynamic",
+            "dynamicParams",
+            "fetchCache",
+            "preferredRegion",
+            "maxDuration",
+            "generateStaticParams",
+            "metadata",
+            "generateMetadata",
+            "getServerSideProps",
+            "getInitialProps",
+            "getStaticProps",
+        ]
+        .into_iter()
+        .map(RcStr::from)
+        .collect::<Vec<RcStr>>(),
+    )
+}

packages/next/src/server/config-schema.ts

@@ -389,6 +389,7 @@ export const configSchema: zod.ZodType<NextConfig> = z.lazy(() =>
               .optional(),
             resolveExtensions: z.array(z.string()).optional(),
             useSwcCss: z.boolean().optional(),
+            treeShaking: z.boolean().optional(),
             memoryLimit: z.number().optional(),
           })
           .optional(),