turbopack_core/module_graph/
chunk_group_info.rs

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

turbopack_core/module_graph/
chunk_group_info.rs
