1use std::{collections::HashSet, sync::atomic::AtomicBool};
2
3use anyhow::{Context, Result};
4use rustc_hash::FxHashMap;
5use smallvec::{SmallVec, smallvec};
6use turbo_tasks::{FxIndexSet, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, Vc};
7
8use super::{
9 Chunk, ChunkGroupContent, ChunkItem, ChunkItemWithAsyncModuleInfo, ChunkingContext,
10 availability_info::AvailabilityInfo, chunking::make_chunks,
11};
12use crate::{
13 chunk::{
14 ChunkableModule, ChunkingType,
15 chunk_item_batch::{ChunkItemBatchGroup, ChunkItemOrBatchWithAsyncModuleInfo},
16 },
17 environment::ChunkLoading,
18 module::Module,
19 module_graph::{
20 GraphTraversalAction, ModuleGraph,
21 module_batch::{ChunkableModuleBatchGroup, ChunkableModuleOrBatch, ModuleOrBatch},
22 module_batches::{BatchingConfig, ModuleBatchesGraphEdge},
23 },
24 output::OutputAssets,
25 reference::ModuleReference,
26 traced_asset::TracedAsset,
27};
28
29pub struct MakeChunkGroupResult {
30 pub chunks: Vec<ResolvedVc<Box<dyn Chunk>>>,
31 pub availability_info: AvailabilityInfo,
32}
33
34pub async fn make_chunk_group(
36 chunk_group_entries: impl IntoIterator<
37 IntoIter = impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + Send,
38 > + Send
39 + Clone,
40 module_graph: Vc<ModuleGraph>,
41 chunking_context: ResolvedVc<Box<dyn ChunkingContext>>,
42 availability_info: AvailabilityInfo,
43) -> Result<MakeChunkGroupResult> {
44 let can_split_async = !matches!(
45 *chunking_context.environment().chunk_loading().await?,
46 ChunkLoading::Edge
47 );
48 let should_trace = *chunking_context.is_tracing_enabled().await?;
49 let should_merge_modules = *chunking_context.is_module_merging_enabled().await?;
50 let batching_config = chunking_context.batching_config();
51
52 let ChunkGroupContent {
53 chunkable_items,
54 batch_groups,
55 async_modules,
56 traced_modules,
57 availability_info,
58 } = chunk_group_content(
59 module_graph,
60 chunk_group_entries.clone(),
61 availability_info,
62 can_split_async,
63 should_trace,
64 should_merge_modules,
65 batching_config,
66 )
67 .await?;
68
69 let async_modules_info = module_graph.async_module_info().await?;
70
71 let mut chunk_items = chunkable_items
73 .iter()
74 .copied()
75 .map(|m| {
76 ChunkItemOrBatchWithAsyncModuleInfo::from_chunkable_module_or_batch(
77 m,
78 &async_modules_info,
79 module_graph,
80 *chunking_context,
81 )
82 })
83 .try_join()
84 .await?
85 .into_iter()
86 .flatten()
87 .collect::<Vec<_>>();
88
89 let chunk_item_batch_groups = batch_groups
90 .iter()
91 .map(|&batch_group| {
92 ChunkItemBatchGroup::from_module_batch_group(
93 ChunkableModuleBatchGroup::from_module_batch_group(*batch_group),
94 module_graph,
95 *chunking_context,
96 )
97 .to_resolved()
98 })
99 .try_join()
100 .await?;
101
102 let async_loaders = async_modules
104 .into_iter()
105 .map(async |module| {
106 chunking_context
107 .async_loader_chunk_item(*module, module_graph, availability_info)
108 .to_resolved()
109 .await
110 })
111 .try_join()
112 .await?;
113 let async_loader_chunk_items = async_loaders.iter().map(|&chunk_item| {
114 ChunkItemOrBatchWithAsyncModuleInfo::ChunkItem(ChunkItemWithAsyncModuleInfo {
115 chunk_item,
116 module: None,
117 async_info: None,
118 })
119 });
120
121 let async_loader_references = async_loaders
123 .iter()
124 .map(|&loader| loader.references())
125 .try_join()
126 .await?;
127
128 let mut referenced_output_assets = traced_modules
129 .into_iter()
130 .map(|module| async move {
131 Ok(ResolvedVc::upcast(
132 TracedAsset::new(*module).to_resolved().await?,
133 ))
134 })
135 .try_join()
136 .await?;
137
138 chunk_items.extend(async_loader_chunk_items);
139 referenced_output_assets.reserve(
140 async_loader_references
141 .iter()
142 .map(|r| r.len())
143 .sum::<usize>(),
144 );
145 referenced_output_assets.extend(async_loader_references.into_iter().flatten());
146
147 let chunks = make_chunks(
149 module_graph,
150 chunking_context,
151 chunk_items,
152 chunk_item_batch_groups,
153 "".into(),
154 ResolvedVc::cell(referenced_output_assets),
155 )
156 .await?;
157
158 Ok(MakeChunkGroupResult {
159 chunks,
160 availability_info,
161 })
162}
163
164pub async fn references_to_output_assets(
165 references: impl IntoIterator<Item = &ResolvedVc<Box<dyn ModuleReference>>>,
166) -> Result<Vc<OutputAssets>> {
167 let output_assets = references
168 .into_iter()
169 .map(|reference| reference.resolve_reference().primary_output_assets())
170 .try_join()
171 .await?;
172 let mut set = HashSet::new();
173 let output_assets = output_assets
174 .iter()
175 .flatten()
176 .copied()
177 .filter(|&asset| set.insert(asset))
178 .map(|asset| *asset)
179 .collect::<Vec<_>>();
180 Ok(OutputAssets::new(output_assets))
181}
182
183pub async fn chunk_group_content(
184 module_graph: Vc<ModuleGraph>,
185 chunk_group_entries: impl IntoIterator<
186 IntoIter = impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + Send,
187 > + Send,
188 availability_info: AvailabilityInfo,
189 can_split_async: bool,
190 should_trace: bool,
191 should_merge_modules: bool,
192 batching_config: Vc<BatchingConfig>,
193) -> Result<ChunkGroupContent> {
194 let module_batches_graph = module_graph.module_batches(batching_config).await?;
195
196 type ModuleToChunkableMap = FxHashMap<ModuleOrBatch, ChunkableModuleOrBatch>;
197
198 struct TraverseState {
199 unsorted_items: ModuleToChunkableMap,
200 chunkable_items: FxIndexSet<ChunkableModuleOrBatch>,
201 async_modules: FxIndexSet<ResolvedVc<Box<dyn ChunkableModule>>>,
202 traced_modules: FxIndexSet<ResolvedVc<Box<dyn Module>>>,
203 }
204
205 let mut state = TraverseState {
206 unsorted_items: FxHashMap::default(),
207 chunkable_items: FxIndexSet::default(),
208 async_modules: FxIndexSet::default(),
209 traced_modules: FxIndexSet::default(),
210 };
211
212 let available_modules = match availability_info.available_modules() {
213 Some(available_modules) => Some(available_modules.snapshot().await?),
214 None => None,
215 };
216
217 let chunk_group_entries = chunk_group_entries.into_iter();
218 let mut entries = Vec::with_capacity(chunk_group_entries.size_hint().0);
219 for entry in chunk_group_entries {
220 entries.push(module_batches_graph.get_entry_index(entry).await?);
221 }
222
223 module_batches_graph.traverse_edges_from_entries_topological(
224 entries,
225 &mut state,
226 |parent_info, &node, state| {
227 if let Some((
229 _,
230 ModuleBatchesGraphEdge {
231 ty: ChunkingType::Traced,
232 ..
233 },
234 )) = parent_info
235 {
236 if should_trace {
237 let ModuleOrBatch::Module(module) = node else {
238 unreachable!();
239 };
240 state.traced_modules.insert(module);
241 }
242 return Ok(GraphTraversalAction::Exclude);
243 }
244
245 let Some(chunkable_node) = ChunkableModuleOrBatch::from_module_or_batch(node) else {
246 return Ok(GraphTraversalAction::Exclude);
247 };
248
249 let is_available = available_modules
250 .as_ref()
251 .is_some_and(|available_modules| available_modules.get(chunkable_node));
252
253 let Some((_, edge)) = parent_info else {
254 return Ok(if is_available {
256 GraphTraversalAction::Exclude
257 } else if state
258 .unsorted_items
259 .try_insert(node, chunkable_node)
260 .is_ok()
261 {
262 GraphTraversalAction::Continue
263 } else {
264 GraphTraversalAction::Exclude
265 });
266 };
267
268 Ok(match edge.ty {
269 ChunkingType::Parallel { .. } | ChunkingType::Shared { .. } => {
270 if is_available {
271 GraphTraversalAction::Exclude
272 } else if state
273 .unsorted_items
274 .try_insert(node, chunkable_node)
275 .is_ok()
276 {
277 GraphTraversalAction::Continue
278 } else {
279 GraphTraversalAction::Exclude
280 }
281 }
282 ChunkingType::Async => {
283 if can_split_async {
284 let chunkable_module = ResolvedVc::try_downcast(edge.module.unwrap())
285 .context("Module in async chunking edge is not chunkable")?;
286 state.async_modules.insert(chunkable_module);
287 GraphTraversalAction::Exclude
288 } else if is_available {
289 GraphTraversalAction::Exclude
290 } else if state
291 .unsorted_items
292 .try_insert(node, chunkable_node)
293 .is_ok()
294 {
295 GraphTraversalAction::Continue
296 } else {
297 GraphTraversalAction::Exclude
298 }
299 }
300 ChunkingType::Traced => {
301 unreachable!();
303 }
304 ChunkingType::Isolated { .. } => {
305 GraphTraversalAction::Exclude
307 }
308 })
309 },
310 |_, node, state| {
311 if let Some(chunkable_module) = state.unsorted_items.get(node).copied() {
313 state.chunkable_items.insert(chunkable_module);
314 }
315 },
316 )?;
317
318 let availability_info = availability_info
320 .with_modules(Vc::cell(state.chunkable_items.clone()))
321 .await?;
322
323 if should_merge_modules {
324 let merged_modules = module_graph.merged_modules().await?;
325 state.chunkable_items = state
326 .chunkable_items
327 .into_iter()
328 .map(async |chunkable_module| match chunkable_module {
329 ChunkableModuleOrBatch::Module(module) => {
330 if !merged_modules.should_create_chunk_item_for(ResolvedVc::upcast(module)) {
331 return Ok(smallvec![]);
332 }
333
334 let module = if let Some(replacement) =
335 merged_modules.should_replace_module(ResolvedVc::upcast(module))
336 {
337 replacement
338 } else {
339 module
340 };
341
342 Ok(smallvec![ChunkableModuleOrBatch::Module(module)])
343 }
344 ChunkableModuleOrBatch::Batch(batch) => {
345 let batch_ref = batch.await?;
346 let modules = &batch_ref.modules;
347
348 let modified = AtomicBool::new(false);
349 let modules = modules
350 .iter()
351 .filter(|module| {
352 if merged_modules
353 .should_create_chunk_item_for(ResolvedVc::upcast(**module))
354 {
355 true
356 } else {
357 modified.store(true, std::sync::atomic::Ordering::Relaxed);
358 false
359 }
360 })
361 .map(|&module| {
362 if let Some(replacement) =
363 merged_modules.should_replace_module(ResolvedVc::upcast(module))
364 {
365 modified.store(true, std::sync::atomic::Ordering::Relaxed);
366 replacement
367 } else {
368 module
369 }
370 })
371 .map(ChunkableModuleOrBatch::Module)
372 .collect::<SmallVec<[_; 1]>>();
373
374 if modified.load(std::sync::atomic::Ordering::Relaxed) {
375 Ok(modules)
376 } else {
377 Ok(smallvec![ChunkableModuleOrBatch::Batch(batch)])
378 }
379 }
380 ChunkableModuleOrBatch::None(i) => Ok(smallvec![ChunkableModuleOrBatch::None(i)]),
381 })
382 .try_flat_join()
383 .await?
384 .into_iter()
385 .collect();
386 };
387
388 let mut batch_groups = FxIndexSet::default();
389 for &module in &state.chunkable_items {
390 if let Some(batch_group) = module_batches_graph.get_batch_group(&module.into()) {
391 batch_groups.insert(batch_group);
392 }
393 }
394
395 Ok(ChunkGroupContent {
396 chunkable_items: state.chunkable_items,
397 batch_groups,
398 async_modules: state.async_modules,
399 traced_modules: state.traced_modules,
400 availability_info,
401 })
402}