Turbopack: find client references layout segment optimization (#70792)

Closes PACK-3273

Before
![Bildschirmfoto 2024-10-04 um 16 13 15
Kopie](https://github.com/user-attachments/assets/6157e94c-bf49-4a59-9341-99244916d3be)

After
![Bildschirmfoto 2024-10-04 um 16 15
02](https://github.com/user-attachments/assets/6ee2d4e3-66f5-42a3-ba37-a3d7e0467cb1)

crates/next-api/src/app.rs

@@ -18,7 +18,10 @@ 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, NextEcmascriptClientReferenceTransition},
+    next_client_reference::{
+        client_reference_graph, find_server_entries, NextEcmascriptClientReferenceTransition,
+        ServerEntries, VisitedClientReferenceGraphNodes,
+    },
     next_config::NextConfig,
     next_dynamic::NextDynamicTransition,
     next_edge::route_regex::get_named_middleware_regex,
@@ -837,8 +840,6 @@ impl AppEndpoint {
 
         let rsc_entry = app_entry.rsc_entry;
 
-        let rsc_entry_asset = Vc::upcast(rsc_entry);
-
         let client_chunking_context = this.app_project.project().client_chunking_context();
 
         let (app_server_reference_modules, client_dynamic_imports, client_references) =
@@ -865,7 +866,33 @@ 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_references = {
+                    let ServerEntries {
+                        server_component_entries,
+                        server_utils,
+                    } = &*find_server_entries(rsc_entry).await?;
+
+                    let mut client_references = client_reference_graph(
+                        server_utils.clone(),
+                        VisitedClientReferenceGraphNodes::empty(),
+                    )
+                    .await?
+                    .clone_value();
+
+                    for module in server_component_entries
+                        .iter()
+                        .map(|m| Vc::upcast::<Box<dyn Module>>(*m))
+                        .chain(std::iter::once(rsc_entry))
+                    {
+                        let current_client_references =
+                            client_reference_graph(vec![module], client_references.visited_nodes)
+                                .await?;
+
+                        client_references.extend(&current_client_references);
+                    }
+                    client_references
+                };
+                let client_references_cell = client_references.clone().cell();
 
                 let ssr_chunking_context = if process_ssr {
                     Some(match runtime {
@@ -886,7 +913,6 @@ impl AppEndpoint {
                     let mut visited_modules = VisitedDynamicImportModules::empty();
 
                     for refs in client_references
-                        .await?
                         .client_references_by_server_component
                         .values()
                     {
@@ -909,7 +935,7 @@ impl AppEndpoint {
                 };
 
                 let client_references_chunks = get_app_client_references_chunks(
-                    client_references,
+                    client_references_cell,
                     client_chunking_context,
                     Value::new(client_shared_availability_info),
                     ssr_chunking_context,
@@ -1009,7 +1035,7 @@ impl AppEndpoint {
                     node_root,
                     client_relative_path,
                     app_entry.original_name.clone(),
-                    client_references,
+                    client_references_cell,
                     client_references_chunks,
                     client_chunking_context,
                     ssr_chunking_context,
@@ -1037,7 +1063,9 @@ impl AppEndpoint {
                 }
 
                 (
-                    Some(get_app_server_reference_modules(client_references.types())),
+                    Some(get_app_server_reference_modules(
+                        client_references_cell.types(),
+                    )),
                     Some(client_dynamic_imports),
                     Some(client_references),
                 )
@@ -1274,8 +1302,7 @@ impl AppEndpoint {
                     let mut current_chunks = OutputAssets::empty();
                     let mut current_availability_info = AvailabilityInfo::Root;
                     if let Some(client_references) = client_references {
-                        let client_references = client_references.await?;
-                        let span = tracing::trace_span!("server utils",);
+                        let span = tracing::trace_span!("server utils");
                         async {
                             let utils_module = IncludeModulesModule::new(
                                 AssetIdent::from_path(this.app_project.project().project_path())

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

@@ -6,6 +6,6 @@ pub use ecmascript_client_reference::{
     ecmascript_client_reference_transition::NextEcmascriptClientReferenceTransition,
 };
 pub use visit_client_reference::{
-    client_reference_graph, ClientReference, ClientReferenceGraphResult, ClientReferenceType,
-    ClientReferenceTypes,
+    client_reference_graph, find_server_entries, ClientReference, ClientReferenceGraphResult,
+    ClientReferenceType, ClientReferenceTypes, ServerEntries, VisitedClientReferenceGraphNodes,
 };

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

@@ -1,21 +1,18 @@
-use std::future::Future;
+use std::{collections::HashSet, future::Future};
 
 use anyhow::Result;
 use indexmap::{IndexMap, IndexSet};
 use serde::{Deserialize, Serialize};
 use tracing::Instrument;
 use turbo_tasks::{
     debug::ValueDebugFormat,
-    graph::{AdjacencyMap, GraphTraversal, Visit, VisitControlFlow},
+    graph::{AdjacencyMap, GraphTraversal, Visit, VisitControlFlow, VisitedNodes},
     trace::TraceRawVcs,
     RcStr, ReadRef, TryJoinIterExt, ValueToString, Vc,
 };
 use turbo_tasks_fs::FileSystemPath;
 use turbopack::css::CssModuleAsset;
-use turbopack_core::{
-    module::{Module, Modules},
-    reference::primary_referenced_modules,
-};
+use turbopack_core::{module::Module, reference::primary_referenced_modules};
 
 use super::ecmascript_client_reference::ecmascript_client_reference_module::EcmascriptClientReferenceModule;
 use crate::next_server_component::server_component_module::NextServerComponentModule;
@@ -46,8 +43,8 @@ pub enum ClientReferenceType {
     CssClientReference(Vc<CssModuleAsset>),
 }
 
-#[turbo_tasks::value]
-#[derive(Debug)]
+#[turbo_tasks::value(shared)]
+#[derive(Clone, Debug)]
 pub struct ClientReferenceGraphResult {
     pub client_references: Vec<ClientReference>,
     /// Only the [`ClientReferenceType::EcmascriptClientReference`]s are listed in this map.
@@ -56,6 +53,30 @@ pub struct ClientReferenceGraphResult {
         IndexMap<Option<Vc<NextServerComponentModule>>, Vec<Vc<Box<dyn Module>>>>,
     pub server_component_entries: Vec<Vc<NextServerComponentModule>>,
     pub server_utils: Vec<Vc<Box<dyn Module>>>,
+    pub visited_nodes: Vc<VisitedClientReferenceGraphNodes>,
+}
+
+impl Default for ClientReferenceGraphResult {
+    fn default() -> Self {
+        ClientReferenceGraphResult {
+            client_references: Default::default(),
+            client_references_by_server_component: Default::default(),
+            server_component_entries: Default::default(),
+            server_utils: Default::default(),
+            visited_nodes: VisitedClientReferenceGraphNodes::empty(),
+        }
+    }
+}
+
+#[turbo_tasks::value(shared)]
+pub struct VisitedClientReferenceGraphNodes(HashSet<VisitClientReferenceNode>);
+
+#[turbo_tasks::value_impl]
+impl VisitedClientReferenceGraphNodes {
+    #[turbo_tasks::function]
+    pub fn empty() -> Vc<Self> {
+        VisitedClientReferenceGraphNodes(Default::default()).cell()
+    }
 }
 
 #[turbo_tasks::value(transparent)]
@@ -74,13 +95,31 @@ impl ClientReferenceGraphResult {
     }
 }
 
+impl ClientReferenceGraphResult {
+    /// Merges multiple return values of client_reference_graph together.
+    pub fn extend(&mut self, other: &Self) {
+        self.client_references
+            .extend(other.client_references.iter().copied());
+        for (k, v) in other.client_references_by_server_component.iter() {
+            self.client_references_by_server_component
+                .entry(*k)
+                .or_insert_with(Vec::new)
+                .extend(v);
+        }
+        self.server_component_entries
+            .extend(other.server_component_entries.iter().copied());
+        self.server_utils.extend(other.server_utils.iter().copied());
+        // This is merged already by `client_reference_graph` itself
+        self.visited_nodes = other.visited_nodes;
+    }
+}
+
 #[turbo_tasks::function]
 pub async fn client_reference_graph(
-    entries: Vc<Modules>,
+    entries: Vec<Vc<Box<dyn Module>>>,
+    visited_nodes: Vc<VisitedClientReferenceGraphNodes>,
 ) -> Result<Vc<ClientReferenceGraphResult>> {
     async move {
-        let entries = entries.await?;
-
         let mut client_references = vec![];
         let mut server_component_entries = vec![];
         let mut server_utils = vec![];
@@ -90,16 +129,25 @@ pub async fn client_reference_graph(
         // first
         client_references_by_server_component.insert(None, Vec::new());
 
-        let graph = AdjacencyMap::new()
-            .skip_duplicates()
+        let (graph, visited_nodes) = AdjacencyMap::new()
+            .skip_duplicates_with_visited_nodes(VisitedNodes(visited_nodes.await?.0.clone()))
             .visit(
                 entries
                     .iter()
                     .copied()
                     .map(|module| async move {
                         Ok(VisitClientReferenceNode {
-                            state: VisitClientReferenceNodeState::Entry {
-                                entry_path: module.ident().path().resolve().await?,
+                            state: if let Some(server_component) =
+                                Vc::try_resolve_downcast_type::<NextServerComponentModule>(module)
+                                    .await?
+                            {
+                                VisitClientReferenceNodeState::InServerComponent {
+                                    server_component,
+                                }
+                            } else {
+                                VisitClientReferenceNodeState::Entry {
+                                    entry_path: module.ident().path().resolve().await?,
+                                }
                             },
                             ty: VisitClientReferenceNodeType::Internal(
                                 module,
@@ -109,14 +157,15 @@ pub async fn client_reference_graph(
                     })
                     .try_join()
                     .await?,
-                VisitClientReference,
+                VisitClientReference {
+                    stop_at_server_entries: false,
+                },
             )
             .await
             .completed()?
-            .into_inner()
-            .into_reverse_topological();
+            .into_inner_with_visited();
 
-        for node in graph {
+        for node in graph.into_reverse_topological() {
             match &node.ty {
                 VisitClientReferenceNodeType::Internal(_asset, _) => {
                     // No-op. These nodes are only useful during graph
@@ -148,14 +197,68 @@ pub async fn client_reference_graph(
             client_references_by_server_component,
             server_component_entries,
             server_utils,
+            visited_nodes: VisitedClientReferenceGraphNodes(visited_nodes.0).cell(),
         }
         .cell())
     }
     .instrument(tracing::info_span!("find client references"))
     .await
 }
 
-struct VisitClientReference;
+#[turbo_tasks::value(shared)]
+#[derive(Clone, Debug)]
+pub struct ServerEntries {
+    pub server_component_entries: Vec<Vc<NextServerComponentModule>>,
+    pub server_utils: Vec<Vc<Box<dyn Module>>>,
+}
+
+#[turbo_tasks::function]
+pub async fn find_server_entries(entry: Vc<Box<dyn Module>>) -> Result<Vc<ServerEntries>> {
+    let graph = AdjacencyMap::new()
+        .skip_duplicates()
+        .visit(
+            vec![VisitClientReferenceNode {
+                state: {
+                    VisitClientReferenceNodeState::Entry {
+                        entry_path: entry.ident().path().resolve().await?,
+                    }
+                },
+                ty: VisitClientReferenceNodeType::Internal(entry, entry.ident().to_string().await?),
+            }],
+            VisitClientReference {
+                stop_at_server_entries: true,
+            },
+        )
+        .await
+        .completed()?
+        .into_inner();
+
+    let mut server_component_entries = vec![];
+    let mut server_utils = vec![];
+    for node in graph.reverse_topological() {
+        match &node.ty {
+            VisitClientReferenceNodeType::ServerUtilEntry(server_util, _) => {
+                server_utils.push(*server_util);
+            }
+            VisitClientReferenceNodeType::ServerComponentEntry(server_component, _) => {
+                server_component_entries.push(*server_component);
+            }
+            VisitClientReferenceNodeType::Internal(_, _)
+            | VisitClientReferenceNodeType::ClientReference(_, _) => {}
+        }
+    }
+
+    Ok(ServerEntries {
+        server_component_entries,
+        server_utils,
+    }
+    .cell())
+}
+
+struct VisitClientReference {
+    /// Used to discover ServerComponents and ServerUtils
+    stop_at_server_entries: bool,
+}
 
 #[derive(
     Clone, Eq, PartialEq, Hash, Serialize, Deserialize, Debug, ValueDebugFormat, TraceRawVcs,
@@ -205,11 +308,21 @@ impl Visit<VisitClientReferenceNode> for VisitClientReference {
     type EdgesFuture = impl Future<Output = Result<Self::EdgesIntoIter>>;
 
     fn visit(&mut self, edge: Self::Edge) -> VisitControlFlow<VisitClientReferenceNode> {
+        if self.stop_at_server_entries
+            && matches!(
+                edge.ty,
+                VisitClientReferenceNodeType::ServerUtilEntry(..)
+                    | VisitClientReferenceNodeType::ServerComponentEntry(..)
+            )
+        {
+            return VisitControlFlow::Skip(edge);
+        }
+
         match edge.ty {
             VisitClientReferenceNodeType::ClientReference(..) => VisitControlFlow::Skip(edge),
-            VisitClientReferenceNodeType::Internal(..) => VisitControlFlow::Continue(edge),
-            VisitClientReferenceNodeType::ServerUtilEntry(..) => VisitControlFlow::Continue(edge),
-            VisitClientReferenceNodeType::ServerComponentEntry(..) => {
+            VisitClientReferenceNodeType::Internal(..)
+            | VisitClientReferenceNodeType::ServerUtilEntry(..)
+            | VisitClientReferenceNodeType::ServerComponentEntry(..) => {
                 VisitControlFlow::Continue(edge)
             }
         }