diff --git a/crates/next-api/src/module_graph.rs b/crates/next-api/src/module_graph.rs index 15bf6c2e6a549..01905846d8f24 100644 --- a/crates/next-api/src/module_graph.rs +++ b/crates/next-api/src/module_graph.rs @@ -64,7 +64,9 @@ struct ModuleSet(pub HashSet>>); impl SingleModuleGraph { /// Walks the graph starting from the given entries and collects all reachable nodes, skipping /// nodes listed in `visited_modules` + /// If passed, `root` is connected to the entries and include in `self.entries`. async fn new_inner( + root: Option>>, entries: &Vec>>, visited_modules: &HashSet>>, ) -> Result> { @@ -73,7 +75,8 @@ impl SingleModuleGraph { let mut modules: HashMap>, NodeIndex> = HashMap::new(); let mut stack: Vec<_> = entries.iter().map(|e| (None, *e)).collect(); while let Some((parent_idx, module)) = stack.pop() { - if visited_modules.contains(&module) { + // Always add entries, even if already visited in other graphs + if parent_idx.is_some() && visited_modules.contains(&module) { continue; } if let Some(idx) = modules.get(&module) { @@ -95,11 +98,25 @@ impl SingleModuleGraph { stack.push((Some(idx), *reference)); } } + + let root_idx = root.and_then(|root| { + if !modules.contains_key(&root) { + let root_idx = graph.add_node(root); + for entry in entries { + graph.add_edge(root_idx, *modules.get(entry).unwrap(), ()); + } + Some((root, root_idx)) + } else { + None + } + }); + Ok(SingleModuleGraph { graph, entries: entries .iter() .map(|e| (*e, *modules.get(e).unwrap())) + .chain(root_idx.into_iter()) .collect(), } .cell()) @@ -175,16 +192,18 @@ impl SingleModuleGraph { impl SingleModuleGraph { #[turbo_tasks::function] async fn new_with_entries(entries: Vc) -> Result> { - SingleModuleGraph::new_inner(&*entries.await?, &Default::default()).await + SingleModuleGraph::new_inner(None, &*entries.await?, &Default::default()).await } + /// `root` is connected to the entries and include in `self.entries`. #[turbo_tasks::function] async fn new_with_entries_visited( + root: ResolvedVc>, // This must not be a Vc> to ensure layout segment optimization hits the cache entries: Vec>>, visited_modules: Vc, ) -> Result> { - SingleModuleGraph::new_inner(&entries, &*visited_modules.await?).await + SingleModuleGraph::new_inner(Some(root), &entries, &*visited_modules.await?).await } } @@ -199,6 +218,7 @@ async fn get_module_graph_for_endpoint( } = &*find_server_entries(*entry).await?; let graph = SingleModuleGraph::new_with_entries_visited( + *entry, server_utils.iter().map(|m| **m).collect(), Vc::cell(Default::default()), ) @@ -210,9 +230,9 @@ async fn get_module_graph_for_endpoint( for module in server_component_entries .iter() .map(|m| ResolvedVc::upcast::>(*m)) - .chain(std::iter::once(entry)) { let graph = SingleModuleGraph::new_with_entries_visited( + *entry, vec![*module], Vc::cell(visited_modules.clone()), ) @@ -221,6 +241,14 @@ async fn get_module_graph_for_endpoint( visited_modules.extend(graph.await?.graph.node_weights().copied()); graphs.push(graph); } + let graph = SingleModuleGraph::new_with_entries_visited( + *entry, + vec![*entry], + Vc::cell(visited_modules.clone()), + ) + .to_resolved() + .await?; + graphs.push(graph); Ok(Vc::cell(graphs)) } @@ -593,20 +621,22 @@ impl ReducedGraphs { // Just a single graph, no need to merge results Ok(graph.get_client_references_for_endpoint(entry)) } else { - todo!("get_client_references_for_endpoint multiple"); - // let result = self - // .client_references - // .iter() - // .map(|graph| async move { - // Ok(graph - // .get_client_references_for_endpoint(entry) - // .await? - // .clone_value()) - // }) - // .try_flat_join() - // .await?; - - // Ok(Vc::cell(result.into_iter().collect())) + let results = self + .client_references + .iter() + .map(|graph| async move { + let get_client_references_for_endpoint = + graph.get_client_references_for_endpoint(entry).await?; + Ok(get_client_references_for_endpoint) + }) + .try_join() + .await?; + + let mut result = results[0].clone_value(); + for r in results.into_iter().skip(1) { + result.extend(&r); + } + Ok(result.cell()) } } .instrument(span) diff --git a/crates/next-core/src/next_client_reference/visit_client_reference.rs b/crates/next-core/src/next_client_reference/visit_client_reference.rs index 36ace7943739e..bd672d695c8c8 100644 --- a/crates/next-core/src/next_client_reference/visit_client_reference.rs +++ b/crates/next-core/src/next_client_reference/visit_client_reference.rs @@ -61,6 +61,25 @@ pub struct ClientReferenceGraphResult { pub visited_nodes: ResolvedVc, } +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; + } +} + impl Default for ClientReferenceGraphResult { fn default() -> Self { ClientReferenceGraphResult {