feat(turbopack): Implement module walking for side effect optimization

turbopack/crates/turbopack-core/src/reference/mod.rs

@@ -8,6 +8,7 @@ use turbo_tasks::{
 };
 
 use crate::{
+    chunk::{ChunkableModuleReference, ChunkingType, ChunkingTypeOption},
     issue::IssueDescriptionExt,
     module::{Module, Modules},
     output::{OutputAsset, OutputAssets},
@@ -84,6 +85,44 @@ impl SingleModuleReference {
     }
 }
 
+#[turbo_tasks::value]
+pub struct SingleChunkableModuleReference {
+    asset: ResolvedVc<Box<dyn Module>>,
+    description: Vc<RcStr>,
+}
+
+#[turbo_tasks::value_impl]
+impl SingleChunkableModuleReference {
+    #[turbo_tasks::function]
+    pub fn new(asset: ResolvedVc<Box<dyn Module>>, description: Vc<RcStr>) -> Vc<Self> {
+        Self::cell(SingleChunkableModuleReference { asset, description })
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ChunkableModuleReference for SingleChunkableModuleReference {
+    #[turbo_tasks::function]
+    fn chunking_type(self: Vc<Self>) -> Vc<ChunkingTypeOption> {
+        Vc::cell(Some(ChunkingType::ParallelInheritAsync))
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ModuleReference for SingleChunkableModuleReference {
+    #[turbo_tasks::function]
+    fn resolve_reference(&self) -> Vc<ModuleResolveResult> {
+        ModuleResolveResult::module(self.asset).cell()
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ValueToString for SingleChunkableModuleReference {
+    #[turbo_tasks::function]
+    fn to_string(&self) -> Vc<RcStr> {
+        self.description
+    }
+}
+
 /// A reference that always resolves to a single module.
 #[turbo_tasks::value]
 pub struct SingleOutputAssetReference {

turbopack/crates/turbopack-ecmascript/src/references/esm/export.rs

@@ -128,10 +128,12 @@ pub async fn follow_reexports(
     module: Vc<Box<dyn EcmascriptChunkPlaceable>>,
     export_name: RcStr,
     side_effect_free_packages: Vc<Glob>,
+    ignore_side_effect_of_entry: bool,
 ) -> Result<Vc<FollowExportsResult>> {
-    if !*module
-        .is_marked_as_side_effect_free(side_effect_free_packages)
-        .await?
+    if !ignore_side_effect_of_entry
+        && !*module
+            .is_marked_as_side_effect_free(side_effect_free_packages)
+            .await?
     {
         return Ok(FollowExportsResult::cell(FollowExportsResult {
             module,

turbopack/crates/turbopack-ecmascript/src/references/esm/mod.rs

@@ -11,7 +11,7 @@ pub use self::{
     base::EsmAssetReference,
     binding::EsmBinding,
     dynamic::EsmAsyncAssetReference,
-    export::{EsmExport, EsmExports},
+    export::{EsmExport, EsmExports, FoundExportType},
     meta::{ImportMetaBinding, ImportMetaRef},
     module_item::EsmModuleItem,
     url::{UrlAssetReference, UrlRewriteBehavior},

turbopack/crates/turbopack-ecmascript/src/tree_shake/asset.rs

@@ -1,12 +1,15 @@
 use anyhow::{Context, Result};
+use turbo_rcstr::RcStr;
 use turbo_tasks::{ResolvedVc, Vc};
+use turbo_tasks_fs::glob::Glob;
 use turbopack_core::{
     asset::{Asset, AssetContent},
     chunk::{AsyncModuleInfo, ChunkableModule, ChunkingContext, EvaluatableAsset},
+    context::AssetContext,
     ident::AssetIdent,
     module::Module,
     reference::{ModuleReference, ModuleReferences, SingleModuleReference},
-    resolve::ModulePart,
+    resolve::{origin::ResolveOrigin, ModulePart},
 };
 
 use super::{
@@ -16,7 +19,11 @@ use super::{
 use crate::{
     chunk::{EcmascriptChunkPlaceable, EcmascriptExports},
     parse::ParseResult,
-    references::analyse_ecmascript_module,
+    references::{
+        analyse_ecmascript_module, esm::FoundExportType, follow_reexports, FollowExportsResult,
+    },
+    side_effect_optimization::facade::module::EcmascriptModuleFacadeModule,
+    tree_shake::{side_effect_module::SideEffectsModule, Key},
     AnalyzeEcmascriptModuleResult, EcmascriptAnalyzable, EcmascriptModuleAsset,
     EcmascriptModuleAssetType, EcmascriptModuleContent, EcmascriptParsable,
 };
@@ -93,15 +100,70 @@ impl EcmascriptModulePartAsset {
     #[turbo_tasks::function]
     pub async fn select_part(
         module: Vc<EcmascriptModuleAsset>,
-        part: Vc<ModulePart>,
+        part: ResolvedVc<ModulePart>,
     ) -> Result<Vc<Box<dyn Module>>> {
-        let split_result = split_module(module).await?;
+        let SplitResult::Ok { entrypoints, .. } = &*split_module(module).await? else {
+            return Ok(Vc::upcast(module));
+        };
 
-        Ok(if matches!(&*split_result, SplitResult::Failed { .. }) {
-            Vc::upcast(module)
-        } else {
-            Vc::upcast(EcmascriptModulePartAsset::new(module, part))
-        })
+        // We follow reexports here
+        if let ModulePart::Export(export) = &*part.await? {
+            let export_name = export.await?.clone_value();
+
+            // If a local binding or reexport with the same name exists, we stop here.
+            // Side effects of the barrel file are preserved.
+            if entrypoints.contains_key(&Key::Export(export_name.clone())) {
+                return Ok(Vc::upcast(EcmascriptModulePartAsset::new(module, *part)));
+            }
+
+            let side_effect_free_packages = module.asset_context().side_effect_free_packages();
+
+            // Exclude local bindings by using exports module part.
+            let source_module = Vc::upcast(module);
+
+            let FollowExportsWithSideEffectsResult {
+                side_effects,
+                result,
+            } = &*follow_reexports_with_side_effects(
+                source_module,
+                export_name.clone(),
+                side_effect_free_packages,
+            )
+            .await?;
+
+            let FollowExportsResult {
+                module: final_module,
+                export_name: new_export,
+                ..
+            } = &*result.await?;
+
+            let final_module = if let Some(new_export) = new_export {
+                if *new_export == export_name {
+                    *final_module
+                } else {
+                    Vc::upcast(EcmascriptModuleFacadeModule::new(
+                        *final_module,
+                        ModulePart::renamed_export(new_export.clone(), export_name.clone()),
+                    ))
+                }
+            } else {
+                Vc::upcast(EcmascriptModuleFacadeModule::new(
+                    *final_module,
+                    ModulePart::renamed_namespace(export_name.clone()),
+                ))
+            };
+
+            if side_effects.is_empty() {
+                return Ok(Vc::upcast(final_module));
+            }
+
+            let side_effects_module =
+                SideEffectsModule::new(module, *part, final_module, side_effects.to_vec());
+
+            return Ok(Vc::upcast(side_effects_module));
+        }
+
+        Ok(Vc::upcast(EcmascriptModulePartAsset::new(module, *part)))
     }
 
     #[turbo_tasks::function]
@@ -117,6 +179,61 @@ impl EcmascriptModulePartAsset {
     }
 }
 
+#[turbo_tasks::value]
+struct FollowExportsWithSideEffectsResult {
+    side_effects: Vec<Vc<Box<dyn EcmascriptChunkPlaceable>>>,
+    result: Vc<FollowExportsResult>,
+}
+
+#[turbo_tasks::function]
+async fn follow_reexports_with_side_effects(
+    module: Vc<Box<dyn EcmascriptChunkPlaceable>>,
+    export_name: RcStr,
+    side_effect_free_packages: Vc<Glob>,
+) -> Result<Vc<FollowExportsWithSideEffectsResult>> {
+    let mut side_effects = vec![];
+
+    let mut current_module = module;
+    let mut current_export_name = export_name;
+    let result = loop {
+        let is_side_effect_free = *current_module
+            .is_marked_as_side_effect_free(side_effect_free_packages)
+            .await?;
+
+        if !is_side_effect_free {
+            side_effects.push(only_effects(current_module));
+        }
+
+        // We ignore the side effect of the entry module here, because we need to proceed.
+        let result = follow_reexports(
+            current_module,
+            current_export_name.clone(),
+            side_effect_free_packages,
+            true,
+        );
+
+        let FollowExportsResult {
+            module,
+            export_name,
+            ty,
+        } = &*result.await?;
+
+        match ty {
+            FoundExportType::SideEffects => {
+                current_module = *module;
+                current_export_name = export_name.clone().unwrap_or(current_export_name);
+            }
+            _ => break result,
+        }
+    };
+
+    Ok(FollowExportsWithSideEffectsResult {
+        side_effects,
+        result,
+    }
+    .cell())
+}
+
 #[turbo_tasks::value_impl]
 impl Module for EcmascriptModulePartAsset {
     #[turbo_tasks::function]
@@ -199,6 +316,21 @@ impl EcmascriptChunkPlaceable for EcmascriptModulePartAsset {
     async fn get_exports(self: Vc<Self>) -> Result<Vc<EcmascriptExports>> {
         Ok(*self.analyze().await?.exports)
     }
+
+    #[turbo_tasks::function]
+    async fn is_marked_as_side_effect_free(
+        self: Vc<Self>,
+        side_effect_free_packages: Vc<Glob>,
+    ) -> Result<Vc<bool>> {
+        let this = self.await?;
+
+        match *this.part.await? {
+            ModulePart::Exports | ModulePart::Export(..) => Ok(Vc::cell(true)),
+            _ => Ok(this
+                .full_module
+                .is_marked_as_side_effect_free(side_effect_free_packages)),
+        }
+    }
 }
 
 #[turbo_tasks::value_impl]
@@ -236,3 +368,15 @@ fn analyze(
 
 #[turbo_tasks::value_impl]
 impl EvaluatableAsset for EcmascriptModulePartAsset {}
+
+#[turbo_tasks::function]
+async fn only_effects(
+    module: Vc<Box<dyn EcmascriptChunkPlaceable>>,
+) -> Result<Vc<Box<dyn EcmascriptChunkPlaceable>>> {
+    if let Some(module) = Vc::try_resolve_downcast_type::<EcmascriptModuleAsset>(module).await? {
+        let module = EcmascriptModulePartAsset::new(module, ModulePart::evaluation());
+        return Ok(Vc::upcast(module));
+    }
+
+    Ok(module)
+}