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