1use std::{
2 hash::Hash,
3 ops::{Deref, DerefMut},
4};
5
6use anyhow::{Context, Result, bail};
7use either::Either;
8use indexmap::map::Entry;
9use roaring::RoaringBitmap;
10use rustc_hash::FxHashMap;
11use serde::{Deserialize, Serialize};
12use tracing::Instrument;
13use turbo_rcstr::RcStr;
14use turbo_tasks::{
15 FxIndexMap, FxIndexSet, NonLocalValue, ResolvedVc, TaskInput, TryJoinIterExt, ValueToString,
16 Vc, debug::ValueDebugFormat, trace::TraceRawVcs,
17};
18
19use crate::{
20 chunk::ChunkingType,
21 module::Module,
22 module_graph::{GraphTraversalAction, ModuleGraph, RefData, SingleModuleGraphModuleNode},
23};
24
25#[derive(
26 Clone, Debug, Default, PartialEq, Serialize, Deserialize, TraceRawVcs, ValueDebugFormat,
27)]
28pub struct RoaringBitmapWrapper(#[turbo_tasks(trace_ignore)] pub RoaringBitmap);
29
30impl TaskInput for RoaringBitmapWrapper {
31 fn is_transient(&self) -> bool {
32 false
33 }
34}
35
36impl RoaringBitmapWrapper {
37 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
50impl 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]
83pub struct ChunkGroupInfo {
84 pub module_chunk_groups: FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper>,
85 #[turbo_tasks(trace_ignore)]
86 pub chunk_groups: FxIndexSet<ChunkGroup>,
87 #[turbo_tasks(trace_ignore)]
88 pub chunk_group_keys: FxIndexSet<ChunkGroupKey>,
89}
90
91#[turbo_tasks::value_impl]
92impl ChunkGroupInfo {
93 #[turbo_tasks::function]
94 pub async fn get_index_of(&self, chunk_group: ChunkGroup) -> Result<Vc<usize>> {
95 if let Some(idx) = self.chunk_groups.get_index_of(&chunk_group) {
96 Ok(Vc::cell(idx))
97 } else {
98 bail!(
99 "Couldn't find chunk group index for {} in {}",
100 chunk_group.debug_str(self).await?,
101 self.chunk_groups
102 .iter()
103 .map(|c| c.debug_str(self))
104 .try_join()
105 .await?
106 .join(", ")
107 );
108 }
109 }
110}
111
112#[derive(
113 Debug, Clone, Hash, TaskInput, PartialEq, Eq, Serialize, Deserialize, TraceRawVcs, NonLocalValue,
114)]
115pub enum ChunkGroupEntry {
116 Entry(Vec<ResolvedVc<Box<dyn Module>>>),
118 Async(ResolvedVc<Box<dyn Module>>),
120 Isolated(ResolvedVc<Box<dyn Module>>),
122 IsolatedMerged {
124 parent: Box<ChunkGroupEntry>,
125 merge_tag: RcStr,
126 entries: Vec<ResolvedVc<Box<dyn Module>>>,
127 },
128 Shared(ResolvedVc<Box<dyn Module>>),
130 SharedMerged {
132 parent: Box<ChunkGroupEntry>,
133 merge_tag: RcStr,
134 entries: Vec<ResolvedVc<Box<dyn Module>>>,
135 },
136}
137impl ChunkGroupEntry {
138 pub fn entries(&self) -> impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + '_ {
139 match self {
140 Self::Async(e) | Self::Isolated(e) | Self::Shared(e) => {
141 Either::Left(std::iter::once(*e))
142 }
143 Self::Entry(entries)
144 | Self::IsolatedMerged { entries, .. }
145 | Self::SharedMerged { entries, .. } => Either::Right(entries.iter().copied()),
146 }
147 }
148}
149
150#[derive(Debug, Clone, Hash, TaskInput, PartialEq, Eq, Serialize, Deserialize, TraceRawVcs)]
151pub enum ChunkGroup {
152 Entry(Vec<ResolvedVc<Box<dyn Module>>>),
154 Async(ResolvedVc<Box<dyn Module>>),
156 Isolated(ResolvedVc<Box<dyn Module>>),
158 IsolatedMerged {
160 parent: usize,
161 merge_tag: RcStr,
162 entries: Vec<ResolvedVc<Box<dyn Module>>>,
163 },
164 Shared(ResolvedVc<Box<dyn Module>>),
166 SharedMerged {
168 parent: usize,
169 merge_tag: RcStr,
170 entries: Vec<ResolvedVc<Box<dyn Module>>>,
171 },
172}
173
174impl ChunkGroup {
175 pub fn get_merged_parent(&self) -> Option<usize> {
178 match self {
179 ChunkGroup::IsolatedMerged { parent, .. } | ChunkGroup::SharedMerged { parent, .. } => {
180 Some(*parent)
181 }
182 _ => None,
183 }
184 }
185
186 pub fn entries(&self) -> impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + Clone + '_ {
189 match self {
190 ChunkGroup::Async(e) | ChunkGroup::Isolated(e) | ChunkGroup::Shared(e) => {
191 Either::Left(std::iter::once(*e))
192 }
193 ChunkGroup::Entry(entries)
194 | ChunkGroup::IsolatedMerged { entries, .. }
195 | ChunkGroup::SharedMerged { entries, .. } => Either::Right(entries.iter().copied()),
196 }
197 }
198
199 pub fn entries_count(&self) -> usize {
200 match self {
201 ChunkGroup::Async(_) | ChunkGroup::Isolated(_) | ChunkGroup::Shared(_) => 1,
202 ChunkGroup::Entry(entries)
203 | ChunkGroup::IsolatedMerged { entries, .. }
204 | ChunkGroup::SharedMerged { entries, .. } => entries.len(),
205 }
206 }
207
208 pub async fn debug_str(&self, chunk_group_info: &ChunkGroupInfo) -> Result<String> {
209 Ok(match self {
210 ChunkGroup::Entry(entries) => format!(
211 "ChunkGroup::Entry({:?})",
212 entries
213 .iter()
214 .map(|m| m.ident().to_string())
215 .try_join()
216 .await?
217 ),
218 ChunkGroup::Async(entry) => {
219 format!("ChunkGroup::Async({:?})", entry.ident().to_string().await?)
220 }
221 ChunkGroup::Isolated(entry) => {
222 format!(
223 "ChunkGroup::Isolated({:?})",
224 entry.ident().to_string().await?
225 )
226 }
227 ChunkGroup::Shared(entry) => {
228 format!("ChunkGroup::Shared({:?})", entry.ident().to_string().await?)
229 }
230 ChunkGroup::IsolatedMerged {
231 parent,
232 merge_tag,
233 entries,
234 } => {
235 format!(
236 "ChunkGroup::IsolatedMerged({}, {}, {:?})",
237 Box::pin(chunk_group_info.chunk_groups[*parent].debug_str(chunk_group_info))
238 .await?,
239 merge_tag,
240 entries
241 .iter()
242 .map(|m| m.ident().to_string())
243 .try_join()
244 .await?
245 )
246 }
247 ChunkGroup::SharedMerged {
248 parent,
249 merge_tag,
250 entries,
251 } => {
252 format!(
253 "ChunkGroup::SharedMerged({}, {}, {:?})",
254 Box::pin(chunk_group_info.chunk_groups[*parent].debug_str(chunk_group_info))
255 .await?,
256 merge_tag,
257 entries
258 .iter()
259 .map(|m| m.ident().to_string())
260 .try_join()
261 .await?
262 )
263 }
264 })
265 }
266}
267
268#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
269pub enum ChunkGroupKey {
270 Entry(Vec<ResolvedVc<Box<dyn Module>>>),
272 Async(ResolvedVc<Box<dyn Module>>),
274 Isolated(ResolvedVc<Box<dyn Module>>),
276 IsolatedMerged {
278 parent: ChunkGroupId,
279 merge_tag: RcStr,
280 },
281 Shared(ResolvedVc<Box<dyn Module>>),
283 SharedMerged {
285 parent: ChunkGroupId,
286 merge_tag: RcStr,
287 },
288}
289
290#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
291pub struct ChunkGroupId(u32);
292
293impl From<usize> for ChunkGroupId {
294 fn from(id: usize) -> Self {
295 Self(id as u32)
296 }
297}
298
299impl Deref for ChunkGroupId {
300 type Target = u32;
301 fn deref(&self) -> &Self::Target {
302 &self.0
303 }
304}
305
306#[derive(Debug, Clone, PartialEq, Eq)]
307struct TraversalPriority {
308 depth: usize,
309 chunk_group_len: u64,
310}
311impl PartialOrd for TraversalPriority {
312 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
313 Some(self.cmp(other))
314 }
315}
316impl Ord for TraversalPriority {
317 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
318 let depth_order = self.depth.cmp(&other.depth).reverse();
322 let chunk_group_len_order = self.chunk_group_len.cmp(&other.chunk_group_len).reverse();
324
325 depth_order.then(chunk_group_len_order)
326 }
327}
328
329pub async fn compute_chunk_group_info(graph: &ModuleGraph) -> Result<Vc<ChunkGroupInfo>> {
330 let span_outer = tracing::info_span!(
331 "compute chunk group info",
332 module_count = tracing::field::Empty,
333 visit_count = tracing::field::Empty,
334 chunk_group_count = tracing::field::Empty
335 );
336
337 let span = span_outer.clone();
338 async move {
339 #[allow(clippy::type_complexity)]
340 let mut chunk_groups_map: FxIndexMap<
341 ChunkGroupKey,
342 (ChunkGroupId, FxIndexSet<ResolvedVc<Box<dyn Module>>>),
343 > = FxIndexMap::default();
344
345 let mut module_chunk_groups: FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper> =
348 FxHashMap::default();
349
350 let graphs = graph.graphs.iter().try_join().await?;
351 let module_count = graphs.iter().map(|g| g.graph.node_count()).sum::<usize>();
352 span.record("module_count", module_count);
353
354 let entries = graphs
356 .iter()
357 .flat_map(|g| g.entries.iter())
358 .collect::<Vec<_>>();
359
360 let module_depth: FxHashMap<ResolvedVc<Box<dyn Module>>, usize> = {
362 let mut module_depth = FxHashMap::default();
363 graph
364 .traverse_edges_from_entries_bfs(
365 entries.iter().flat_map(|e| e.entries()),
366 |parent, node| {
367 if let Some((parent, _)) = parent {
368 let parent_depth = *module_depth
369 .get(&parent.module)
370 .context("Module depth not found")?;
371 module_depth.entry(node.module).or_insert(parent_depth + 1);
372 } else {
373 module_depth.insert(node.module, 0);
374 };
375
376 module_chunk_groups.insert(node.module, RoaringBitmapWrapper::default());
377
378 Ok(GraphTraversalAction::Continue)
379 },
380 )
381 .await?;
382 module_depth
383 };
384
385 #[allow(clippy::type_complexity)]
388 fn entry_to_chunk_group_id(
389 entry: ChunkGroupEntry,
390 chunk_groups_map: &mut FxIndexMap<
391 ChunkGroupKey,
392 (ChunkGroupId, FxIndexSet<ResolvedVc<Box<dyn Module>>>),
393 >,
394 ) -> ChunkGroupKey {
395 match entry {
396 ChunkGroupEntry::Entry(entries) => ChunkGroupKey::Entry(entries),
397 ChunkGroupEntry::Async(entry) => ChunkGroupKey::Async(entry),
398 ChunkGroupEntry::Isolated(entry) => ChunkGroupKey::Isolated(entry),
399 ChunkGroupEntry::Shared(entry) => ChunkGroupKey::Shared(entry),
400 ChunkGroupEntry::IsolatedMerged {
401 parent,
402 merge_tag,
403 entries: _,
404 } => {
405 let parent = entry_to_chunk_group_id(*parent, chunk_groups_map);
406 let len = chunk_groups_map.len();
407 let parent = chunk_groups_map
408 .entry(parent)
409 .or_insert_with(|| (ChunkGroupId(len as u32), FxIndexSet::default()))
410 .0;
411
412 ChunkGroupKey::IsolatedMerged {
413 parent: ChunkGroupId(*parent as u32),
414 merge_tag,
415 }
416 }
417 ChunkGroupEntry::SharedMerged {
418 parent,
419 merge_tag,
420 entries: _,
421 } => {
422 let parent = entry_to_chunk_group_id(*parent, chunk_groups_map);
423 let len = chunk_groups_map.len();
424 let parent = chunk_groups_map
425 .entry(parent)
426 .or_insert_with(|| (ChunkGroupId(len as u32), FxIndexSet::default()))
427 .0;
428
429 ChunkGroupKey::SharedMerged {
430 parent: ChunkGroupId(*parent as u32),
431 merge_tag,
432 }
433 }
434 }
435 }
436
437 let entry_chunk_group_keys = graphs
438 .iter()
439 .flat_map(|g| g.entries.iter())
440 .flat_map(|chunk_group| {
441 let chunk_group_key =
442 entry_to_chunk_group_id(chunk_group.clone(), &mut chunk_groups_map);
443 chunk_group
444 .entries()
445 .map(move |e| (e, chunk_group_key.clone()))
446 })
447 .collect::<FxHashMap<_, _>>();
448
449 let visit_count = graph
450 .traverse_edges_fixed_point_with_priority(
451 entries
452 .iter()
453 .flat_map(|e| e.entries())
454 .map(|e| {
455 Ok((
456 e,
457 TraversalPriority {
458 depth: *module_depth.get(&e).context("Module depth not found")?,
459 chunk_group_len: 0,
460 },
461 ))
462 })
463 .collect::<Result<Vec<_>>>()?,
464 &mut module_chunk_groups,
465 |parent_info: Option<(&'_ SingleModuleGraphModuleNode, &'_ RefData)>,
466 node: &'_ SingleModuleGraphModuleNode,
467 module_chunk_groups: &mut FxHashMap<
468 ResolvedVc<Box<dyn Module>>,
469 RoaringBitmapWrapper,
470 >|
471 -> Result<GraphTraversalAction> {
472 enum ChunkGroupInheritance<It: Iterator<Item = ChunkGroupKey>> {
473 Inherit(ResolvedVc<Box<dyn Module>>),
474 ChunkGroup(It),
475 }
476 let chunk_groups = if let Some((parent, ref_data)) = parent_info {
477 match &ref_data.chunking_type {
478 ChunkingType::Parallel { .. } => {
479 ChunkGroupInheritance::Inherit(parent.module)
480 }
481 ChunkingType::Async => ChunkGroupInheritance::ChunkGroup(Either::Left(
482 std::iter::once(ChunkGroupKey::Async(node.module)),
483 )),
484 ChunkingType::Isolated {
485 merge_tag: None, ..
486 } => ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
487 ChunkGroupKey::Isolated(node.module),
488 ))),
489 ChunkingType::Shared {
490 merge_tag: None, ..
491 } => ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
492 ChunkGroupKey::Shared(node.module),
493 ))),
494 ChunkingType::Isolated {
495 merge_tag: Some(merge_tag),
496 ..
497 } => {
498 let parents = module_chunk_groups
499 .get(&parent.module)
500 .context("Module chunk group not found")?;
501 let chunk_groups =
502 parents.iter().map(|parent| ChunkGroupKey::IsolatedMerged {
503 parent: ChunkGroupId(parent),
504 merge_tag: merge_tag.clone(),
505 });
506 ChunkGroupInheritance::ChunkGroup(Either::Right(Either::Left(
507 chunk_groups,
508 )))
509 }
510 ChunkingType::Shared {
511 merge_tag: Some(merge_tag),
512 ..
513 } => {
514 let parents = module_chunk_groups
515 .get(&parent.module)
516 .context("Module chunk group not found")?;
517 let chunk_groups =
518 parents.iter().map(|parent| ChunkGroupKey::SharedMerged {
519 parent: ChunkGroupId(parent),
520 merge_tag: merge_tag.clone(),
521 });
522 ChunkGroupInheritance::ChunkGroup(Either::Right(Either::Right(
523 chunk_groups,
524 )))
525 }
526 ChunkingType::Traced => {
527 return Ok(GraphTraversalAction::Skip);
529 }
530 }
531 } else {
532 ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
533 entry_chunk_group_keys
535 .get(&node.module)
536 .context("Module chunk group not found")?
537 .clone(),
538 )))
539 };
540
541 Ok(match chunk_groups {
542 ChunkGroupInheritance::ChunkGroup(chunk_groups) => {
543 let chunk_group_ids = chunk_groups.map(|chunk_group| {
545 let len = chunk_groups_map.len();
546 let is_merged = matches!(
547 chunk_group,
548 ChunkGroupKey::IsolatedMerged { .. }
549 | ChunkGroupKey::SharedMerged { .. }
550 );
551 match chunk_groups_map.entry(chunk_group) {
552 Entry::Occupied(mut e) => {
553 let (id, merged_entries) = e.get_mut();
554 if is_merged {
555 merged_entries.insert(node.module);
556 }
557 **id
558 }
559 Entry::Vacant(e) => {
560 let chunk_group_id = len as u32;
561 let mut set = FxIndexSet::default();
562 if is_merged {
563 set.insert(node.module);
564 }
565 e.insert((ChunkGroupId(chunk_group_id), set));
566 chunk_group_id
567 }
568 }
569 });
570
571 let chunk_groups =
572 RoaringBitmapWrapper(RoaringBitmap::from_iter(chunk_group_ids));
573
574 let bitset = module_chunk_groups
576 .get_mut(&node.module)
577 .context("Module chunk group not found")?;
578 if chunk_groups.is_proper_superset(bitset) {
579 **bitset |= chunk_groups.into_inner();
581
582 GraphTraversalAction::Continue
583 } else {
584 GraphTraversalAction::Skip
586 }
587 }
588 ChunkGroupInheritance::Inherit(parent) => {
589 if parent == node.module {
593 GraphTraversalAction::Skip
595 } else {
596 let [Some(parent_chunk_groups), Some(current_chunk_groups)] =
597 module_chunk_groups.get_disjoint_mut([&parent, &node.module])
598 else {
599 bail!("Module chunk groups not found");
603 };
604
605 if current_chunk_groups.is_empty() {
606 *current_chunk_groups = parent_chunk_groups.clone();
608 GraphTraversalAction::Continue
609 } else if parent_chunk_groups
610 .is_proper_superset(current_chunk_groups)
611 {
612 **current_chunk_groups |= &**parent_chunk_groups;
614 GraphTraversalAction::Continue
615 } else {
616 GraphTraversalAction::Skip
618 }
619 }
620 }
621 })
622 },
623 |successor, module_chunk_groups| {
631 Ok(TraversalPriority {
632 depth: *module_depth
633 .get(&successor.module)
634 .context("Module depth not found")?,
635 chunk_group_len: module_chunk_groups
636 .get(&successor.module)
637 .context("Module chunk group not found")?
638 .len(),
639 })
640 },
641 )
642 .await?;
643
644 span.record("visit_count", visit_count);
645 span.record("chunk_group_count", chunk_groups_map.len());
646
647 Ok(ChunkGroupInfo {
648 module_chunk_groups,
649 chunk_group_keys: chunk_groups_map.keys().cloned().collect(),
650 chunk_groups: chunk_groups_map
651 .into_iter()
652 .map(|(k, (_, merged_entries))| match k {
653 ChunkGroupKey::Entry(entries) => ChunkGroup::Entry(entries),
654 ChunkGroupKey::Async(module) => ChunkGroup::Async(module),
655 ChunkGroupKey::Isolated(module) => ChunkGroup::Isolated(module),
656 ChunkGroupKey::IsolatedMerged { parent, merge_tag } => {
657 ChunkGroup::IsolatedMerged {
658 parent: parent.0 as usize,
659 merge_tag,
660 entries: merged_entries.into_iter().collect(),
661 }
662 }
663 ChunkGroupKey::Shared(module) => ChunkGroup::Shared(module),
664 ChunkGroupKey::SharedMerged { parent, merge_tag } => ChunkGroup::SharedMerged {
665 parent: parent.0 as usize,
666 merge_tag,
667 entries: merged_entries.into_iter().collect(),
668 },
669 })
670 .collect(),
671 }
672 .cell())
673 }
674 .instrument(span_outer)
675 .await
676}