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