1use std::collections::hash_map::Entry;
2
3use anyhow::{Context, Result, bail};
4use rustc_hash::{FxHashMap, FxHashSet};
5use tracing::Instrument;
6use turbo_tasks::{FxIndexMap, FxIndexSet, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, Vc};
7
8use crate::{
9 chunk::{
10 ChunkableModule, ChunkingType, MergeableModule, MergeableModuleExposure, MergeableModules,
11 MergeableModulesExposed,
12 },
13 module::Module,
14 module_graph::{
15 GraphTraversalAction, ModuleGraph, RefData, chunk_group_info::RoaringBitmapWrapper,
16 },
17 resolve::ExportUsage,
18};
19
20#[turbo_tasks::value(transparent, cell = "keyed")]
21#[allow(clippy::type_complexity)]
22pub struct MergedModulesReplacements(
23 FxHashMap<ResolvedVc<Box<dyn Module>>, ResolvedVc<Box<dyn ChunkableModule>>>,
24);
25
26#[turbo_tasks::value(transparent, cell = "keyed")]
27#[allow(clippy::type_complexity)]
28pub struct MergedModulesOriginalModules(
29 FxHashMap<ResolvedVc<Box<dyn Module>>, ResolvedVc<Box<dyn Module>>>,
30);
31
32#[turbo_tasks::value(transparent, cell = "keyed")]
33#[allow(clippy::type_complexity)]
34pub struct MergedModulesIncluded(FxHashSet<ResolvedVc<Box<dyn Module>>>);
35
36#[turbo_tasks::value]
37pub struct MergedModuleInfo {
38 pub replacements: ResolvedVc<MergedModulesReplacements>,
40 pub replacements_to_original: ResolvedVc<MergedModulesOriginalModules>,
43 pub included: ResolvedVc<MergedModulesIncluded>,
45}
46
47impl MergedModuleInfo {
48 pub async fn should_replace_module(
50 &self,
51 module: ResolvedVc<Box<dyn Module>>,
52 ) -> Result<Option<ResolvedVc<Box<dyn ChunkableModule>>>> {
53 Ok(self.replacements.get(&module).await?.as_deref().copied())
54 }
55
56 pub async fn get_original_module(
59 &self,
60 module: ResolvedVc<Box<dyn Module>>,
61 ) -> Result<Option<ResolvedVc<Box<dyn Module>>>> {
62 Ok(self
63 .replacements_to_original
64 .get(&module)
65 .await?
66 .as_deref()
67 .copied())
68 }
69
70 pub async fn should_create_chunk_item_for(
73 &self,
74 module: ResolvedVc<Box<dyn Module>>,
75 ) -> Result<bool> {
76 Ok(!self.included.contains_key(&module).await?)
77 }
78}
79
80pub async fn compute_merged_modules(module_graph: Vc<ModuleGraph>) -> Result<Vc<MergedModuleInfo>> {
86 let span_outer = tracing::info_span!(
87 "compute merged modules",
88 module_count = tracing::field::Empty,
89 visit_count = tracing::field::Empty,
90 merged_groups = tracing::field::Empty,
91 included_modules = tracing::field::Empty
92 );
93
94 let span = span_outer.clone();
95 async move {
96 let async_module_info = module_graph.async_module_info();
97 let chunk_group_info = module_graph.chunk_group_info().await?;
98 let module_graph = module_graph.await?;
99
100 let graphs = &module_graph.graphs;
101 let module_count = graphs.iter().map(|g| g.graph.node_count()).sum::<usize>();
102 span.record("module_count", module_count);
103
104 let entries = graphs
106 .iter()
107 .flat_map(|g| g.entries.iter())
108 .flat_map(|g| g.entries())
109 .collect::<Vec<_>>();
110
111 let module_depth = {
113 let _inner_span = tracing::info_span!("compute depth").entered();
114
115 let mut module_depth =
116 FxHashMap::with_capacity_and_hasher(module_count, Default::default());
117 module_graph.traverse_edges_bfs(entries.iter().copied(), |parent, node| {
118 if let Some((parent, _)) = parent {
119 let parent_depth = *module_depth
120 .get(&parent)
121 .context("Module depth not found")?;
122 module_depth.entry(node).or_insert(parent_depth + 1);
123 } else {
124 module_depth.insert(node, 0);
125 };
126
127 Ok(GraphTraversalAction::Continue)
128 })?;
129 module_depth
130 };
131
132 let mut module_merged_groups: FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper> =
136 FxHashMap::with_capacity_and_hasher(module_count, Default::default());
137 let mut entry_modules =
139 FxHashSet::with_capacity_and_hasher(module_count, Default::default());
140
141 let inner_span = tracing::info_span!("collect mergeable modules");
142 let mergeable = graphs
143 .iter()
144 .flat_map(|g| g.iter_nodes())
145 .map(async |module| {
146 if let Some(mergeable) =
147 ResolvedVc::try_downcast::<Box<dyn MergeableModule>>(module)
148 && *mergeable.is_mergeable().await?
149 {
150 return Ok(Some(module));
151 }
152 Ok(None)
153 })
154 .try_flat_join()
155 .instrument(inner_span)
156 .await?
157 .into_iter()
158 .collect::<FxHashSet<_>>();
159
160 let inner_span = tracing::info_span!("pre-fetch async module status");
163 let async_modules: FxHashSet<_> = mergeable
164 .iter()
165 .map(async |&module| Ok(async_module_info.is_async(module).await?.then_some(module)))
166 .try_flat_join()
167 .instrument(inner_span)
168 .await?
169 .into_iter()
170 .collect();
171
172 let inner_span = tracing::info_span!("fixed point traversal").entered();
173
174 let mut next_index = 0u32;
175 let visit_count = module_graph.traverse_edges_fixed_point_with_priority(
176 entries
177 .iter()
178 .map(|e| Ok((*e, -*module_depth.get(e).context("Module depth not found")?)))
179 .collect::<Result<Vec<_>>>()?,
180 &mut (),
181 |parent_info: Option<(ResolvedVc<Box<dyn Module>>, &'_ RefData, _)>,
182 node: ResolvedVc<Box<dyn Module>>,
183 _|
184 -> Result<GraphTraversalAction> {
185 let (parent_module, hoisted) =
188 parent_info.map_or((None, false), |(node, ty, _)| {
189 (
190 Some(node),
191 match &ty.chunking_type {
192 ChunkingType::Parallel { hoisted, .. } => *hoisted,
193 _ => false,
194 },
195 )
196 });
197 let module = node;
198
199 Ok(if parent_module.is_some_and(|p| p == module) {
200 GraphTraversalAction::Skip
202 } else if hoisted
203 && let Some(parent_module) = parent_module
204 && mergeable.contains(&parent_module)
205 && mergeable.contains(&module)
206 && !async_modules.contains(&parent_module)
207 && !async_modules.contains(&module)
208 {
209 module_merged_groups.entry(node).or_default();
214 let [Some(parent_merged_groups), Some(current_merged_groups)] =
215 module_merged_groups.get_disjoint_mut([&parent_module, &node])
216 else {
217 bail!("unreachable except for eventual consistency");
219 };
220
221 if current_merged_groups.is_empty() {
222 *current_merged_groups = parent_merged_groups.clone();
224 GraphTraversalAction::Continue
225 } else if parent_merged_groups.is_proper_superset(current_merged_groups) {
226 **current_merged_groups |= &**parent_merged_groups;
228 GraphTraversalAction::Continue
229 } else {
230 GraphTraversalAction::Skip
232 }
233 } else {
234 if entry_modules.insert(module) {
237 let idx = next_index;
239 next_index += 1;
240
241 if module_merged_groups.entry(module).or_default().insert(idx) {
242 GraphTraversalAction::Continue
244 } else {
245 GraphTraversalAction::Skip
247 }
248 } else {
249 GraphTraversalAction::Skip
252 }
253 })
254 },
255 |successor, _| {
256 Ok(-*module_depth
259 .get(&successor)
260 .context("Module depth not found")?)
261 },
262 )?;
263
264 drop(inner_span);
265 let inner_span = tracing::info_span!("chunk group collection").entered();
266
267 span.record("visit_count", visit_count);
268
269 #[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
270 struct ListOccurrence {
271 entry: usize,
274 list: usize,
275 chunk_group: usize,
276 }
277
278 let mut lists;
283 let mut lists_reverse_indices: FxIndexMap<
284 ResolvedVc<Box<dyn MergeableModule>>,
285 FxIndexSet<ListOccurrence>,
286 > = FxIndexMap::default();
287
288 #[allow(non_snake_case)]
291 let LISTS_COMMON_IDX: usize;
292
293 #[allow(clippy::type_complexity)]
297 let mut intra_group_references: FxIndexMap<
298 ResolvedVc<Box<dyn Module>>,
299 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
300 > = FxIndexMap::default();
301 #[allow(clippy::type_complexity)]
304 let mut intra_group_references_rev: FxIndexMap<
305 ResolvedVc<Box<dyn Module>>,
306 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
307 > = FxIndexMap::default();
308
309 {
310 struct ChunkGroupResult {
311 first_chunk_group_idx: usize,
312 #[allow(clippy::type_complexity)]
313 list_lists: Vec<Vec<Vec<ResolvedVc<Box<dyn MergeableModule>>>>>,
314 lists_reverse_indices:
315 FxIndexMap<ResolvedVc<Box<dyn MergeableModule>>, FxIndexSet<ListOccurrence>>,
316 #[allow(clippy::type_complexity)]
317 intra_group_references_rev: FxIndexMap<
318 ResolvedVc<Box<dyn Module>>,
319 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
320 >,
321 }
322 let span = tracing::info_span!("map chunk groups").entered();
323
324 let result = turbo_tasks::parallel::map_collect_chunked_owned::<_, _, Result<Vec<_>>>(
325 chunk_group_info.chunk_groups.iter().enumerate().collect(),
327 |chunk| {
328 let mut list_lists = vec![];
329 let mut lists_reverse_indices: FxIndexMap<
330 ResolvedVc<Box<dyn MergeableModule>>,
331 FxIndexSet<ListOccurrence>,
332 > = FxIndexMap::default();
333 #[allow(clippy::type_complexity)]
334 let mut intra_group_references_rev: FxIndexMap<
335 ResolvedVc<Box<dyn Module>>,
336 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
337 > = FxIndexMap::default();
338
339 let mut chunk = chunk.peekable();
340 let first_chunk_group_idx = chunk.peek().unwrap().0;
341
342 for (chunk_group_idx, chunk_group) in chunk {
343 let mut lists = vec![];
344
345 let mut chunk_lists: FxHashMap<&RoaringBitmapWrapper, usize> =
349 FxHashMap::with_capacity_and_hasher(
350 module_merged_groups.len() / chunk_group_info.chunk_groups.len(),
351 Default::default(),
352 );
353
354 let mut visited = FxHashSet::default();
358
359 module_graph.traverse_edges_dfs(
360 chunk_group.entries(),
361 &mut (),
362 |parent_info, node, _| {
363 if parent_info.is_none_or(|(_, r)| r.chunking_type.is_parallel())
364 && visited.insert(node)
365 {
366 Ok(GraphTraversalAction::Continue)
367 } else {
368 Ok(GraphTraversalAction::Exclude)
369 }
370 },
371 |parent_info, node, _| {
372 let module = node;
373 let bitmap = module_merged_groups
374 .get(&module)
375 .context("every module should have a bitmap")?;
376
377 if mergeable.contains(&module) {
378 let mergeable_module =
379 ResolvedVc::try_downcast::<Box<dyn MergeableModule>>(
380 module,
381 )
382 .context(
383 "found mergeable module which is not a MergeableModule",
384 )?;
385 match chunk_lists.entry(bitmap) {
386 Entry::Vacant(e) => {
387 let idx = lists.len();
389 e.insert(idx);
390 lists.push(vec![mergeable_module]);
391 lists_reverse_indices
392 .entry(mergeable_module)
393 .or_default()
394 .insert(ListOccurrence {
395 chunk_group: chunk_group_idx,
396 list: idx,
397 entry: 0,
398 });
399 }
400 Entry::Occupied(e) => {
401 let list_idx = *e.get();
402 let list = &mut lists[list_idx];
403 list.push(mergeable_module);
404 lists_reverse_indices
405 .entry(mergeable_module)
406 .or_default()
407 .insert(ListOccurrence {
408 chunk_group: chunk_group_idx,
409 list: list_idx,
410 entry: list.len() - 1,
411 });
412 }
413 }
414 }
415
416 if let Some((parent, _)) = parent_info {
417 let same_bitmap = module_merged_groups
418 .get(&parent)
419 .context("every module should have a bitmap")?
420 == module_merged_groups
421 .get(&module)
422 .context("every module should have a bitmap")?;
423
424 if same_bitmap {
425 intra_group_references_rev
426 .entry(module)
427 .or_default()
428 .insert(parent);
429 }
430 }
431 Ok(())
432 },
433 )?;
434
435 list_lists.push(lists);
436 }
437 Ok(ChunkGroupResult {
438 first_chunk_group_idx,
439 list_lists,
440 lists_reverse_indices,
441 intra_group_references_rev,
442 })
443 },
444 )?;
445
446 drop(span);
447 let _span = tracing::info_span!("merging chunk group lists").entered();
448
449 lists_reverse_indices
450 .reserve_exact(result.iter().map(|r| r.lists_reverse_indices.len()).sum());
451 intra_group_references_rev.reserve_exact(
452 result
453 .iter()
454 .map(|r| r.intra_group_references_rev.len())
455 .sum(),
456 );
457
458 lists = vec![Default::default(); chunk_group_info.chunk_groups.len() + 1];
459 LISTS_COMMON_IDX = result.len();
460 for ChunkGroupResult {
461 first_chunk_group_idx,
462 list_lists: result_lists,
463 lists_reverse_indices: result_lists_reverse_indices,
464 intra_group_references_rev: result_intra_group_references_rev,
465 } in result
466 {
467 lists.splice(
468 first_chunk_group_idx..(first_chunk_group_idx + result_lists.len()),
469 result_lists,
470 );
471 for (module, occurrences) in result_lists_reverse_indices {
472 lists_reverse_indices
473 .entry(module)
474 .or_default()
475 .extend(occurrences);
476 }
477 for (module, occurrences) in result_intra_group_references_rev {
478 intra_group_references_rev
479 .entry(module)
480 .or_default()
481 .extend(occurrences);
482 }
483 }
484 }
485
486 drop(inner_span);
487 let inner_span = tracing::info_span!("exposed computation").entered();
488
489 lists_reverse_indices
491 .sort_by_cached_key(|_, b| b.iter().map(|o| o.entry).min().map(|v| -(v as i64)));
492
493 let mut exposed_modules_imported: FxHashSet<ResolvedVc<Box<dyn Module>>> =
497 FxHashSet::with_capacity_and_hasher(module_merged_groups.len(), Default::default());
498 let mut exposed_modules_namespace: FxHashSet<ResolvedVc<Box<dyn Module>>> =
499 FxHashSet::with_capacity_and_hasher(module_merged_groups.len(), Default::default());
500
501 module_graph.traverse_edges_dfs(
502 entries,
503 &mut (),
504 |_, _, _| Ok(GraphTraversalAction::Continue),
505 |parent_info, node, _| {
506 let module = node;
507
508 if let Some((parent, _)) = parent_info {
509 let same_bitmap = module_merged_groups
510 .get(&parent)
511 .context("every module should have a bitmap")?
512 == module_merged_groups
513 .get(&module)
514 .context("every module should have a bitmap")?;
515
516 if same_bitmap {
517 intra_group_references
518 .entry(parent)
519 .or_default()
520 .insert(module);
521 }
522 }
523
524 if match parent_info {
525 None => true,
526 Some((parent, _)) => {
527 module_merged_groups
528 .get(&parent)
529 .context("every module should have a bitmap")?
530 != module_merged_groups
531 .get(&module)
532 .context("every module should have a bitmap")?
533 }
534 } {
535 exposed_modules_imported.insert(module);
540 }
541 if parent_info.is_some_and(|(_, r)| {
542 matches!(
543 r.binding_usage.export,
544 ExportUsage::All | ExportUsage::PartialNamespaceObject(_)
545 )
546 }) {
547 exposed_modules_namespace.insert(module);
550 }
551 Ok(())
552 },
553 )?;
554
555 drop(inner_span);
556 let inner_span = tracing::info_span!("reconciliation").entered();
557 while let Some((_, common_occurrences)) = lists_reverse_indices.pop() {
558 if common_occurrences.len() < 2 {
559 continue;
561 }
562 let first_occurrence = &common_occurrences[0];
566
567 let mut common_length = 2;
569 loop {
570 let m = lists[first_occurrence.chunk_group][first_occurrence.list]
571 .get(first_occurrence.entry + common_length - 1);
572 if m.is_some()
573 && common_occurrences.iter().skip(1).all(
574 |ListOccurrence {
575 chunk_group,
576 list,
577 entry,
578 }| {
579 lists[*chunk_group][*list].get(*entry + common_length - 1) == m
580 },
581 )
582 {
583 common_length += 1;
584 continue;
585 }
586
587 common_length -= 1;
589 break;
590 }
591
592 let common_list = lists[first_occurrence.chunk_group][first_occurrence.list]
597 [first_occurrence.entry..first_occurrence.entry + common_length]
598 .to_vec();
599
600 let common_list_index = lists[LISTS_COMMON_IDX].len();
601 lists[LISTS_COMMON_IDX].push(common_list.clone());
602
603 for (i, &m) in common_list.iter().enumerate().skip(1) {
606 let occurrences = lists_reverse_indices
607 .get_mut(&m)
608 .context("every module should have occurrences")?;
609 for common_occurrence in &common_occurrences {
610 let removed = occurrences.swap_remove(&ListOccurrence {
611 chunk_group: common_occurrence.chunk_group,
612 list: common_occurrence.list,
613 entry: common_occurrence.entry + i,
614 });
615 debug_assert!(removed);
616 }
617 occurrences.insert(ListOccurrence {
618 chunk_group: LISTS_COMMON_IDX,
619 list: common_list_index,
620 entry: i,
621 });
622 }
623
624 for common_occurrence in &common_occurrences {
625 let list = &mut lists[common_occurrence.chunk_group][common_occurrence.list];
626 let after_list = list.split_off(common_occurrence.entry + common_length);
627 list.truncate(common_occurrence.entry);
628 let before_list = &*list;
629
630 {
637 let before_list =
638 FxHashSet::from_iter(before_list.iter().map(|m| ResolvedVc::upcast(*m)));
639 let common_list =
640 FxHashSet::from_iter(common_list.iter().map(|m| ResolvedVc::upcast(*m)));
641 let after_list =
642 FxHashSet::from_iter(after_list.iter().map(|m| ResolvedVc::upcast(*m)));
643
644 let references_from_before = before_list
645 .iter()
646 .filter_map(|m| intra_group_references.get(m))
647 .flatten()
648 .copied()
649 .filter(|m| common_list.contains(m) || after_list.contains(m))
650 .collect::<FxHashSet<_>>();
651 let references_from_common = common_list
652 .iter()
653 .filter_map(|m| intra_group_references.get(m))
654 .flatten()
655 .filter(|m| before_list.contains(m) || after_list.contains(m))
656 .collect::<FxHashSet<_>>();
657 let references_from_after = after_list
658 .iter()
659 .filter_map(|m| intra_group_references.get(m))
660 .flatten()
661 .copied()
662 .filter(|m| before_list.contains(m) || common_list.contains(m))
663 .collect::<FxHashSet<_>>();
664
665 let modules_to_expose = before_list
666 .iter()
667 .chain(common_list.iter())
668 .chain(after_list.iter())
669 .copied()
670 .filter(|m| {
671 references_from_before.contains(m)
672 || references_from_common.contains(m)
673 || references_from_after.contains(m)
674 });
675
676 exposed_modules_imported.extend(modules_to_expose);
677 }
678
679 if !after_list.is_empty() {
682 let after_index = lists[LISTS_COMMON_IDX].len();
683 lists[LISTS_COMMON_IDX].push(after_list.clone());
684 for (i, &m) in after_list.iter().enumerate() {
685 let Some(occurrences) = lists_reverse_indices.get_mut(&m) else {
686 bail!("Couldn't find module in reverse list");
687 };
688
689 let removed = occurrences.swap_remove(&ListOccurrence {
690 chunk_group: common_occurrence.chunk_group,
691 list: common_occurrence.list,
692 entry: common_occurrence.entry + common_length + i,
693 });
694 debug_assert!(removed);
695
696 occurrences.insert(ListOccurrence {
697 chunk_group: LISTS_COMMON_IDX,
698 list: after_index,
699 entry: i,
700 });
701 }
702 }
703 }
704 }
705
706 let lists = lists.into_iter().flatten().collect::<FxHashSet<_>>();
708
709 drop(inner_span);
710 let inner_span = tracing::info_span!("merging");
711 let result = lists
713 .into_iter()
714 .map(async |list| {
715 if list.len() < 2 {
716 return Ok(None);
718 }
719
720 let list_set = list
721 .iter()
722 .map(|&m| ResolvedVc::upcast::<Box<dyn Module>>(m))
723 .collect::<FxIndexSet<_>>();
724
725 let entry_points = list
726 .iter()
727 .filter(|m| {
728 intra_group_references_rev
729 .get(&ResolvedVc::upcast(**m))
730 .is_none_or(|refs| refs.is_disjoint(&list_set))
731 })
732 .map(|m| **m)
733 .collect::<Vec<_>>();
734 debug_assert_ne!(entry_points.len(), 0);
735
736 let list_exposed = list
737 .iter()
738 .map(|&m| {
739 (
740 m,
741 if exposed_modules_imported.contains(&ResolvedVc::upcast(m)) {
742 MergeableModuleExposure::External
743 } else if exposed_modules_namespace.contains(&ResolvedVc::upcast(m)) {
744 MergeableModuleExposure::Internal
745 } else {
746 MergeableModuleExposure::None
747 },
748 )
749 })
750 .collect::<Vec<_>>();
751
752 let entry = *list.last().unwrap();
753 let result = entry
754 .merge(
755 MergeableModulesExposed::interned(list_exposed),
756 MergeableModules::interned(entry_points),
757 )
758 .to_resolved()
759 .await?;
760
761 let list_len = list.len();
762 Ok(Some((
763 ResolvedVc::upcast::<Box<dyn Module>>(entry),
764 result,
765 list.into_iter()
766 .take(list_len - 1)
767 .map(ResolvedVc::upcast::<Box<dyn Module>>)
768 .collect::<Vec<_>>(),
769 )))
770 })
771 .try_join()
772 .instrument(inner_span)
773 .await?;
774
775 #[allow(clippy::type_complexity)]
776 let mut replacements: FxHashMap<
777 ResolvedVc<Box<dyn Module>>,
778 ResolvedVc<Box<dyn ChunkableModule>>,
779 > = Default::default();
780 #[allow(clippy::type_complexity)]
781 let mut replacements_to_original: FxHashMap<
782 ResolvedVc<Box<dyn Module>>,
783 ResolvedVc<Box<dyn Module>>,
784 > = Default::default();
785 let mut included: FxHashSet<ResolvedVc<Box<dyn Module>>> = FxHashSet::default();
786
787 for (original, replacement, replacement_included) in result.into_iter().flatten() {
788 replacements.insert(original, replacement);
789 replacements_to_original.insert(ResolvedVc::upcast(replacement), original);
790 included.extend(replacement_included);
791 }
792
793 span.record("merged_groups", replacements.len());
794 span.record("included_modules", included.len());
795
796 Ok(MergedModuleInfo {
797 replacements: ResolvedVc::cell(replacements),
798 replacements_to_original: ResolvedVc::cell(replacements_to_original),
799 included: ResolvedVc::cell(included),
800 }
801 .cell())
802 }
803 .instrument(span_outer)
804 .await
805}