1use std::{
2 hash::Hash,
3 ops::{Deref, DerefMut},
4};
5
6use anyhow::{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, 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.get(&parent.module).unwrap();
369 module_depth.entry(node.module).or_insert(parent_depth + 1);
370 } else {
371 module_depth.insert(node.module, 0);
372 };
373
374 module_chunk_groups.insert(node.module, RoaringBitmapWrapper::default());
375
376 GraphTraversalAction::Continue
377 },
378 )
379 .await?;
380 module_depth
381 };
382
383 #[allow(clippy::type_complexity)]
386 fn entry_to_chunk_group_id(
387 entry: ChunkGroupEntry,
388 chunk_groups_map: &mut FxIndexMap<
389 ChunkGroupKey,
390 (ChunkGroupId, FxIndexSet<ResolvedVc<Box<dyn Module>>>),
391 >,
392 ) -> ChunkGroupKey {
393 match entry {
394 ChunkGroupEntry::Entry(entries) => ChunkGroupKey::Entry(entries),
395 ChunkGroupEntry::Async(entry) => ChunkGroupKey::Async(entry),
396 ChunkGroupEntry::Isolated(entry) => ChunkGroupKey::Isolated(entry),
397 ChunkGroupEntry::Shared(entry) => ChunkGroupKey::Shared(entry),
398 ChunkGroupEntry::IsolatedMerged {
399 parent,
400 merge_tag,
401 entries: _,
402 } => {
403 let parent = entry_to_chunk_group_id(*parent, chunk_groups_map);
404 let len = chunk_groups_map.len();
405 let parent = chunk_groups_map
406 .entry(parent)
407 .or_insert_with(|| (ChunkGroupId(len as u32), FxIndexSet::default()))
408 .0;
409
410 ChunkGroupKey::IsolatedMerged {
411 parent: ChunkGroupId(*parent as u32),
412 merge_tag,
413 }
414 }
415 ChunkGroupEntry::SharedMerged {
416 parent,
417 merge_tag,
418 entries: _,
419 } => {
420 let parent = entry_to_chunk_group_id(*parent, chunk_groups_map);
421 let len = chunk_groups_map.len();
422 let parent = chunk_groups_map
423 .entry(parent)
424 .or_insert_with(|| (ChunkGroupId(len as u32), FxIndexSet::default()))
425 .0;
426
427 ChunkGroupKey::SharedMerged {
428 parent: ChunkGroupId(*parent as u32),
429 merge_tag,
430 }
431 }
432 }
433 }
434
435 let entry_chunk_group_keys = graphs
436 .iter()
437 .flat_map(|g| g.entries.iter())
438 .flat_map(|chunk_group| {
439 let chunk_group_key =
440 entry_to_chunk_group_id(chunk_group.clone(), &mut chunk_groups_map);
441 chunk_group
442 .entries()
443 .map(move |e| (e, chunk_group_key.clone()))
444 })
445 .collect::<FxHashMap<_, _>>();
446
447 let visit_count = graph
448 .traverse_edges_fixed_point_with_priority(
449 entries.iter().flat_map(|e| e.entries()).map(|e| {
450 (
451 e,
452 TraversalPriority {
453 depth: *module_depth.get(&e).unwrap(),
454 chunk_group_len: 0,
455 },
456 )
457 }),
458 &mut module_chunk_groups,
459 |parent_info: Option<(&'_ SingleModuleGraphModuleNode, &'_ ChunkingType)>,
460 node: &'_ SingleModuleGraphModuleNode,
461 module_chunk_groups: &mut FxHashMap<
462 ResolvedVc<Box<dyn Module>>,
463 RoaringBitmapWrapper,
464 >|
465 -> GraphTraversalAction {
466 enum ChunkGroupInheritance<It: Iterator<Item = ChunkGroupKey>> {
467 Inherit(ResolvedVc<Box<dyn Module>>),
468 ChunkGroup(It),
469 }
470 let chunk_groups = if let Some((parent, chunking_type)) = parent_info {
471 match chunking_type {
472 ChunkingType::Parallel { .. } => {
473 ChunkGroupInheritance::Inherit(parent.module)
474 }
475 ChunkingType::Async => ChunkGroupInheritance::ChunkGroup(Either::Left(
476 std::iter::once(ChunkGroupKey::Async(node.module)),
477 )),
478 ChunkingType::Isolated {
479 merge_tag: None, ..
480 } => ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
481 ChunkGroupKey::Isolated(node.module),
482 ))),
483 ChunkingType::Shared {
484 merge_tag: None, ..
485 } => ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
486 ChunkGroupKey::Shared(node.module),
487 ))),
488 ChunkingType::Isolated {
489 merge_tag: Some(merge_tag),
490 ..
491 } => {
492 let parents = module_chunk_groups.get(&parent.module).unwrap();
493 let chunk_groups =
494 parents.iter().map(|parent| ChunkGroupKey::IsolatedMerged {
495 parent: ChunkGroupId(parent),
496 merge_tag: merge_tag.clone(),
497 });
498 ChunkGroupInheritance::ChunkGroup(Either::Right(Either::Left(
499 chunk_groups,
500 )))
501 }
502 ChunkingType::Shared {
503 merge_tag: Some(merge_tag),
504 ..
505 } => {
506 let parents = module_chunk_groups.get(&parent.module).unwrap();
507 let chunk_groups =
508 parents.iter().map(|parent| ChunkGroupKey::SharedMerged {
509 parent: ChunkGroupId(parent),
510 merge_tag: merge_tag.clone(),
511 });
512 ChunkGroupInheritance::ChunkGroup(Either::Right(Either::Right(
513 chunk_groups,
514 )))
515 }
516 ChunkingType::Traced => {
517 return GraphTraversalAction::Skip;
519 }
520 }
521 } else {
522 ChunkGroupInheritance::ChunkGroup(Either::Left(std::iter::once(
523 entry_chunk_group_keys.get(&node.module).unwrap().clone(),
524 )))
525 };
526
527 match chunk_groups {
528 ChunkGroupInheritance::ChunkGroup(chunk_groups) => {
529 let chunk_group_ids = chunk_groups.map(|chunk_group| {
531 let len = chunk_groups_map.len();
532 let is_merged = matches!(
533 chunk_group,
534 ChunkGroupKey::IsolatedMerged { .. }
535 | ChunkGroupKey::SharedMerged { .. }
536 );
537 match chunk_groups_map.entry(chunk_group) {
538 Entry::Occupied(mut e) => {
539 let (id, merged_entries) = e.get_mut();
540 if is_merged {
541 merged_entries.insert(node.module);
542 }
543 **id
544 }
545 Entry::Vacant(e) => {
546 let chunk_group_id = len as u32;
547 let mut set = FxIndexSet::default();
548 if is_merged {
549 set.insert(node.module);
550 }
551 e.insert((ChunkGroupId(chunk_group_id), set));
552 chunk_group_id
553 }
554 }
555 });
556
557 let chunk_groups =
558 RoaringBitmapWrapper(RoaringBitmap::from_iter(chunk_group_ids));
559
560 let bitset = module_chunk_groups.get_mut(&node.module).unwrap();
562 if chunk_groups.is_proper_superset(bitset) {
563 **bitset |= chunk_groups.into_inner();
565
566 GraphTraversalAction::Continue
567 } else {
568 GraphTraversalAction::Skip
570 }
571 }
572 ChunkGroupInheritance::Inherit(parent) => {
573 if parent == node.module {
577 GraphTraversalAction::Skip
579 } else {
580 let [Some(parent_chunk_groups), Some(current_chunk_groups)] =
582 module_chunk_groups.get_disjoint_mut([&parent, &node.module])
583 else {
584 unreachable!()
586 };
587
588 if current_chunk_groups.is_empty() {
589 *current_chunk_groups = parent_chunk_groups.clone();
591 GraphTraversalAction::Continue
592 } else if parent_chunk_groups
593 .is_proper_superset(current_chunk_groups)
594 {
595 **current_chunk_groups |= &**parent_chunk_groups;
597 GraphTraversalAction::Continue
598 } else {
599 GraphTraversalAction::Skip
601 }
602 }
603 }
604 }
605 },
606 |successor, module_chunk_groups| TraversalPriority {
614 depth: *module_depth.get(&successor.module).unwrap(),
615 chunk_group_len: module_chunk_groups.get(&successor.module).unwrap().len(),
616 },
617 )
618 .await?;
619
620 span.record("visit_count", visit_count);
621 span.record("chunk_group_count", chunk_groups_map.len());
622
623 Ok(ChunkGroupInfo {
624 module_chunk_groups,
625 chunk_group_keys: chunk_groups_map.keys().cloned().collect(),
626 chunk_groups: chunk_groups_map
627 .into_iter()
628 .map(|(k, (_, merged_entries))| match k {
629 ChunkGroupKey::Entry(entries) => ChunkGroup::Entry(entries),
630 ChunkGroupKey::Async(module) => ChunkGroup::Async(module),
631 ChunkGroupKey::Isolated(module) => ChunkGroup::Isolated(module),
632 ChunkGroupKey::IsolatedMerged { parent, merge_tag } => {
633 ChunkGroup::IsolatedMerged {
634 parent: parent.0 as usize,
635 merge_tag,
636 entries: merged_entries.into_iter().collect(),
637 }
638 }
639 ChunkGroupKey::Shared(module) => ChunkGroup::Shared(module),
640 ChunkGroupKey::SharedMerged { parent, merge_tag } => ChunkGroup::SharedMerged {
641 parent: parent.0 as usize,
642 merge_tag,
643 entries: merged_entries.into_iter().collect(),
644 },
645 })
646 .collect(),
647 }
648 .cell())
649 }
650 .instrument(span_outer)
651 .await
652}