Skip to content

Commit 645ef16

Browse files
sokrawbinnssmithmischnic
authored
[Turbopack] implement chunking based on the module graph (#74979)
### What? Visit the module graph by using the computed module graph during chunking --------- Co-authored-by: Will Binns-Smith <wbinnssmith@gmail.com> Co-authored-by: Niklas Mischkulnig <4586894+mischnic@users.noreply.github.com>
1 parent a0de92e commit 645ef16

File tree

13 files changed

+330
-628
lines changed

13 files changed

+330
-628
lines changed

‎crates/next-api/src/app.rs‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -728,28 +728,35 @@ impl AppProject {
728728
&self,
729729
endpoint: Vc<AppEndpoint>,
730730
rsc_entry: ResolvedVc<Box<dyn Module>>,
731+
extra_entries: Vc<EvaluatableAssets>,
731732
has_layout_segments: bool,
732733
) -> Result<Vc<ModuleGraphs>> {
734+
let extra_entries = extra_entries
735+
.await?
736+
.into_iter()
737+
.map(|m| *ResolvedVc::upcast(*m));
738+
733739
if *self.project.per_page_module_graph().await? {
734740
// Implements layout segment optimization to compute a graph "chain" for each layout
735741
// segment
736742
async move {
737743
let mut graphs = vec![];
738-
let mut visited_modules = VisitedModules::empty();
739-
740-
if has_layout_segments {
744+
let mut visited_modules = if has_layout_segments {
741745
let ServerEntries {
742746
server_utils,
743747
server_component_entries,
744748
} = &*find_server_entries(*rsc_entry).await?;
745-
if !server_utils.is_empty() {
746-
let graph = SingleModuleGraph::new_with_entries_visited(
747-
server_utils.iter().map(|m| **m).collect(),
748-
visited_modules,
749-
);
750-
graphs.push(graph);
751-
visited_modules = VisitedModules::from_graph(graph)
752-
}
749+
750+
let graph = SingleModuleGraph::new_with_entries_visited(
751+
server_utils
752+
.iter()
753+
.map(|m| **m)
754+
.chain(extra_entries)
755+
.collect(),
756+
VisitedModules::empty(),
757+
);
758+
graphs.push(graph);
759+
let mut visited_modules = VisitedModules::from_graph(graph);
753760

754761
for module in server_component_entries.iter() {
755762
let graph = SingleModuleGraph::new_with_entries_visited(
@@ -759,14 +766,27 @@ impl AppProject {
759766
graphs.push(graph);
760767
let is_layout =
761768
module.server_path().file_stem().await?.as_deref() == Some("layout");
762-
if is_layout {
769+
visited_modules = if is_layout {
763770
// Only propagate the visited_modules of the parent layout(s), not
764771
// across siblings such as loading.js and
765772
// page.js.
766-
visited_modules = visited_modules.concatenate(graph);
767-
}
773+
visited_modules.concatenate(graph)
774+
} else {
775+
// Prevents graph index from getting out of sync.
776+
// TODO We should remove VisitedModule entirely in favor of lookups
777+
// in SingleModuleGraph
778+
visited_modules.with_incremented_index()
779+
};
768780
}
769-
}
781+
visited_modules
782+
} else {
783+
let graph = SingleModuleGraph::new_with_entries_visited(
784+
extra_entries.collect::<_>(),
785+
VisitedModules::empty(),
786+
);
787+
graphs.push(graph);
788+
VisitedModules::from_graph(graph)
789+
};
770790

771791
let graph =
772792
SingleModuleGraph::new_with_entries_visited(vec![*rsc_entry], visited_modules);
@@ -1021,6 +1041,7 @@ impl AppEndpoint {
10211041
.app_module_graphs(
10221042
self,
10231043
*rsc_entry,
1044+
this.app_project.client_runtime_entries(),
10241045
matches!(this.ty, AppEndpointType::Page { .. }),
10251046
)
10261047
.await?;

‎crates/next-api/src/middleware.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,6 @@ impl MiddlewareEndpoint {
102102
)
103103
.resolve_entries(*this.asset_context);
104104

105-
let module_graph = this.project.module_graph_for_entries(evaluatable_assets);
106-
107105
let mut evaluatable_assets = evaluatable_assets.await?.clone_value();
108106

109107
let Some(module) =
@@ -116,11 +114,14 @@ impl MiddlewareEndpoint {
116114
.context("Entry module must be evaluatable")?;
117115
evaluatable_assets.push(evaluatable.to_resolved().await?);
118116

117+
let evaluatable_assets = Vc::cell(evaluatable_assets);
118+
let module_graph = this.project.module_graph_for_entries(evaluatable_assets);
119+
119120
let edge_chunking_context = this.project.edge_chunking_context(false);
120121

121122
let edge_files = edge_chunking_context.evaluated_chunk_group_assets(
122123
module.ident(),
123-
Vc::cell(evaluatable_assets),
124+
evaluatable_assets,
124125
module_graph,
125126
Value::new(AvailabilityInfo::Root),
126127
);

‎crates/next-api/src/module_graph.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -328,13 +328,13 @@ impl ClientReferencesGraph {
328328

329329
state_map.insert(module, current_server_component);
330330

331-
match module_type {
331+
Ok(match module_type {
332332
Some(
333333
ClientReferenceMapType::EcmascriptClientReference { .. }
334334
| ClientReferenceMapType::CssClientReference { .. },
335335
) => GraphTraversalAction::Skip,
336336
_ => GraphTraversalAction::Continue,
337-
}
337+
})
338338
},
339339
|parent_info, node, state_map| {
340340
let Some((parent_node, _)) = parent_info else {

‎crates/next-api/src/pages.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -941,7 +941,7 @@ impl PageEndpoint {
941941
)) = next_dynamic_imports
942942
{
943943
collect_next_dynamic_chunks(
944-
module_graph,
944+
self.client_module_graph(),
945945
Vc::upcast(project.client_chunking_context()),
946946
next_dynamic_imports,
947947
NextDynamicChunkAvailability::AvailabilityInfo(client_availability_info),

‎crates/next-core/src/next_font/google/mod.rs‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ use turbopack_core::{
1919
context::AssetContext,
2020
ident::AssetIdent,
2121
issue::{IssueExt, IssueSeverity},
22-
module_graph::ModuleGraph,
2322
reference_type::{InnerAssets, ReferenceType},
2423
resolve::{
2524
options::{ImportMapResult, ImportMappingReplacement, ReplacedImportMapping},
@@ -695,14 +694,12 @@ async fn get_mock_stylesheet(
695694
.module();
696695

697696
let root = mock_fs.root();
698-
let module_graph = ModuleGraph::from_module(mocked_response_asset);
699697
let val = evaluate(
700698
mocked_response_asset,
701699
root,
702700
*env,
703701
AssetIdent::from_path(loader_path),
704702
asset_context,
705-
module_graph,
706703
*chunking_context,
707704
None,
708705
vec![],

‎turbopack/crates/turbopack-core/src/chunk/chunk_group.rs‎

Lines changed: 168 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,26 @@
11
use std::collections::HashSet;
22

3-
use anyhow::Result;
3+
use anyhow::{Context, Result};
44
use auto_hash_map::AutoSet;
55
use futures::future::Either;
6-
use turbo_tasks::{FxIndexMap, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, Value, Vc};
6+
use swc_core::alloc::collections::FxHashMap;
7+
use turbo_tasks::{
8+
FxIndexMap, FxIndexSet, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, Value, Vc,
9+
};
710

811
use super::{
9-
availability_info::AvailabilityInfo, available_modules::AvailableModulesInfo, chunk_content,
10-
chunking::make_chunks, AsyncModuleInfo, Chunk, ChunkContentResult, ChunkItem, ChunkItemTy,
12+
availability_info::AvailabilityInfo, available_modules::AvailableModulesInfo,
13+
chunking::make_chunks, AsyncModuleInfo, Chunk, ChunkGroupContent, ChunkItem, ChunkItemTy,
1114
ChunkItemWithAsyncModuleInfo, ChunkableModule, ChunkingContext,
1215
};
1316
use crate::{
14-
environment::ChunkLoading, module::Module, module_graph::ModuleGraph, output::OutputAssets,
15-
rebase::RebasedAsset, reference::ModuleReference,
17+
chunk::ChunkingType,
18+
environment::ChunkLoading,
19+
module::Module,
20+
module_graph::{GraphTraversalAction, ModuleGraph},
21+
output::OutputAssets,
22+
rebase::RebasedAsset,
23+
reference::ModuleReference,
1624
};
1725

1826
pub struct MakeChunkGroupResult {
@@ -33,15 +41,16 @@ pub async fn make_chunk_group(
3341
);
3442
let should_trace = *chunking_context.is_tracing_enabled().await?;
3543

36-
let ChunkContentResult {
44+
let ChunkGroupContent {
3745
chunkable_modules,
3846
async_modules,
3947
traced_modules,
4048
passthrough_modules,
4149
forward_edges_inherit_async,
4250
local_back_edges_inherit_async,
4351
available_async_modules_back_edges_inherit_async,
44-
} = chunk_content(
52+
} = chunk_group_content(
53+
&*module_graph.await?,
4554
chunk_group_entries,
4655
availability_info,
4756
can_split_async,
@@ -272,3 +281,154 @@ pub async fn references_to_output_assets(
272281
.collect::<Vec<_>>();
273282
Ok(OutputAssets::new(output_assets))
274283
}
284+
285+
pub async fn chunk_group_content(
286+
module_graph: &ModuleGraph,
287+
chunk_group_entries: impl IntoIterator<Item = ResolvedVc<Box<dyn Module>>>,
288+
availability_info: AvailabilityInfo,
289+
can_split_async: bool,
290+
should_trace: bool,
291+
) -> Result<ChunkGroupContent> {
292+
type ModuleToChunkableMap =
293+
FxHashMap<ResolvedVc<Box<dyn Module>>, ResolvedVc<Box<dyn ChunkableModule>>>;
294+
295+
struct TraverseState {
296+
unsorted_chunkable_modules: ModuleToChunkableMap,
297+
result: ChunkGroupContent,
298+
}
299+
300+
let mut state = TraverseState {
301+
unsorted_chunkable_modules: FxHashMap::default(),
302+
result: ChunkGroupContent {
303+
chunkable_modules: FxIndexSet::default(),
304+
async_modules: FxIndexSet::default(),
305+
traced_modules: FxIndexSet::default(),
306+
passthrough_modules: FxIndexSet::default(),
307+
forward_edges_inherit_async: FxIndexMap::default(),
308+
local_back_edges_inherit_async: FxIndexMap::default(),
309+
available_async_modules_back_edges_inherit_async: FxIndexMap::default(),
310+
},
311+
};
312+
313+
let available_modules = match availability_info.available_modules() {
314+
Some(available_modules) => Some(available_modules.snapshot().await?),
315+
None => None,
316+
};
317+
318+
module_graph
319+
.traverse_edges_from_entries_topological(
320+
chunk_group_entries,
321+
&mut state,
322+
|parent_info,
323+
node,
324+
TraverseState {
325+
unsorted_chunkable_modules,
326+
result,
327+
}| {
328+
if let Some((_, ChunkingType::Traced)) = parent_info {
329+
if should_trace {
330+
result.traced_modules.insert(node.module);
331+
}
332+
return Ok(GraphTraversalAction::Skip);
333+
}
334+
335+
let Some(chunkable_module) =
336+
ResolvedVc::try_sidecast_sync::<Box<dyn ChunkableModule>>(node.module)
337+
else {
338+
return Ok(GraphTraversalAction::Skip);
339+
};
340+
341+
let available_info = available_modules
342+
.as_ref()
343+
.and_then(|available_modules| available_modules.get(chunkable_module));
344+
345+
let Some((parent_node, edge)) = parent_info else {
346+
return Ok(if available_info.is_some() {
347+
GraphTraversalAction::Skip
348+
} else {
349+
unsorted_chunkable_modules.insert(node.module, chunkable_module);
350+
GraphTraversalAction::Continue
351+
});
352+
};
353+
354+
let parent_module =
355+
ResolvedVc::try_sidecast_sync::<Box<dyn ChunkableModule>>(parent_node.module)
356+
.context("Expected parent module to be chunkable")?;
357+
358+
Ok(match edge {
359+
ChunkingType::Passthrough => {
360+
result.passthrough_modules.insert(chunkable_module);
361+
GraphTraversalAction::Continue
362+
}
363+
ChunkingType::Parallel => {
364+
if available_info.is_some() {
365+
GraphTraversalAction::Skip
366+
} else {
367+
unsorted_chunkable_modules.insert(node.module, chunkable_module);
368+
GraphTraversalAction::Continue
369+
}
370+
}
371+
ChunkingType::ParallelInheritAsync => {
372+
result
373+
.forward_edges_inherit_async
374+
.entry(parent_module)
375+
.or_default()
376+
.push(chunkable_module);
377+
if let Some(info) = available_info {
378+
if info.is_async {
379+
result
380+
.available_async_modules_back_edges_inherit_async
381+
.entry(chunkable_module)
382+
.or_default()
383+
.push(parent_module);
384+
}
385+
GraphTraversalAction::Skip
386+
} else {
387+
result
388+
.local_back_edges_inherit_async
389+
.entry(chunkable_module)
390+
.or_default()
391+
.push(parent_module);
392+
unsorted_chunkable_modules.insert(node.module, chunkable_module);
393+
GraphTraversalAction::Continue
394+
}
395+
}
396+
ChunkingType::Async => {
397+
if can_split_async {
398+
result.async_modules.insert(chunkable_module);
399+
GraphTraversalAction::Skip
400+
} else if available_info.is_some() {
401+
GraphTraversalAction::Skip
402+
} else {
403+
unsorted_chunkable_modules.insert(node.module, chunkable_module);
404+
GraphTraversalAction::Continue
405+
}
406+
}
407+
ChunkingType::Traced => {
408+
// handled above before the sidecast
409+
unreachable!();
410+
}
411+
ChunkingType::Isolated { .. } => {
412+
// TODO currently not implemented
413+
GraphTraversalAction::Skip
414+
}
415+
})
416+
},
417+
|_,
418+
node,
419+
TraverseState {
420+
unsorted_chunkable_modules,
421+
result,
422+
}| {
423+
// Insert modules in topological order
424+
if let Some(chunkable_module) =
425+
unsorted_chunkable_modules.get(&node.module).copied()
426+
{
427+
result.chunkable_modules.insert(chunkable_module);
428+
}
429+
},
430+
)
431+
.await?;
432+
433+
Ok(state.result)
434+
}

0 commit comments

Comments
 (0)