turbopack_core/chunk/
chunk_group.rs

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
34/// Creates a chunk group from a set of entries.
35pub 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    // Attach async info to chunkable modules
72    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    // Insert async chunk loaders for every referenced async module
103    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    // And also add output assets referenced by async chunk loaders
122    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    // Pass chunk items to chunking algorithm
148    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            // Traced modules need to have a special handling
228            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                // An entry from the entries list
255                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                    // handled above before the sidecast
302                    unreachable!();
303                }
304                ChunkingType::Isolated { .. } => {
305                    // TODO currently not implemented
306                    GraphTraversalAction::Exclude
307                }
308            })
309        },
310        |_, node, state| {
311            // Insert modules in topological order
312            if let Some(chunkable_module) = state.unsorted_items.get(node).copied() {
313                state.chunkable_items.insert(chunkable_module);
314            }
315        },
316    )?;
317
318    // This needs to use the unmerged items
319    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}