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 _|
176 -> Result<GraphTraversalAction> {
177 let (parent_module, hoisted) =
180 parent_info.map_or((None, false), |(node, ty, _)| {
181 (
182 Some(node),
183 match &ty.chunking_type {
184 ChunkingType::Parallel { hoisted, .. } => *hoisted,
185 _ => false,
186 },
187 )
188 });
189 let module = node;
190
191 Ok(if parent_module.is_some_and(|p| p == module) {
192 GraphTraversalAction::Skip
194 } else if hoisted
195 && let Some(parent_module) = parent_module
196 && mergeable.contains(&parent_module)
197 && mergeable.contains(&module)
198 && !async_modules.contains(&parent_module)
199 && !async_modules.contains(&module)
200 {
201 module_merged_groups.entry(node).or_default();
206 let [Some(parent_merged_groups), Some(current_merged_groups)] =
207 module_merged_groups.get_disjoint_mut([&parent_module, &node])
208 else {
209 bail!("unreachable except for eventual consistency");
211 };
212
213 if current_merged_groups.is_empty() {
214 *current_merged_groups = parent_merged_groups.clone();
216 GraphTraversalAction::Continue
217 } else if parent_merged_groups.is_proper_superset(current_merged_groups) {
218 **current_merged_groups |= &**parent_merged_groups;
220 GraphTraversalAction::Continue
221 } else {
222 GraphTraversalAction::Skip
224 }
225 } else {
226 if entry_modules.insert(module) {
229 let idx = next_index;
231 next_index += 1;
232
233 if module_merged_groups.entry(module).or_default().insert(idx) {
234 GraphTraversalAction::Continue
236 } else {
237 GraphTraversalAction::Skip
239 }
240 } else {
241 GraphTraversalAction::Skip
244 }
245 })
246 },
247 |successor, _| {
248 Ok(-*module_depth
251 .get(&successor)
252 .context("Module depth not found")?)
253 },
254 )?;
255
256 drop(inner_span);
257 let inner_span = tracing::info_span!("chunk group collection").entered();
258
259 span.record("visit_count", visit_count);
260
261 #[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
262 struct ListOccurrence {
263 entry: usize,
266 list: usize,
267 chunk_group: usize,
268 }
269
270 let mut lists;
275 let mut lists_reverse_indices: FxIndexMap<
276 ResolvedVc<Box<dyn MergeableModule>>,
277 FxIndexSet<ListOccurrence>,
278 > = FxIndexMap::default();
279
280 #[allow(non_snake_case)]
283 let LISTS_COMMON_IDX: usize;
284
285 #[allow(clippy::type_complexity)]
289 let mut intra_group_references: FxIndexMap<
290 ResolvedVc<Box<dyn Module>>,
291 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
292 > = FxIndexMap::default();
293 #[allow(clippy::type_complexity)]
296 let mut intra_group_references_rev: FxIndexMap<
297 ResolvedVc<Box<dyn Module>>,
298 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
299 > = FxIndexMap::default();
300
301 {
302 struct ChunkGroupResult {
303 first_chunk_group_idx: usize,
304 #[allow(clippy::type_complexity)]
305 list_lists: Vec<Vec<Vec<ResolvedVc<Box<dyn MergeableModule>>>>>,
306 lists_reverse_indices:
307 FxIndexMap<ResolvedVc<Box<dyn MergeableModule>>, FxIndexSet<ListOccurrence>>,
308 #[allow(clippy::type_complexity)]
309 intra_group_references_rev: FxIndexMap<
310 ResolvedVc<Box<dyn Module>>,
311 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
312 >,
313 }
314 let span = tracing::info_span!("map chunk groups").entered();
315
316 let result = turbo_tasks::parallel::map_collect_chunked_owned::<_, _, Result<Vec<_>>>(
317 chunk_group_info.chunk_groups.iter().enumerate().collect(),
319 |chunk| {
320 let mut list_lists = vec![];
321 let mut lists_reverse_indices: FxIndexMap<
322 ResolvedVc<Box<dyn MergeableModule>>,
323 FxIndexSet<ListOccurrence>,
324 > = FxIndexMap::default();
325 #[allow(clippy::type_complexity)]
326 let mut intra_group_references_rev: FxIndexMap<
327 ResolvedVc<Box<dyn Module>>,
328 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
329 > = FxIndexMap::default();
330
331 let mut chunk = chunk.peekable();
332 let first_chunk_group_idx = chunk.peek().unwrap().0;
333
334 for (chunk_group_idx, chunk_group) in chunk {
335 let mut lists = vec![];
336
337 let mut chunk_lists: FxHashMap<&RoaringBitmapWrapper, usize> =
341 FxHashMap::with_capacity_and_hasher(
342 module_merged_groups.len() / chunk_group_info.chunk_groups.len(),
343 Default::default(),
344 );
345
346 let mut visited = FxHashSet::default();
350
351 module_graph.traverse_edges_dfs(
352 chunk_group.entries(),
353 &mut (),
354 |parent_info, node, _| {
355 if parent_info.is_none_or(|(_, r)| r.chunking_type.is_parallel())
356 && visited.insert(node)
357 {
358 Ok(GraphTraversalAction::Continue)
359 } else {
360 Ok(GraphTraversalAction::Exclude)
361 }
362 },
363 |parent_info, node, _| {
364 let module = node;
365 let bitmap = module_merged_groups
366 .get(&module)
367 .context("every module should have a bitmap")?;
368
369 if mergeable.contains(&module) {
370 let mergeable_module =
371 ResolvedVc::try_downcast::<Box<dyn MergeableModule>>(
372 module,
373 )
374 .context(
375 "found mergeable module which is not a MergeableModule",
376 )?;
377 match chunk_lists.entry(bitmap) {
378 Entry::Vacant(e) => {
379 let idx = lists.len();
381 e.insert(idx);
382 lists.push(vec![mergeable_module]);
383 lists_reverse_indices
384 .entry(mergeable_module)
385 .or_default()
386 .insert(ListOccurrence {
387 chunk_group: chunk_group_idx,
388 list: idx,
389 entry: 0,
390 });
391 }
392 Entry::Occupied(e) => {
393 let list_idx = *e.get();
394 let list = &mut lists[list_idx];
395 list.push(mergeable_module);
396 lists_reverse_indices
397 .entry(mergeable_module)
398 .or_default()
399 .insert(ListOccurrence {
400 chunk_group: chunk_group_idx,
401 list: list_idx,
402 entry: list.len() - 1,
403 });
404 }
405 }
406 }
407
408 if let Some((parent, _)) = parent_info {
409 let same_bitmap = module_merged_groups
410 .get(&parent)
411 .context("every module should have a bitmap")?
412 == module_merged_groups
413 .get(&module)
414 .context("every module should have a bitmap")?;
415
416 if same_bitmap {
417 intra_group_references_rev
418 .entry(module)
419 .or_default()
420 .insert(parent);
421 }
422 }
423 Ok(())
424 },
425 false,
426 )?;
427
428 list_lists.push(lists);
429 }
430 Ok(ChunkGroupResult {
431 first_chunk_group_idx,
432 list_lists,
433 lists_reverse_indices,
434 intra_group_references_rev,
435 })
436 },
437 )?;
438
439 drop(span);
440 let _span = tracing::info_span!("merging chunk group lists").entered();
441
442 lists_reverse_indices
443 .reserve_exact(result.iter().map(|r| r.lists_reverse_indices.len()).sum());
444 intra_group_references_rev.reserve_exact(
445 result
446 .iter()
447 .map(|r| r.intra_group_references_rev.len())
448 .sum(),
449 );
450
451 lists = vec![Default::default(); chunk_group_info.chunk_groups.len() + 1];
452 LISTS_COMMON_IDX = result.len();
453 for ChunkGroupResult {
454 first_chunk_group_idx,
455 list_lists: result_lists,
456 lists_reverse_indices: result_lists_reverse_indices,
457 intra_group_references_rev: result_intra_group_references_rev,
458 } in result
459 {
460 lists.splice(
461 first_chunk_group_idx..(first_chunk_group_idx + result_lists.len()),
462 result_lists,
463 );
464 for (module, occurrences) in result_lists_reverse_indices {
465 lists_reverse_indices
466 .entry(module)
467 .or_default()
468 .extend(occurrences);
469 }
470 for (module, occurrences) in result_intra_group_references_rev {
471 intra_group_references_rev
472 .entry(module)
473 .or_default()
474 .extend(occurrences);
475 }
476 }
477 }
478
479 drop(inner_span);
480 let inner_span = tracing::info_span!("exposed computation").entered();
481
482 lists_reverse_indices
484 .sort_by_cached_key(|_, b| b.iter().map(|o| o.entry).min().map(|v| -(v as i64)));
485
486 let mut exposed_modules_imported: FxHashSet<ResolvedVc<Box<dyn Module>>> =
490 FxHashSet::with_capacity_and_hasher(module_merged_groups.len(), Default::default());
491 let mut exposed_modules_namespace: FxHashSet<ResolvedVc<Box<dyn Module>>> =
492 FxHashSet::with_capacity_and_hasher(module_merged_groups.len(), Default::default());
493
494 module_graph.traverse_edges_dfs(
495 entries,
496 &mut (),
497 |_, _, _| Ok(GraphTraversalAction::Continue),
498 |parent_info, node, _| {
499 let module = node;
500
501 if let Some((parent, _)) = parent_info {
502 let same_bitmap = module_merged_groups
503 .get(&parent)
504 .context("every module should have a bitmap")?
505 == module_merged_groups
506 .get(&module)
507 .context("every module should have a bitmap")?;
508
509 if same_bitmap {
510 intra_group_references
511 .entry(parent)
512 .or_default()
513 .insert(module);
514 }
515 }
516
517 if match parent_info {
518 None => true,
519 Some((parent, _)) => {
520 module_merged_groups
521 .get(&parent)
522 .context("every module should have a bitmap")?
523 != module_merged_groups
524 .get(&module)
525 .context("every module should have a bitmap")?
526 }
527 } {
528 exposed_modules_imported.insert(module);
533 }
534 if parent_info.is_some_and(|(_, r)| {
535 matches!(
536 r.binding_usage.export,
537 ExportUsage::All | ExportUsage::PartialNamespaceObject(_)
538 )
539 }) {
540 exposed_modules_namespace.insert(module);
543 }
544 Ok(())
545 },
546 false,
547 )?;
548
549 drop(inner_span);
550 let inner_span = tracing::info_span!("reconciliation").entered();
551 while let Some((_, common_occurrences)) = lists_reverse_indices.pop() {
552 if common_occurrences.len() < 2 {
553 continue;
555 }
556 let first_occurrence = &common_occurrences[0];
560
561 let mut common_length = 2;
563 loop {
564 let m = lists[first_occurrence.chunk_group][first_occurrence.list]
565 .get(first_occurrence.entry + common_length - 1);
566 if m.is_some()
567 && common_occurrences.iter().skip(1).all(
568 |ListOccurrence {
569 chunk_group,
570 list,
571 entry,
572 }| {
573 lists[*chunk_group][*list].get(*entry + common_length - 1) == m
574 },
575 )
576 {
577 common_length += 1;
578 continue;
579 }
580
581 common_length -= 1;
583 break;
584 }
585
586 let common_list = lists[first_occurrence.chunk_group][first_occurrence.list]
591 [first_occurrence.entry..first_occurrence.entry + common_length]
592 .to_vec();
593
594 let common_list_index = lists[LISTS_COMMON_IDX].len();
595 lists[LISTS_COMMON_IDX].push(common_list.clone());
596
597 for (i, &m) in common_list.iter().enumerate().skip(1) {
600 let occurrences = lists_reverse_indices
601 .get_mut(&m)
602 .context("every module should have occurrences")?;
603 for common_occurrence in &common_occurrences {
604 let removed = occurrences.swap_remove(&ListOccurrence {
605 chunk_group: common_occurrence.chunk_group,
606 list: common_occurrence.list,
607 entry: common_occurrence.entry + i,
608 });
609 debug_assert!(removed);
610 }
611 occurrences.insert(ListOccurrence {
612 chunk_group: LISTS_COMMON_IDX,
613 list: common_list_index,
614 entry: i,
615 });
616 }
617
618 for common_occurrence in &common_occurrences {
619 let list = &mut lists[common_occurrence.chunk_group][common_occurrence.list];
620 let after_list = list.split_off(common_occurrence.entry + common_length);
621 list.truncate(common_occurrence.entry);
622 let before_list = &*list;
623
624 {
631 let before_list =
632 FxHashSet::from_iter(before_list.iter().map(|m| ResolvedVc::upcast(*m)));
633 let common_list =
634 FxHashSet::from_iter(common_list.iter().map(|m| ResolvedVc::upcast(*m)));
635 let after_list =
636 FxHashSet::from_iter(after_list.iter().map(|m| ResolvedVc::upcast(*m)));
637
638 let references_from_before = before_list
639 .iter()
640 .filter_map(|m| intra_group_references.get(m))
641 .flatten()
642 .copied()
643 .filter(|m| common_list.contains(m) || after_list.contains(m))
644 .collect::<FxHashSet<_>>();
645 let references_from_common = common_list
646 .iter()
647 .filter_map(|m| intra_group_references.get(m))
648 .flatten()
649 .filter(|m| before_list.contains(m) || after_list.contains(m))
650 .collect::<FxHashSet<_>>();
651 let references_from_after = after_list
652 .iter()
653 .filter_map(|m| intra_group_references.get(m))
654 .flatten()
655 .copied()
656 .filter(|m| before_list.contains(m) || common_list.contains(m))
657 .collect::<FxHashSet<_>>();
658
659 let modules_to_expose = before_list
660 .iter()
661 .chain(common_list.iter())
662 .chain(after_list.iter())
663 .copied()
664 .filter(|m| {
665 references_from_before.contains(m)
666 || references_from_common.contains(m)
667 || references_from_after.contains(m)
668 });
669
670 exposed_modules_imported.extend(modules_to_expose);
671 }
672
673 if !after_list.is_empty() {
676 let after_index = lists[LISTS_COMMON_IDX].len();
677 lists[LISTS_COMMON_IDX].push(after_list.clone());
678 for (i, &m) in after_list.iter().enumerate() {
679 let Some(occurrences) = lists_reverse_indices.get_mut(&m) else {
680 bail!("Couldn't find module in reverse list");
681 };
682
683 let removed = occurrences.swap_remove(&ListOccurrence {
684 chunk_group: common_occurrence.chunk_group,
685 list: common_occurrence.list,
686 entry: common_occurrence.entry + common_length + i,
687 });
688 debug_assert!(removed);
689
690 occurrences.insert(ListOccurrence {
691 chunk_group: LISTS_COMMON_IDX,
692 list: after_index,
693 entry: i,
694 });
695 }
696 }
697 }
698 }
699
700 let lists = lists.into_iter().flatten().collect::<FxHashSet<_>>();
702
703 drop(inner_span);
704 let inner_span = tracing::info_span!("merging");
705 let result = lists
707 .into_iter()
708 .map(async |list| {
709 if list.len() < 2 {
710 return Ok(None);
712 }
713
714 let list_set = list
715 .iter()
716 .map(|&m| ResolvedVc::upcast::<Box<dyn Module>>(m))
717 .collect::<FxIndexSet<_>>();
718
719 let entry_points = list
720 .iter()
721 .filter(|m| {
722 intra_group_references_rev
723 .get(&ResolvedVc::upcast(**m))
724 .is_none_or(|refs| refs.is_disjoint(&list_set))
725 })
726 .map(|m| **m)
727 .collect::<Vec<_>>();
728 debug_assert_ne!(entry_points.len(), 0);
729
730 let list_exposed = list
731 .iter()
732 .map(|&m| {
733 (
734 m,
735 if exposed_modules_imported.contains(&ResolvedVc::upcast(m)) {
736 MergeableModuleExposure::External
737 } else if exposed_modules_namespace.contains(&ResolvedVc::upcast(m)) {
738 MergeableModuleExposure::Internal
739 } else {
740 MergeableModuleExposure::None
741 },
742 )
743 })
744 .collect::<Vec<_>>();
745
746 let entry = *list.last().unwrap();
747 let result = entry
748 .merge(
749 MergeableModulesExposed::interned(list_exposed),
750 MergeableModules::interned(entry_points),
751 )
752 .to_resolved()
753 .await?;
754
755 let list_len = list.len();
756 Ok(Some((
757 ResolvedVc::upcast::<Box<dyn Module>>(entry),
758 result,
759 list.into_iter()
760 .take(list_len - 1)
761 .map(ResolvedVc::upcast::<Box<dyn Module>>)
762 .collect::<Vec<_>>(),
763 )))
764 })
765 .try_join()
766 .instrument(inner_span)
767 .await?;
768
769 #[allow(clippy::type_complexity)]
770 let mut replacements: FxHashMap<
771 ResolvedVc<Box<dyn Module>>,
772 Option<ResolvedVc<Box<dyn ChunkableModule>>>,
773 > = Default::default();
774 #[allow(clippy::type_complexity)]
775 let mut replacements_to_original: FxHashMap<
776 ResolvedVc<Box<dyn Module>>,
777 ResolvedVc<Box<dyn Module>>,
778 > = Default::default();
779 let mut merged_groups = 0;
780 let mut included_modules = 0;
781
782 for (original, replacement, replacement_included) in result.into_iter().flatten() {
783 replacements.insert(original, Some(replacement));
784 replacements_to_original.insert(ResolvedVc::upcast(replacement), original);
785 merged_groups += 1;
786 included_modules += replacement_included.len();
787 for included in replacement_included {
788 replacements.insert(included, None);
789 }
790 }
791
792 span.record("merged_groups", merged_groups);
793 span.record("included_modules", included_modules);
794
795 Ok(MergedModuleInfo {
796 replacements: ResolvedVc::cell(replacements),
797 replacements_to_original: ResolvedVc::cell(replacements_to_original),
798 }
799 .cell())
800 }
801 .instrument(span_outer)
802 .await
803}