turbopack_core/module_graph/
chunk_group_info.rs

1use std::{
2    hash::Hash,
3    ops::{Deref, DerefMut},
4    sync::LazyLock,
5};
6
7use anyhow::{Context, Result, bail};
8use bincode::{Decode, Encode};
9use either::Either;
10use indexmap::map::Entry;
11use roaring::RoaringBitmap;
12use rustc_hash::FxHashMap;
13use tracing::Instrument;
14use turbo_rcstr::RcStr;
15use turbo_tasks::{
16    FxIndexMap, FxIndexSet, NonLocalValue, ResolvedVc, TaskInput, TryJoinIterExt, ValueToString,
17    Vc, debug::ValueDebugFormat, turbofmt,
18};
19
20use crate::{
21    chunk::ChunkingType,
22    module::Module,
23    module_graph::{GraphTraversalAction, ModuleGraph, RefData},
24};
25
26#[derive(Clone, Debug, Default, PartialEq, ValueDebugFormat, Encode, Decode)]
27#[repr(transparent)]
28pub struct RoaringBitmapWrapper(#[bincode(with_serde)] pub RoaringBitmap);
29
30impl TaskInput for RoaringBitmapWrapper {
31    fn is_transient(&self) -> bool {
32        false
33    }
34}
35
36impl RoaringBitmapWrapper {
37    /// Whether `self` contains bits that are not in `other`
38    ///
39    /// The existing `is_superset` method also returns true for equal sets
40    pub fn is_proper_superset(&self, other: &Self) -> bool {
41        !self.is_subset(other)
42    }
43
44    pub fn into_inner(self) -> RoaringBitmap {
45        self.0
46    }
47}
48unsafe impl NonLocalValue for RoaringBitmapWrapper {}
49
50// RoaringBitmap doesn't impl Eq: https://github.com/RoaringBitmap/roaring-rs/issues/302
51// PartialEq can only return true if both bitmaps have the same internal representation, but two
52// bitmaps with the same content should always have the same internal representation
53impl Eq for RoaringBitmapWrapper {}
54
55impl Deref for RoaringBitmapWrapper {
56    type Target = RoaringBitmap;
57    fn deref(&self) -> &Self::Target {
58        &self.0
59    }
60}
61impl DerefMut for RoaringBitmapWrapper {
62    fn deref_mut(&mut self) -> &mut Self::Target {
63        &mut self.0
64    }
65}
66impl Hash for RoaringBitmapWrapper {
67    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
68        struct HasherWriter<'a, H: std::hash::Hasher>(&'a mut H);
69        impl<H: std::hash::Hasher> std::io::Write for HasherWriter<'_, H> {
70            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
71                self.0.write(buf);
72                Ok(buf.len())
73            }
74            fn flush(&mut self) -> std::io::Result<()> {
75                Ok(())
76            }
77        }
78        self.0.serialize_into(HasherWriter(state)).unwrap();
79    }
80}
81
82#[turbo_tasks::value(transparent, cell = "keyed")]
83pub struct ModuleToChunkGroups(FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper>);
84
85#[turbo_tasks::value]
86pub struct ChunkGroupInfo {
87    pub module_chunk_groups: ResolvedVc<ModuleToChunkGroups>,
88    #[bincode(with = "turbo_bincode::indexset")]
89    pub chunk_groups: FxIndexSet<ChunkGroup>,
90    #[turbo_tasks(unsafe_ignore)]
91    #[bincode(with = "turbo_bincode::indexset")]
92    pub chunk_group_keys: FxIndexSet<ChunkGroupKey>,
93    pub chunking_heuristics: ChunkingHeuristicsInfo,
94}
95
96/// Chunking heuristics computed by [`compute_chunk_group_info`]. `clusters` is indexed by
97/// chunk-group index (same length and order as [`ChunkGroupInfo::chunk_groups`]); `priority_routes`
98/// is a set of those indices.
99#[derive(Debug, Default, Clone, PartialEq, Eq, ValueDebugFormat, NonLocalValue, Encode, Decode)]
100pub struct ChunkingHeuristicsInfo {
101    /// For each chunk group (by index), the set of cluster IDs it belongs to. A cluster ID is the
102    /// index of a configured cluster. A route's chunk group carries that route's clusters; chunk
103    /// groups it pulls in inherit them.
104    ///
105    /// Example: `clusters[5] = [0, 2]` — chunk group 5 is part of clusters 0 and 2.
106    pub clusters: Vec<Vec<u16>>,
107    /// The set of chunk-group indices that belong to a priority route: the priority
108    /// routes themselves, plus every chunk group they pull in.
109    ///
110    /// Example: `priority_routes = {3, 7}` — chunk groups 3 and 7 are served by a priority
111    /// route; any group not in the set (e.g. 4) is not.
112    pub priority_routes: RoaringBitmapWrapper,
113}
114
115#[turbo_tasks::value_impl]
116impl ChunkGroupInfo {
117    #[turbo_tasks::function]
118    pub fn module_chunk_groups(&self) -> Vc<ModuleToChunkGroups> {
119        *self.module_chunk_groups
120    }
121
122    #[turbo_tasks::function]
123    pub async fn get_index_of(&self, chunk_group: ChunkGroup) -> Result<Vc<usize>> {
124        if let Some(idx) = self.chunk_groups.get_index_of(&chunk_group) {
125            Ok(Vc::cell(idx))
126        } else {
127            if cfg!(debug_assertions) {
128                bail!(
129                    "Couldn't find chunk group index for {} in {}",
130                    chunk_group.debug_str(self).await?,
131                    self.chunk_groups
132                        .iter()
133                        .map(|c| c.debug_str(self))
134                        .try_join()
135                        .await?
136                        .join(", ")
137                );
138            } else {
139                bail!("Couldn't find chunk group index")
140            }
141        }
142    }
143}
144
145/// Per-entry chunking heuristics.
146#[turbo_tasks::task_input]
147#[derive(Debug, Default, Clone, Hash, PartialEq, Eq, Encode, Decode)]
148pub struct EntryHeuristics {
149    /// Cluster indices this route belongs to.
150    pub clusters: Vec<u16>,
151    pub high_priority: bool,
152}
153
154impl EntryHeuristics {
155    /// Heuristics for an entry that is a high-priority route: belongs to no clusters and is marked
156    /// as high priority.
157    pub fn high_priority() -> Self {
158        Self {
159            clusters: Vec::new(),
160            high_priority: true,
161        }
162    }
163}
164
165/// See [ChunkGroup] for documentation
166#[turbo_tasks::task_input]
167#[derive(Debug, Clone, Hash, PartialEq, Eq, Encode, Decode)]
168pub enum ChunkGroupEntry {
169    Entry {
170        modules: Vec<ResolvedVc<Box<dyn Module>>>,
171        heuristics: EntryHeuristics,
172    },
173    Async(ResolvedVc<Box<dyn Module>>),
174    Isolated(ResolvedVc<Box<dyn Module>>),
175    IsolatedMerged {
176        parent: Box<ChunkGroupEntry>,
177        merge_tag: RcStr,
178        entries: Vec<ResolvedVc<Box<dyn Module>>>,
179    },
180    Shared(ResolvedVc<Box<dyn Module>>),
181    SharedMultiple(Vec<ResolvedVc<Box<dyn Module>>>),
182    SharedMerged {
183        parent: Box<ChunkGroupEntry>,
184        merge_tag: RcStr,
185        entries: Vec<ResolvedVc<Box<dyn Module>>>,
186    },
187}
188impl ChunkGroupEntry {
189    pub fn entries(&self) -> impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + '_ {
190        match self {
191            Self::Async(e) | Self::Isolated(e) | Self::Shared(e) => {
192                Either::Left(std::iter::once(*e))
193            }
194            Self::Entry {
195                modules: entries, ..
196            }
197            | Self::IsolatedMerged { entries, .. }
198            | Self::SharedMultiple(entries)
199            | Self::SharedMerged { entries, .. } => Either::Right(entries.iter().copied()),
200        }
201    }
202}
203
204#[turbo_tasks::task_input]
205#[derive(Debug, Clone, Hash, PartialEq, Eq, Encode, Decode)]
206pub enum ChunkGroup {
207    /// The entry chunk group of the compilation, e.g. src/index.js for a SPA, or app/foo/page.js
208    /// for Next.js.
209    Entry(Vec<ResolvedVc<Box<dyn Module>>>),
210    /// An async chunk group. Corresponds to an incoming [ChunkingType::Async] reference
211    Async(ResolvedVc<Box<dyn Module>>),
212    /// An isolated chunk group. Corresponds to an incoming [ChunkingType::Isolated] reference with
213    /// `merge_tag: None`
214    Isolated(ResolvedVc<Box<dyn Module>>),
215    /// An isolated chunk group. Corresponds to an incoming [ChunkingType::Isolated] reference with
216    /// `merge_tag: Some(_)`
217    IsolatedMerged {
218        parent: usize,
219        merge_tag: RcStr,
220        entries: Vec<ResolvedVc<Box<dyn Module>>>,
221    },
222    /// A shared chunk group. Corresponds to an incoming [ChunkingType::Shared] reference with
223    /// `merge_tag: None`
224    Shared(ResolvedVc<Box<dyn Module>>),
225    /// A shared chunk group with multiple entries.
226    SharedMultiple(Vec<ResolvedVc<Box<dyn Module>>>),
227    /// A shared chunk group. Corresponds to an incoming [ChunkingType::Shared] reference with
228    /// `merge_tag: Some(_)`
229    SharedMerged {
230        parent: usize,
231        merge_tag: RcStr,
232        entries: Vec<ResolvedVc<Box<dyn Module>>>,
233    },
234    /// a module with an incoming collected edge
235    Collected(ResolvedVc<Box<dyn Module>>),
236}
237
238impl ChunkGroup {
239    /// Returns the parent group when this chunk group is a merged group. In that case `entries()`
240    /// are in unspecified order.
241    pub fn get_merged_parent(&self) -> Option<usize> {
242        match self {
243            ChunkGroup::IsolatedMerged { parent, .. } | ChunkGroup::SharedMerged { parent, .. } => {
244                Some(*parent)
245            }
246            _ => None,
247        }
248    }
249
250    /// Iterates over the entries of the chunk group. When `get_merged_parent` is Some, the order is
251    /// unspecified.
252    pub fn entries(&self) -> impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + Clone + '_ {
253        match self {
254            ChunkGroup::Async(e)
255            | ChunkGroup::Isolated(e)
256            | ChunkGroup::Shared(e)
257            | ChunkGroup::Collected(e) => Either::Left(std::iter::once(*e)),
258            ChunkGroup::Entry(entries)
259            | ChunkGroup::IsolatedMerged { entries, .. }
260            | ChunkGroup::SharedMultiple(entries)
261            | ChunkGroup::SharedMerged { entries, .. } => Either::Right(entries.iter().copied()),
262        }
263    }
264
265    pub fn entries_count(&self) -> usize {
266        match self {
267            ChunkGroup::Async(_)
268            | ChunkGroup::Isolated(_)
269            | ChunkGroup::Shared(_)
270            | ChunkGroup::Collected(_) => 1,
271            ChunkGroup::Entry(entries)
272            | ChunkGroup::IsolatedMerged { entries, .. }
273            | ChunkGroup::SharedMultiple(entries)
274            | ChunkGroup::SharedMerged { entries, .. } => entries.len(),
275        }
276    }
277
278    pub async fn debug_str(&self, chunk_group_info: &ChunkGroupInfo) -> Result<String> {
279        Ok(match self {
280            ChunkGroup::Entry(entries) => format!(
281                "ChunkGroup::Entry({:?})",
282                entries
283                    .iter()
284                    .map(|m| m.ident().to_string())
285                    .try_join()
286                    .await?
287            ),
288            ChunkGroup::Async(entry) => turbofmt!("ChunkGroup::Async({:?})", entry.ident())
289                .await?
290                .to_string(),
291            ChunkGroup::Collected(entry) => turbofmt!("ChunkGroup::Collected({:?})", entry.ident())
292                .await?
293                .to_string(),
294            ChunkGroup::Isolated(entry) => turbofmt!("ChunkGroup::Isolated({:?})", entry.ident())
295                .await?
296                .to_string(),
297            ChunkGroup::Shared(entry) => turbofmt!("ChunkGroup::Shared({:?})", entry.ident())
298                .await?
299                .to_string(),
300            ChunkGroup::SharedMultiple(entries) => format!(
301                "ChunkGroup::SharedMultiple({:?})",
302                entries
303                    .iter()
304                    .map(|m| m.ident().to_string())
305                    .try_join()
306                    .await?
307            ),
308            ChunkGroup::IsolatedMerged {
309                parent,
310                merge_tag,
311                entries,
312            } => {
313                format!(
314                    "ChunkGroup::IsolatedMerged({}, {}, {:?})",
315                    Box::pin(chunk_group_info.chunk_groups[*parent].debug_str(chunk_group_info))
316                        .await?,
317                    merge_tag,
318                    entries
319                        .iter()
320                        .map(|m| m.ident().to_string())
321                        .try_join()
322                        .await?
323                )
324            }
325            ChunkGroup::SharedMerged {
326                parent,
327                merge_tag,
328                entries,
329            } => {
330                format!(
331                    "ChunkGroup::SharedMerged({}, {}, {:?})",
332                    Box::pin(chunk_group_info.chunk_groups[*parent].debug_str(chunk_group_info))
333                        .await?,
334                    merge_tag,
335                    entries
336                        .iter()
337                        .map(|m| m.ident().to_string())
338                        .try_join()
339                        .await?
340                )
341            }
342        })
343    }
344}
345
346/// See [ChunkGroup] for documentation
347#[derive(Debug, Clone, Hash, PartialEq, Eq, Encode, Decode)]
348pub enum ChunkGroupKey {
349    Entry(Vec<ResolvedVc<Box<dyn Module>>>),
350    Async(ResolvedVc<Box<dyn Module>>),
351    Isolated(ResolvedVc<Box<dyn Module>>),
352    IsolatedMerged {
353        parent: ChunkGroupId,
354        merge_tag: RcStr,
355    },
356    Shared(ResolvedVc<Box<dyn Module>>),
357    SharedMultiple(Vec<ResolvedVc<Box<dyn Module>>>),
358    SharedMerged {
359        parent: ChunkGroupId,
360        merge_tag: RcStr,
361    },
362    /// a module with an incoming collected edge
363    Collected(ResolvedVc<Box<dyn Module>>),
364}
365
366impl ChunkGroupKey {
367    pub async fn debug_str(
368        &self,
369        keys: impl std::ops::Index<usize, Output = Self>,
370    ) -> Result<String> {
371        Ok(match self {
372            ChunkGroupKey::Entry(entries) => format!(
373                "Entry({:?})",
374                entries
375                    .iter()
376                    .map(|m| m.ident().to_string())
377                    .try_join()
378                    .await?
379            ),
380            ChunkGroupKey::Async(module) => {
381                turbofmt!("Async({:?})", module.ident()).await?.to_string()
382            }
383            ChunkGroupKey::Collected(module) => turbofmt!("Collected({:?})", module.ident())
384                .await?
385                .to_string(),
386            ChunkGroupKey::Isolated(module) => turbofmt!("Isolated({:?})", module.ident())
387                .await?
388                .to_string(),
389            ChunkGroupKey::IsolatedMerged { parent, merge_tag } => {
390                format!(
391                    "IsolatedMerged {{ parent: {}, merge_tag: {:?} }}",
392                    Box::pin(keys.index(parent.0 as usize).clone().debug_str(keys)).await?,
393                    merge_tag
394                )
395            }
396            ChunkGroupKey::Shared(module) => {
397                turbofmt!("Shared({:?})", module.ident()).await?.to_string()
398            }
399            ChunkGroupKey::SharedMultiple(entries) => format!(
400                "SharedMultiple({:?})",
401                entries
402                    .iter()
403                    .map(|m| m.ident().to_string())
404                    .try_join()
405                    .await?
406            ),
407            ChunkGroupKey::SharedMerged { parent, merge_tag } => {
408                format!(
409                    "SharedMerged {{ parent: {}, merge_tag: {:?} }}",
410                    Box::pin(keys.index(parent.0 as usize).clone().debug_str(keys)).await?,
411                    merge_tag
412                )
413            }
414        })
415    }
416}
417
418#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Encode, Decode)]
419pub struct ChunkGroupId(u32);
420
421impl From<usize> for ChunkGroupId {
422    fn from(id: usize) -> Self {
423        Self(id as u32)
424    }
425}
426
427impl Deref for ChunkGroupId {
428    type Target = u32;
429    fn deref(&self) -> &Self::Target {
430        &self.0
431    }
432}
433
434#[derive(Debug, Clone, PartialEq, Eq)]
435pub(crate) struct TraversalPriority {
436    pub depth: usize,
437    pub chunk_group_len: u64,
438}
439impl PartialOrd for TraversalPriority {
440    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
441        Some(self.cmp(other))
442    }
443}
444// TODO automatically derive this using Reverse<_>
445impl Ord for TraversalPriority {
446    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
447        // BinaryHeap prioritizes high values
448
449        // Smaller depth has higher priority
450        let depth_order = self.depth.cmp(&other.depth).reverse();
451        // Smaller group length has higher priority
452        let chunk_group_len_order = self.chunk_group_len.cmp(&other.chunk_group_len).reverse();
453
454        depth_order.then(chunk_group_len_order)
455    }
456}
457
458pub async fn compute_chunk_group_info(graph: &ModuleGraph) -> Result<Vc<ChunkGroupInfo>> {
459    let span_outer = tracing::info_span!(
460        "compute chunk group info",
461        module_count = tracing::field::Empty,
462        visit_count = tracing::field::Empty,
463        chunk_group_count = tracing::field::Empty
464    );
465
466    let span = span_outer.clone();
467    async move {
468        let mut chunk_groups_map: FxIndexMap<
469            ChunkGroupKey,
470            FxIndexSet<ResolvedVc<Box<dyn Module>>>,
471        > = FxIndexMap::default();
472
473        // For each module, the indices in the bitmap store which chunk groups in `chunk_groups_map`
474        // that module is part of.
475        let mut module_chunk_groups: FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper> =
476            FxHashMap::default();
477
478        let module_count = graph
479            .graphs
480            .iter()
481            .map(|g| g.graph.node_count())
482            .sum::<usize>();
483        span.record("module_count", module_count);
484
485        // use all entries from all graphs
486        let entries = graph.all_chunk_group_entries().collect::<Vec<_>>();
487
488        // First, compute the depth for each module in the graph and initialize
489        // `module_chunk_groups` with empty bitmaps.
490        let module_depth: FxHashMap<ResolvedVc<Box<dyn Module>>, usize> = {
491            let mut module_depth =
492                FxHashMap::with_capacity_and_hasher(module_count, Default::default());
493            graph.traverse_edges_bfs(
494                entries.iter().flat_map(|e| e.entries()),
495                |parent, node| {
496                    if let Some((parent, _)) = parent {
497                        let parent_depth = *module_depth
498                            .get(&parent)
499                            .context("Module depth not found")?;
500                        module_depth.entry(node).or_insert(parent_depth + 1);
501                    } else {
502                        module_depth.insert(node, 0);
503                    };
504
505                    module_chunk_groups.insert(node, RoaringBitmapWrapper::default());
506
507                    Ok(GraphTraversalAction::Continue)
508                },
509            )?;
510            module_depth
511        };
512
513        // ----
514
515        fn entry_to_chunk_group_id(
516            entry: ChunkGroupEntry,
517            chunk_groups_map: &mut FxIndexMap<
518                ChunkGroupKey,
519                FxIndexSet<ResolvedVc<Box<dyn Module>>>,
520            >,
521        ) -> ChunkGroupKey {
522            match entry {
523                ChunkGroupEntry::Entry { modules, .. } => ChunkGroupKey::Entry(modules),
524                ChunkGroupEntry::Async(entry) => ChunkGroupKey::Async(entry),
525                ChunkGroupEntry::Isolated(entry) => ChunkGroupKey::Isolated(entry),
526                ChunkGroupEntry::Shared(entry) => ChunkGroupKey::Shared(entry),
527                ChunkGroupEntry::SharedMultiple(entries) => ChunkGroupKey::SharedMultiple(entries),
528                ChunkGroupEntry::IsolatedMerged {
529                    parent,
530                    merge_tag,
531                    entries: _,
532                } => {
533                    let parent = entry_to_chunk_group_id(*parent, chunk_groups_map);
534                    let parent_entry = chunk_groups_map.entry(parent);
535                    let parent_id = parent_entry.index();
536                    parent_entry.or_default();
537
538                    ChunkGroupKey::IsolatedMerged {
539                        parent: ChunkGroupId::from(parent_id),
540                        merge_tag,
541                    }
542                }
543                ChunkGroupEntry::SharedMerged {
544                    parent,
545                    merge_tag,
546                    entries: _,
547                } => {
548                    let parent = entry_to_chunk_group_id(*parent, chunk_groups_map);
549                    let parent_entry = chunk_groups_map.entry(parent);
550                    let parent_id = parent_entry.index();
551                    parent_entry.or_default();
552
553                    ChunkGroupKey::SharedMerged {
554                        parent: ChunkGroupId::from(parent_id),
555                        merge_tag,
556                    }
557                }
558            }
559        }
560
561        let entry_chunk_group_keys = entries
562            .iter()
563            .flat_map(|&chunk_group| {
564                let chunk_group_key =
565                    entry_to_chunk_group_id(chunk_group.clone(), &mut chunk_groups_map);
566                chunk_group
567                    .entries()
568                    .map(move |e| (e, chunk_group_key.clone()))
569            })
570            .collect::<FxHashMap<_, _>>();
571
572        // `inherits_from[source]` is the set of chunk groups that inherit heuristics from `source`.
573        let mut inherits_from: FxHashMap<u32, RoaringBitmap> = FxHashMap::default();
574
575        let visit_count = graph.traverse_edges_fixed_point_with_priority(
576            entries
577                .iter()
578                .flat_map(|e| e.entries())
579                .map(|e| {
580                    Ok((
581                        e,
582                        TraversalPriority {
583                            depth: *module_depth.get(&e).context("Module depth not found")?,
584                            chunk_group_len: 0,
585                        },
586                    ))
587                })
588                .collect::<Result<Vec<_>>>()?,
589            &mut module_chunk_groups,
590            |parent_info: Option<(ResolvedVc<Box<dyn Module>>, &'_ RefData, _)>,
591             node: ResolvedVc<Box<dyn Module>>,
592             _,
593             module_chunk_groups: &mut FxHashMap<
594                ResolvedVc<Box<dyn Module>>,
595                RoaringBitmapWrapper,
596            >|
597             -> Result<GraphTraversalAction> {
598                enum ChunkGroupInheritance<It: Iterator<Item = ChunkGroupKey>> {
599                    Inherit(ResolvedVc<Box<dyn Module>>),
600                    ChunkGroup(It),
601                }
602                let chunk_groups = if let Some((parent, ref_data, _)) = parent_info {
603                    match &ref_data.chunking_type {
604                        ChunkingType::Parallel { .. } => ChunkGroupInheritance::Inherit(parent),
605                        ChunkingType::Async => ChunkGroupInheritance::ChunkGroup(Either::Left(
606                            std::iter::once(ChunkGroupKey::Async(node)),
607                        )),
608                        ChunkingType::Isolated {
609                            merge_tag: None, ..
610                        } => ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
611                            ChunkGroupKey::Isolated(node),
612                        ))),
613                        ChunkingType::Shared {
614                            merge_tag: None, ..
615                        } => ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
616                            ChunkGroupKey::Shared(node),
617                        ))),
618                        ChunkingType::Isolated {
619                            merge_tag: Some(merge_tag),
620                            ..
621                        } => {
622                            let parents = module_chunk_groups
623                                .get(&parent)
624                                .context("Module chunk group not found")?;
625                            let chunk_groups =
626                                parents.iter().map(|parent| ChunkGroupKey::IsolatedMerged {
627                                    parent: ChunkGroupId(parent),
628                                    merge_tag: merge_tag.clone(),
629                                });
630                            ChunkGroupInheritance::ChunkGroup(Either::Right(Either::Left(
631                                chunk_groups,
632                            )))
633                        }
634                        ChunkingType::Shared {
635                            merge_tag: Some(merge_tag),
636                            ..
637                        } => {
638                            let parents = module_chunk_groups
639                                .get(&parent)
640                                .context("Module chunk group not found")?;
641                            let chunk_groups =
642                                parents.iter().map(|parent| ChunkGroupKey::SharedMerged {
643                                    parent: ChunkGroupId(parent),
644                                    merge_tag: merge_tag.clone(),
645                                });
646                            ChunkGroupInheritance::ChunkGroup(Either::Right(Either::Right(
647                                chunk_groups,
648                            )))
649                        }
650                        ChunkingType::Traced { .. } => {
651                            // Traced modules are not placed in chunk groups
652                            return Ok(GraphTraversalAction::Skip);
653                        }
654                        ChunkingType::Emitted {
655                            namespace: _,
656                            emit_to_all_entries: _,
657                        } => {
658                            // TODO ideally this would get the chunk group bitset of the parent
659                            // ChunkGroup::Entry
660                            ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
661                                ChunkGroupKey::Collected(node),
662                            )))
663                        }
664                        ChunkingType::Collected { .. } => {
665                            // These only exist in the module_batches graph.
666                            unreachable!();
667                        }
668                        ChunkingType::PerEntry => {
669                            // A collecting module belongs wherever its importer belongs. The
670                            // per-entry behavior is carried by the `chunk_group` argument of
671                            // `CollectingModule::as_chunk_item`, not by chunk group membership.
672                            // This matches `collect.rs`, which inherits the parent's entry
673                            // membership across this edge.
674                            ChunkGroupInheritance::Inherit(parent)
675                        }
676                    }
677                } else {
678                    ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
679                        // TODO remove clone
680                        entry_chunk_group_keys
681                            .get(&node)
682                            .context("Module chunk group not found")?
683                            .clone(),
684                    )))
685                };
686
687                Ok(match chunk_groups {
688                    ChunkGroupInheritance::ChunkGroup(chunk_groups) => {
689                        // Start of a new chunk group, don't inherit anything from parent
690                        let chunk_group_ids = chunk_groups.map(|chunk_group| {
691                            // For merged groups, the parent group id whose heuristics they inherit.
692                            let merged_parent = match &chunk_group {
693                                ChunkGroupKey::IsolatedMerged { parent, .. }
694                                | ChunkGroupKey::SharedMerged { parent, .. } => Some(parent.0),
695                                _ => None,
696                            };
697                            let id = match chunk_groups_map.entry(chunk_group) {
698                                Entry::Occupied(mut e) => {
699                                    let id = e.index() as u32;
700                                    if merged_parent.is_some() {
701                                        e.get_mut().insert(node);
702                                    }
703                                    id
704                                }
705                                Entry::Vacant(e) => {
706                                    let id = e.index() as u32;
707                                    let mut set = FxIndexSet::default();
708                                    if merged_parent.is_some() {
709                                        set.insert(node);
710                                    }
711                                    e.insert(set);
712                                    id
713                                }
714                            };
715                            // Record heuristics-inheritance edges into this chunk group: merged
716                            // groups inherit from their specific parent group; all other groups
717                            // inherit from every chunk group of the referencing module.
718                            if let Some(parent) = merged_parent {
719                                inherits_from.entry(parent).or_default().insert(id);
720                            } else if let Some((parent_module, _, _)) = parent_info
721                                && let Some(parent_groups) = module_chunk_groups.get(&parent_module)
722                            {
723                                for source in parent_groups.iter() {
724                                    inherits_from.entry(source).or_default().insert(id);
725                                }
726                            }
727                            id
728                        });
729
730                        let chunk_groups =
731                            RoaringBitmapWrapper(RoaringBitmap::from_iter(chunk_group_ids));
732
733                        // Assign chunk group to the target node (the entry of the chunk group)
734                        let bitset = module_chunk_groups
735                            .get_mut(&node)
736                            .context("Module chunk group not found")?;
737                        if chunk_groups.is_proper_superset(bitset) {
738                            // Add bits from parent, and continue traversal because changed
739                            **bitset |= chunk_groups.into_inner();
740
741                            GraphTraversalAction::Continue
742                        } else {
743                            // Unchanged, no need to forward to children
744                            GraphTraversalAction::Skip
745                        }
746                    }
747                    ChunkGroupInheritance::Inherit(parent) => {
748                        // Inherit chunk groups from parent, merge parent chunk groups into
749                        // current
750
751                        if parent == node {
752                            // A self-reference
753                            GraphTraversalAction::Skip
754                        } else {
755                            let [Some(parent_chunk_groups), Some(current_chunk_groups)] =
756                                module_chunk_groups.get_disjoint_mut([&parent, &node])
757                            else {
758                                // All modules are inserted in the previous iteration
759                                // Technically unreachable, but could be reached due to eventual
760                                // consistency
761                                bail!("Module chunk groups not found");
762                            };
763
764                            if current_chunk_groups.is_empty() {
765                                // Initial visit, clone instead of merging
766                                *current_chunk_groups = parent_chunk_groups.clone();
767                                GraphTraversalAction::Continue
768                            } else if parent_chunk_groups.is_proper_superset(current_chunk_groups) {
769                                // Add bits from parent, and continue traversal because changed
770                                **current_chunk_groups |= &**parent_chunk_groups;
771                                GraphTraversalAction::Continue
772                            } else {
773                                // Unchanged, no need to forward to children
774                                GraphTraversalAction::Skip
775                            }
776                        }
777                    }
778                })
779            },
780            // This priority is used as a heuristic to keep the number of retraversals down, by
781            // - keeping it similar to a BFS via the depth priority
782            // - prioritizing smaller chunk groups which are expected to themselves reference
783            //   bigger chunk groups (i.e. shared code deeper down in the graph).
784            //
785            // Both try to first visit modules with a large dependency subgraph first (which
786            // would be higher in the graph and are included by few chunks themselves).
787            |successor, module_chunk_groups| {
788                Ok(TraversalPriority {
789                    depth: *module_depth
790                        .get(&successor)
791                        .context("Module depth not found")?,
792                    chunk_group_len: module_chunk_groups
793                        .get(&successor)
794                        .context("Module chunk group not found")?
795                        .len(),
796                })
797            },
798        )?;
799
800        span.record("visit_count", visit_count);
801        span.record("chunk_group_count", chunk_groups_map.len());
802
803        static PRINT_CHUNK_GROUP_INFO: LazyLock<bool> =
804            LazyLock::new(|| match std::env::var_os("TURBOPACK_PRINT_CHUNK_GROUPS") {
805                Some(v) => v == "1",
806                None => false,
807            });
808        if *PRINT_CHUNK_GROUP_INFO {
809            use std::{
810                collections::{BTreeMap, BTreeSet},
811                path::absolute,
812            };
813
814            let mut buckets = BTreeMap::default();
815            for (module, key) in &module_chunk_groups {
816                if !key.is_empty() {
817                    buckets
818                        .entry(key.iter().collect::<Vec<_>>())
819                        .or_insert(BTreeSet::new())
820                        .insert(module.ident().to_string().await?);
821                }
822            }
823
824            let mut result = vec![];
825            result.push("Chunk Groups:".to_string());
826            for (i, (key, _)) in chunk_groups_map.iter().enumerate() {
827                result.push(format!(
828                    "  {:?}: {}",
829                    i,
830                    key.debug_str(chunk_groups_map.keys()).await?
831                ));
832            }
833            result.push("# Module buckets:".to_string());
834            for (key, modules) in buckets.iter() {
835                result.push(format!("## {:?}:", key.iter().collect::<Vec<_>>()));
836                for module in modules {
837                    result.push(format!("  {module}"));
838                }
839                result.push("".to_string());
840            }
841            let f = absolute(format!("chunk_group_info_{}.log", visit_count))?;
842            println!("Wrote Chunk Group Info to {}", f.display());
843            std::fs::write(f, result.join("\n"))?;
844        }
845
846        // Resolve per-chunk-group chunking heuristics. Entry
847        // chunk groups carry their route's clusters / priority-route flag; other chunk groups
848        // inherit the union of clusters (and OR of the flag) from their referencing chunk groups.
849        let mut clusters: Vec<RoaringBitmap> = vec![RoaringBitmap::new(); chunk_groups_map.len()];
850        let mut priority_routes = RoaringBitmap::new();
851
852        let mut worklist: Vec<usize> = Vec::new();
853
854        for chunk_group in &entries {
855            let ChunkGroupEntry::Entry {
856                modules,
857                heuristics,
858            } = chunk_group
859            else {
860                continue;
861            };
862            if heuristics.clusters.is_empty() && !heuristics.high_priority {
863                continue;
864            }
865            if let Some(index) =
866                chunk_groups_map.get_index_of(&ChunkGroupKey::Entry(modules.clone()))
867            {
868                if clusters[index].is_empty() && !priority_routes.contains(index as u32) {
869                    worklist.push(index);
870                }
871                clusters[index].extend(heuristics.clusters.iter().map(|&c| c as u32));
872                if heuristics.high_priority {
873                    priority_routes.insert(index as u32);
874                }
875            }
876        }
877
878        while let Some(source) = worklist.pop() {
879            let source_priority_route = priority_routes.contains(source as u32);
880            let Some(targets) = inherits_from.get(&(source as u32)) else {
881                continue;
882            };
883            for target in targets.iter() {
884                let target = target as usize;
885                if target == source {
886                    continue;
887                }
888                let [source_clusters, target_clusters] =
889                    clusters.get_disjoint_mut([source, target]).unwrap();
890                let previous_target_clusters_len = target_clusters.len();
891                *target_clusters |= &*source_clusters;
892                let changed = (source_priority_route && priority_routes.insert(target as u32))
893                    || previous_target_clusters_len != target_clusters.len();
894                if changed {
895                    worklist.push(target);
896                }
897            }
898        }
899
900        let chunk_group_clusters: Vec<Vec<u16>> = clusters
901            .into_iter()
902            .map(|bm| bm.iter().map(|id| id as u16).collect())
903            .collect();
904        let chunk_group_priority_routes = RoaringBitmapWrapper(priority_routes);
905
906        Ok(ChunkGroupInfo {
907            module_chunk_groups: ResolvedVc::cell(module_chunk_groups),
908            chunk_group_keys: chunk_groups_map.keys().cloned().collect(),
909            chunking_heuristics: ChunkingHeuristicsInfo {
910                clusters: chunk_group_clusters,
911                priority_routes: chunk_group_priority_routes,
912            },
913            chunk_groups: chunk_groups_map
914                .into_iter()
915                .map(|(k, merged_entries)| match k {
916                    ChunkGroupKey::Entry(entries) => ChunkGroup::Entry(entries),
917                    ChunkGroupKey::Async(module) => ChunkGroup::Async(module),
918                    ChunkGroupKey::Collected(module) => ChunkGroup::Collected(module),
919                    ChunkGroupKey::Isolated(module) => ChunkGroup::Isolated(module),
920                    ChunkGroupKey::IsolatedMerged { parent, merge_tag } => {
921                        ChunkGroup::IsolatedMerged {
922                            parent: parent.0 as usize,
923                            merge_tag,
924                            entries: merged_entries.into_iter().collect(),
925                        }
926                    }
927                    ChunkGroupKey::Shared(module) => ChunkGroup::Shared(module),
928                    ChunkGroupKey::SharedMultiple(entries) => ChunkGroup::SharedMultiple(entries),
929                    ChunkGroupKey::SharedMerged { parent, merge_tag } => ChunkGroup::SharedMerged {
930                        parent: parent.0 as usize,
931                        merge_tag,
932                        entries: merged_entries.into_iter().collect(),
933                    },
934                })
935                .collect(),
936        }
937        .cell())
938    }
939    .instrument(span_outer)
940    .await
941}
turbopack_core/module_graph/chunk_group_info.rs

turbopack_core/module_graph/
chunk_group_info.rs
