Turbopack: next/dynamic layout segment optimization (vercel#70708)

Closes PACK-3274

This deduplicates work by only searching for `next/dynamic` once per
layout segment.

- `collect_next_dynamic_imports` is a turbotask now, executed once per
layout segment
- `collect_next_dynamic_imports` also takes a list of visited modules,
to prevent revisiting

Before
<img width="1722" alt="Bildschirmfoto 2024-10-01 um 15 28 34"
src="https://github.com/user-attachments/assets/0db86796-289e-4d8a-95b9-4164bcaabc71">

After
<img width="1728" alt="Bildschirmfoto 2024-10-02 um 15 17 10"
src="https://github.com/user-attachments/assets/bfb99b03-26d8-4c0f-b97e-f9e437f7fc4c">

crates/next-api/src/app.rs

@@ -1,5 +1,5 @@
 use anyhow::{Context, Result};
-use indexmap::IndexSet;
+use indexmap::{IndexMap, IndexSet};
 use next_core::{
     all_assets_from_entries,
     app_segment_config::NextSegmentConfig,
@@ -18,9 +18,7 @@ use next_core::{
         get_client_module_options_context, get_client_resolve_options_context,
         get_client_runtime_entries, ClientContextType, RuntimeEntries,
     },
-    next_client_reference::{
-        client_reference_graph, ClientReferenceType, NextEcmascriptClientReferenceTransition,
-    },
+    next_client_reference::{client_reference_graph, NextEcmascriptClientReferenceTransition},
     next_config::NextConfig,
     next_dynamic::NextDynamicTransition,
     next_edge::route_regex::get_named_middleware_regex,
@@ -69,6 +67,7 @@ use turbopack_ecmascript::resolve::cjs_resolve;
 use crate::{
     dynamic_imports::{
         collect_chunk_group, collect_evaluated_chunk_group, collect_next_dynamic_imports,
+        VisitedDynamicImportModules,
     },
     font::create_font_manifest,
     loadable_manifest::create_react_loadable_manifest,
@@ -864,7 +863,6 @@ impl AppEndpoint {
                 let client_shared_availability_info = client_shared_chunk_group.availability_info;
 
                 let client_references = client_reference_graph(Vc::cell(vec![rsc_entry_asset]));
-                let client_reference_types = client_references.types();
 
                 let ssr_chunking_context = if process_ssr {
                     Some(match runtime {
@@ -880,21 +878,32 @@ impl AppEndpoint {
                     None
                 };
 
-                let client_dynamic_imports = collect_next_dynamic_imports(
-                    client_references
+                let client_dynamic_imports = {
+                    let mut client_dynamic_imports = IndexMap::new();
+                    let mut visited_modules = VisitedDynamicImportModules::empty();
+
+                    for refs in client_references
                         .await?
-                        .client_references
-                        .iter()
-                        .filter_map(|r| match r.ty() {
-                            ClientReferenceType::EcmascriptClientReference(entry) => Some(entry),
-                            ClientReferenceType::CssClientReference(_) => None,
-                        })
-                        .map(|entry| async move { Ok(Vc::upcast(entry.await?.ssr_module)) })
-                        .try_join()
-                        .await?,
-                    Vc::upcast(this.app_project.client_module_context()),
-                )
-                .await?;
+                        .client_references_by_server_component
+                        .values()
+                    {
+                        let result = collect_next_dynamic_imports(
+                            refs.clone(),
+                            Vc::upcast(this.app_project.client_module_context()),
+                            visited_modules,
+                        )
+                        .await?;
+                        client_dynamic_imports.extend(
+                            result
+                                .client_dynamic_imports
+                                .iter()
+                                .map(|(k, v)| (*k, v.clone())),
+                        );
+                        visited_modules = result.visited_modules;
+                    }
+
+                    client_dynamic_imports
+                };
 
                 let client_references_chunks = get_app_client_references_chunks(
                     client_references,
@@ -1025,7 +1034,7 @@ impl AppEndpoint {
                 }
 
                 (
-                    Some(get_app_server_reference_modules(client_reference_types)),
+                    Some(get_app_server_reference_modules(client_references.types())),
                     Some(client_dynamic_imports),
                     Some(client_references),
                 )
@@ -1199,10 +1208,13 @@ impl AppEndpoint {
 
                 // create react-loadable-manifest for next/dynamic
                 let mut dynamic_import_modules = collect_next_dynamic_imports(
-                    [Vc::upcast(app_entry.rsc_entry)],
+                    vec![Vc::upcast(app_entry.rsc_entry)],
                     Vc::upcast(this.app_project.client_module_context()),
+                    VisitedDynamicImportModules::empty(),
                 )
-                .await?;
+                .await?
+                .client_dynamic_imports
+                .clone();
                 dynamic_import_modules.extend(client_dynamic_imports.into_iter().flatten());
                 let dynamic_import_entries = collect_evaluated_chunk_group(
                     Vc::upcast(client_chunking_context),
@@ -1351,10 +1363,13 @@ impl AppEndpoint {
                 // create react-loadable-manifest for next/dynamic
                 let availability_info = Value::new(AvailabilityInfo::Root);
                 let mut dynamic_import_modules = collect_next_dynamic_imports(
-                    [Vc::upcast(app_entry.rsc_entry)],
+                    vec![Vc::upcast(app_entry.rsc_entry)],
                     Vc::upcast(this.app_project.client_module_context()),
+                    VisitedDynamicImportModules::empty(),
                 )
-                .await?;
+                .await?
+                .client_dynamic_imports
+                .clone();
                 dynamic_import_modules.extend(client_dynamic_imports.into_iter().flatten());
                 let dynamic_import_entries = collect_chunk_group(
                     Vc::upcast(client_chunking_context),

crates/next-api/src/dynamic_imports.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use anyhow::{bail, Result};
 use futures::Future;
@@ -8,9 +8,9 @@ use swc_core::ecma::{
     ast::{CallExpr, Callee, Expr, Ident, Lit},
     visit::{Visit, VisitWith},
 };
-use tracing::Level;
+use tracing::{Instrument, Level};
 use turbo_tasks::{
-    graph::{GraphTraversal, NonDeterministic, VisitControlFlow},
+    graph::{GraphTraversal, NonDeterministic, VisitControlFlow, VisitedNodes},
     trace::TraceRawVcs,
     RcStr, ReadRef, TryJoinIterExt, Value, ValueToString, Vc,
 };
@@ -103,6 +103,23 @@ pub(crate) async fn collect_evaluated_chunk_group(
     .await
 }
 
+#[turbo_tasks::value(shared)]
+pub struct NextDynamicImportsResult {
+    pub client_dynamic_imports: IndexMap<Vc<Box<dyn Module>>, DynamicImportedModules>,
+    pub visited_modules: Vc<VisitedDynamicImportModules>,
+}
+
+#[turbo_tasks::value(shared)]
+pub struct VisitedDynamicImportModules(HashSet<NextDynamicVisitEntry>);
+
+#[turbo_tasks::value_impl]
+impl VisitedDynamicImportModules {
+    #[turbo_tasks::function]
+    pub fn empty() -> Vc<Self> {
+        VisitedDynamicImportModules(Default::default()).cell()
+    }
+}
+
 /// Returns a mapping of the dynamic imports for each module, if the import is
 /// wrapped in `next/dynamic`'s `dynamic()`. Refer [documentation](https://nextjs.org/docs/pages/building-your-application/optimizing/lazy-loading#with-named-exports) for the usecases.
 ///
@@ -125,60 +142,71 @@ pub(crate) async fn collect_evaluated_chunk_group(
 ///    - Loadable runtime [injects preload fn](https://github.com/vercel/next.js/blob/ad42b610c25b72561ad367b82b1c7383fd2a5dd2/packages/next/src/shared/lib/loadable.shared-runtime.tsx#L281)
 ///      to wait until all the dynamic components are being loaded, this ensures hydration mismatch
 ///      won't occur
-#[tracing::instrument(level = Level::INFO, name = "collecting next/dynamic imports", skip_all)]
+#[turbo_tasks::function]
 pub(crate) async fn collect_next_dynamic_imports(
-    server_entries: impl IntoIterator<Item = Vc<Box<dyn Module>>>,
+    // `server_entries` cannot be a `Vc<Vec<_>>` because that would compare by cell identity and
+    // not by value, breaking memoization.
+    server_entries: Vec<Vc<Box<dyn Module>>>,
     client_asset_context: Vc<Box<dyn AssetContext>>,
-) -> Result<IndexMap<Vc<Box<dyn Module>>, DynamicImportedModules>> {
-    // Traverse referenced modules graph, collect all of the dynamic imports:
-    // - Read the Program AST of the Module, this is the origin (A)
-    //  - If there's `dynamic(import(B))`, then B is the module that is being imported
-    // Returned import mappings are in the form of
-    // (Module<A>, Vec<(B, Module<B>)>) (where B is the raw import source string,
-    // and Module<B> is the actual resolved Module)
-    let imported_modules_mapping = NonDeterministic::new()
-        .skip_duplicates()
-        .visit(
-            server_entries
-                .into_iter()
-                .map(|module| async move {
-                    Ok(NextDynamicVisitEntry::Module(
-                        module.resolve().await?,
-                        module.ident().to_string().await?,
-                    ))
-                })
-                .try_join()
-                .await?
-                .into_iter(),
-            NextDynamicVisit {
-                client_asset_context: client_asset_context.resolve().await?,
-            },
-        )
-        .await
-        .completed()?
-        .into_inner()
-        .into_iter()
-        .filter_map(|entry| {
+    visited_modules: Vc<VisitedDynamicImportModules>,
+) -> Result<Vc<NextDynamicImportsResult>> {
+    async move {
+        // Traverse referenced modules graph, collect all of the dynamic imports:
+        // - Read the Program AST of the Module, this is the origin (A)
+        //  - If there's `dynamic(import(B))`, then B is the module that is being imported
+        // Returned import mappings are in the form of
+        // (Module<A>, Vec<(B, Module<B>)>) (where B is the raw import source string,
+        // and Module<B> is the actual resolved Module)
+        let (result, visited_modules) = NonDeterministic::new()
+            .skip_duplicates_with_visited_nodes(VisitedNodes(visited_modules.await?.0.clone()))
+            .visit(
+                server_entries
+                    .iter()
+                    .map(|module| async move {
+                        Ok(NextDynamicVisitEntry::Module(
+                            module.resolve().await?,
+                            module.ident().to_string().await?,
+                        ))
+                    })
+                    .try_join()
+                    .await?
+                    .into_iter(),
+                NextDynamicVisit {
+                    client_asset_context: client_asset_context.resolve().await?,
+                },
+            )
+            .await
+            .completed()?
+            .into_inner_with_visited();
+
+        let imported_modules_mapping = result.into_iter().filter_map(|entry| {
             if let NextDynamicVisitEntry::DynamicImportsMap(dynamic_imports_map) = entry {
                 Some(dynamic_imports_map)
             } else {
                 None
             }
         });
 
-    // Consolifate import mappings into a single indexmap
-    let mut import_mappings: IndexMap<Vc<Box<dyn Module>>, DynamicImportedModules> =
-        IndexMap::new();
+        // Consolidate import mappings into a single indexmap
+        let mut import_mappings: IndexMap<Vc<Box<dyn Module>>, DynamicImportedModules> =
+            IndexMap::new();
 
-    for module_mapping in imported_modules_mapping {
-        let (origin_module, dynamic_imports) = &*module_mapping.await?;
-        import_mappings
-            .entry(*origin_module)
-            .or_insert_with(Vec::new)
-            .append(&mut dynamic_imports.clone())
-    }
+        for module_mapping in imported_modules_mapping {
+            let (origin_module, dynamic_imports) = &*module_mapping.await?;
+            import_mappings
+                .entry(*origin_module)
+                .or_insert_with(Vec::new)
+                .append(&mut dynamic_imports.clone())
+        }
 
-    Ok(import_mappings)
+        Ok(NextDynamicImportsResult {
+            client_dynamic_imports: import_mappings,
+            visited_modules: VisitedDynamicImportModules(visited_modules.0).cell(),
+        }
+        .cell())
+    }
+    .instrument(tracing::info_span!("collecting next/dynamic imports"))
+    .await
 }
 
 #[derive(Debug, PartialEq, Eq, Hash, Clone, TraceRawVcs, Serialize, Deserialize)]

crates/next-api/src/pages.rs

@@ -61,7 +61,7 @@ use turbopack_nodejs::NodeJsChunkingContext;
 use crate::{
     dynamic_imports::{
         collect_chunk_group, collect_evaluated_chunk_group, collect_next_dynamic_imports,
-        DynamicImportedChunks,
+        DynamicImportedChunks, VisitedDynamicImportModules,
     },
     font::create_font_manifest,
     loadable_manifest::create_react_loadable_manifest,
@@ -810,8 +810,16 @@ impl PageEndpoint {
                 runtime,
             } = *self.internal_ssr_chunk_module().await?;
 
-            let is_edge = matches!(runtime, NextRuntime::Edge);
+            let dynamic_import_modules = collect_next_dynamic_imports(
+                vec![Vc::upcast(ssr_module)],
+                this.pages_project.client_module_context(),
+                VisitedDynamicImportModules::empty(),
+            )
+            .await?
+            .client_dynamic_imports
+            .clone();
 
+            let is_edge = matches!(runtime, NextRuntime::Edge);
             if is_edge {
                 let mut evaluatable_assets = edge_runtime_entries.await?.clone_value();
                 let evaluatable = Vc::try_resolve_sidecast(ssr_module)
@@ -825,11 +833,6 @@ impl PageEndpoint {
                     Value::new(AvailabilityInfo::Root),
                 );
 
-                let dynamic_import_modules = collect_next_dynamic_imports(
-                    [Vc::upcast(ssr_module)],
-                    this.pages_project.client_module_context(),
-                )
-                .await?;
                 let client_chunking_context =
                     this.pages_project.project().client_chunking_context();
                 let dynamic_import_entries = collect_evaluated_chunk_group(
@@ -863,18 +866,12 @@ impl PageEndpoint {
                     )
                     .await?;
 
-                let availability_info = Value::new(AvailabilityInfo::Root);
-                let dynamic_import_modules = collect_next_dynamic_imports(
-                    [Vc::upcast(ssr_module)],
-                    this.pages_project.client_module_context(),
-                )
-                .await?;
                 let client_chunking_context =
                     this.pages_project.project().client_chunking_context();
                 let dynamic_import_entries = collect_chunk_group(
                     Vc::upcast(client_chunking_context),
                     dynamic_import_modules,
-                    availability_info,
+                    Value::new(AvailabilityInfo::Root),
                 )
                 .await?;
 

crates/next-core/src/next_client_reference/visit_client_reference.rs

@@ -1,7 +1,7 @@
 use std::future::Future;
 
 use anyhow::Result;
-use indexmap::IndexSet;
+use indexmap::{IndexMap, IndexSet};
 use serde::{Deserialize, Serialize};
 use tracing::Instrument;
 use turbo_tasks::{
@@ -50,6 +50,10 @@ pub enum ClientReferenceType {
 #[derive(Debug)]
 pub struct ClientReferenceGraphResult {
     pub client_references: Vec<ClientReference>,
+    /// Only the [`ClientReferenceType::EcmascriptClientReference`]s are listed in this map.
+    #[allow(clippy::type_complexity)]
+    pub client_references_by_server_component:
+        IndexMap<Option<Vc<NextServerComponentModule>>, Vec<Vc<Box<dyn Module>>>>,
     pub server_component_entries: Vec<Vc<NextServerComponentModule>>,
     pub server_utils: Vec<Vc<Box<dyn Module>>>,
 }
@@ -81,6 +85,11 @@ pub async fn client_reference_graph(
         let mut server_component_entries = vec![];
         let mut server_utils = vec![];
 
+        let mut client_references_by_server_component = IndexMap::new();
+        // Make sure None (for the various internal next/dist/esm/client/components/*) is listed
+        // first
+        client_references_by_server_component.insert(None, Vec::new());
+
         let graph = AdjacencyMap::new()
             .skip_duplicates()
             .visit(
@@ -115,6 +124,15 @@ pub async fn client_reference_graph(
                 }
                 VisitClientReferenceNodeType::ClientReference(client_reference, _) => {
                     client_references.push(*client_reference);
+
+                    if let ClientReferenceType::EcmascriptClientReference(entry) =
+                        client_reference.ty()
+                    {
+                        client_references_by_server_component
+                            .entry(client_reference.server_component)
+                            .or_insert_with(Vec::new)
+                            .push(Vc::upcast::<Box<dyn Module>>(entry.await?.ssr_module));
+                    }
                 }
                 VisitClientReferenceNodeType::ServerUtilEntry(server_util, _) => {
                     server_utils.push(*server_util);
@@ -127,6 +145,7 @@ pub async fn client_reference_graph(
 
         Ok(ClientReferenceGraphResult {
             client_references,
+            client_references_by_server_component,
             server_component_entries,
             server_utils,
         }