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]
21pub struct MergedModuleInfo {
22 #[allow(clippy::type_complexity)]
24 pub replacements: FxHashMap<ResolvedVc<Box<dyn Module>>, ResolvedVc<Box<dyn ChunkableModule>>>,
25 #[allow(clippy::type_complexity)]
28 pub replacements_to_original:
29 FxHashMap<ResolvedVc<Box<dyn Module>>, ResolvedVc<Box<dyn Module>>>,
30 pub included: FxHashSet<ResolvedVc<Box<dyn Module>>>,
32}
33
34impl MergedModuleInfo {
35 pub fn should_replace_module(
37 &self,
38 module: ResolvedVc<Box<dyn Module>>,
39 ) -> Option<ResolvedVc<Box<dyn ChunkableModule>>> {
40 self.replacements.get(&module).copied()
41 }
42
43 pub fn get_original_module(
46 &self,
47 module: ResolvedVc<Box<dyn Module>>,
48 ) -> Option<ResolvedVc<Box<dyn Module>>> {
49 self.replacements_to_original.get(&module).copied()
50 }
51
52 pub fn should_create_chunk_item_for(&self, module: ResolvedVc<Box<dyn Module>>) -> bool {
55 !self.included.contains(&module)
56 }
57}
58
59pub async fn compute_merged_modules(module_graph: Vc<ModuleGraph>) -> Result<Vc<MergedModuleInfo>> {
65 let span_outer = tracing::info_span!(
66 "compute merged modules",
67 module_count = tracing::field::Empty,
68 visit_count = tracing::field::Empty,
69 merged_groups = tracing::field::Empty,
70 included_modules = tracing::field::Empty
71 );
72
73 let span = span_outer.clone();
74 async move {
75 let async_module_info = module_graph.async_module_info().await?;
76 let chunk_group_info = module_graph.chunk_group_info().await?;
77 let module_graph = module_graph.read_graphs().await?;
78
79 let graphs = &module_graph.graphs;
80 let module_count = graphs.iter().map(|g| g.graph.node_count()).sum::<usize>();
81 span.record("module_count", module_count);
82
83 let entries = graphs
85 .iter()
86 .flat_map(|g| g.entries.iter())
87 .flat_map(|g| g.entries())
88 .collect::<Vec<_>>();
89
90 let module_depth = {
92 let _inner_span = tracing::info_span!("compute depth").entered();
93
94 let mut module_depth =
95 FxHashMap::with_capacity_and_hasher(module_count, Default::default());
96 module_graph.traverse_edges_from_entries_bfs(
97 entries.iter().copied(),
98 |parent, node| {
99 if let Some((parent, _)) = parent {
100 let parent_depth = *module_depth
101 .get(&parent)
102 .context("Module depth not found")?;
103 module_depth.entry(node).or_insert(parent_depth + 1);
104 } else {
105 module_depth.insert(node, 0);
106 };
107
108 Ok(GraphTraversalAction::Continue)
109 },
110 )?;
111 module_depth
112 };
113
114 let mut module_merged_groups: FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper> =
118 FxHashMap::with_capacity_and_hasher(module_count, Default::default());
119 let mut entry_modules =
121 FxHashSet::with_capacity_and_hasher(module_count, Default::default());
122
123 let inner_span = tracing::info_span!("collect mergeable modules");
124 let mergeable = graphs
125 .iter()
126 .flat_map(|g| g.iter_nodes())
127 .map(async |module| {
128 if let Some(mergeable) =
129 ResolvedVc::try_downcast::<Box<dyn MergeableModule>>(module)
130 && *mergeable.is_mergeable().await?
131 {
132 return Ok(Some(module));
133 }
134 Ok(None)
135 })
136 .try_flat_join()
137 .instrument(inner_span)
138 .await?
139 .into_iter()
140 .collect::<FxHashSet<_>>();
141
142 let inner_span = tracing::info_span!("fixed point traversal").entered();
143
144 let mut next_index = 0u32;
145 let visit_count = module_graph.traverse_edges_fixed_point_with_priority(
146 entries
147 .iter()
148 .map(|e| Ok((*e, -*module_depth.get(e).context("Module depth not found")?)))
149 .collect::<Result<Vec<_>>>()?,
150 &mut (),
151 |parent_info: Option<(ResolvedVc<Box<dyn Module>>, &'_ RefData, _)>,
152 node: ResolvedVc<Box<dyn Module>>,
153 _|
154 -> Result<GraphTraversalAction> {
155 let (parent_module, hoisted) =
158 parent_info.map_or((None, false), |(node, ty, _)| {
159 (
160 Some(node),
161 match &ty.chunking_type {
162 ChunkingType::Parallel { hoisted, .. } => *hoisted,
163 _ => false,
164 },
165 )
166 });
167 let module = node;
168
169 Ok(if parent_module.is_some_and(|p| p == module) {
170 GraphTraversalAction::Skip
172 } else if hoisted
173 && let Some(parent_module) = parent_module
174 && mergeable.contains(&parent_module)
175 && mergeable.contains(&module)
176 && !async_module_info.contains(&parent_module)
177 && !async_module_info.contains(&module)
178 {
179 module_merged_groups.entry(node).or_default();
184 let [Some(parent_merged_groups), Some(current_merged_groups)] =
185 module_merged_groups.get_disjoint_mut([&parent_module, &node])
186 else {
187 bail!("unreachable except for eventual consistency");
189 };
190
191 if current_merged_groups.is_empty() {
192 *current_merged_groups = parent_merged_groups.clone();
194 GraphTraversalAction::Continue
195 } else if parent_merged_groups.is_proper_superset(current_merged_groups) {
196 **current_merged_groups |= &**parent_merged_groups;
198 GraphTraversalAction::Continue
199 } else {
200 GraphTraversalAction::Skip
202 }
203 } else {
204 if entry_modules.insert(module) {
207 let idx = next_index;
209 next_index += 1;
210
211 if module_merged_groups.entry(module).or_default().insert(idx) {
212 GraphTraversalAction::Continue
214 } else {
215 GraphTraversalAction::Skip
217 }
218 } else {
219 GraphTraversalAction::Skip
222 }
223 })
224 },
225 |successor, _| {
226 Ok(-*module_depth
229 .get(&successor)
230 .context("Module depth not found")?)
231 },
232 )?;
233
234 drop(inner_span);
235 let inner_span = tracing::info_span!("chunk group collection").entered();
236
237 span.record("visit_count", visit_count);
238
239 #[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
240 struct ListOccurrence {
241 entry: usize,
244 list: usize,
245 chunk_group: usize,
246 }
247
248 let mut lists;
253 let mut lists_reverse_indices: FxIndexMap<
254 ResolvedVc<Box<dyn MergeableModule>>,
255 FxIndexSet<ListOccurrence>,
256 > = FxIndexMap::default();
257
258 #[allow(non_snake_case)]
261 let LISTS_COMMON_IDX: usize;
262
263 #[allow(clippy::type_complexity)]
267 let mut intra_group_references: FxIndexMap<
268 ResolvedVc<Box<dyn Module>>,
269 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
270 > = FxIndexMap::default();
271 #[allow(clippy::type_complexity)]
274 let mut intra_group_references_rev: FxIndexMap<
275 ResolvedVc<Box<dyn Module>>,
276 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
277 > = FxIndexMap::default();
278
279 {
280 struct ChunkGroupResult {
281 first_chunk_group_idx: usize,
282 #[allow(clippy::type_complexity)]
283 list_lists: Vec<Vec<Vec<ResolvedVc<Box<dyn MergeableModule>>>>>,
284 lists_reverse_indices:
285 FxIndexMap<ResolvedVc<Box<dyn MergeableModule>>, FxIndexSet<ListOccurrence>>,
286 #[allow(clippy::type_complexity)]
287 intra_group_references_rev: FxIndexMap<
288 ResolvedVc<Box<dyn Module>>,
289 FxIndexSet<ResolvedVc<Box<dyn Module>>>,
290 >,
291 }
292 let span = tracing::info_span!("map chunk groups").entered();
293
294 let result = turbo_tasks::parallel::map_collect_chunked_owned::<_, _, Result<Vec<_>>>(
295 chunk_group_info.chunk_groups.iter().enumerate().collect(),
297 |chunk| {
298 let mut list_lists = vec![];
299 let mut lists_reverse_indices: FxIndexMap<
300 ResolvedVc<Box<dyn MergeableModule>>,
301 FxIndexSet<ListOccurrence>,
302 > = FxIndexMap::default();
303 #[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 let mut chunk = chunk.peekable();
310 let first_chunk_group_idx = chunk.peek().unwrap().0;
311
312 for (chunk_group_idx, chunk_group) in chunk {
313 let mut lists = vec![];
314
315 let mut chunk_lists: FxHashMap<&RoaringBitmapWrapper, usize> =
319 FxHashMap::with_capacity_and_hasher(
320 module_merged_groups.len() / chunk_group_info.chunk_groups.len(),
321 Default::default(),
322 );
323
324 let mut visited = FxHashSet::default();
328
329 module_graph.traverse_edges_from_entries_dfs(
330 chunk_group.entries(),
331 &mut (),
332 |parent_info, node, _| {
333 if parent_info.is_none_or(|(_, r)| r.chunking_type.is_parallel())
334 && visited.insert(node)
335 {
336 Ok(GraphTraversalAction::Continue)
337 } else {
338 Ok(GraphTraversalAction::Exclude)
339 }
340 },
341 |parent_info, node, _| {
342 let module = node;
343 let bitmap = module_merged_groups
344 .get(&module)
345 .context("every module should have a bitmap")?;
346
347 if mergeable.contains(&module) {
348 let mergeable_module =
349 ResolvedVc::try_downcast::<Box<dyn MergeableModule>>(
350 module,
351 )
352 .context(
353 "found mergeable module which is not a MergeableModule",
354 )?;
355 match chunk_lists.entry(bitmap) {
356 Entry::Vacant(e) => {
357 let idx = lists.len();
359 e.insert(idx);
360 lists.push(vec![mergeable_module]);
361 lists_reverse_indices
362 .entry(mergeable_module)
363 .or_default()
364 .insert(ListOccurrence {
365 chunk_group: chunk_group_idx,
366 list: idx,
367 entry: 0,
368 });
369 }
370 Entry::Occupied(e) => {
371 let list_idx = *e.get();
372 let list = &mut lists[list_idx];
373 list.push(mergeable_module);
374 lists_reverse_indices
375 .entry(mergeable_module)
376 .or_default()
377 .insert(ListOccurrence {
378 chunk_group: chunk_group_idx,
379 list: list_idx,
380 entry: list.len() - 1,
381 });
382 }
383 }
384 }
385
386 if let Some((parent, _)) = parent_info {
387 let same_bitmap = module_merged_groups
388 .get(&parent)
389 .context("every module should have a bitmap")?
390 == module_merged_groups
391 .get(&module)
392 .context("every module should have a bitmap")?;
393
394 if same_bitmap {
395 intra_group_references_rev
396 .entry(module)
397 .or_default()
398 .insert(parent);
399 }
400 }
401 Ok(())
402 },
403 )?;
404
405 list_lists.push(lists);
406 }
407 Ok(ChunkGroupResult {
408 first_chunk_group_idx,
409 list_lists,
410 lists_reverse_indices,
411 intra_group_references_rev,
412 })
413 },
414 )?;
415
416 drop(span);
417 let _span = tracing::info_span!("merging chunk group lists").entered();
418
419 lists_reverse_indices
420 .reserve_exact(result.iter().map(|r| r.lists_reverse_indices.len()).sum());
421 intra_group_references_rev.reserve_exact(
422 result
423 .iter()
424 .map(|r| r.intra_group_references_rev.len())
425 .sum(),
426 );
427
428 lists = vec![Default::default(); chunk_group_info.chunk_groups.len() + 1];
429 LISTS_COMMON_IDX = result.len();
430 for ChunkGroupResult {
431 first_chunk_group_idx,
432 list_lists: result_lists,
433 lists_reverse_indices: result_lists_reverse_indices,
434 intra_group_references_rev: result_intra_group_references_rev,
435 } in result
436 {
437 lists.splice(
438 first_chunk_group_idx..(first_chunk_group_idx + result_lists.len()),
439 result_lists,
440 );
441 for (module, occurrences) in result_lists_reverse_indices {
442 lists_reverse_indices
443 .entry(module)
444 .or_default()
445 .extend(occurrences);
446 }
447 for (module, occurrences) in result_intra_group_references_rev {
448 intra_group_references_rev
449 .entry(module)
450 .or_default()
451 .extend(occurrences);
452 }
453 }
454 }
455
456 drop(inner_span);
457 let inner_span = tracing::info_span!("exposed computation").entered();
458
459 lists_reverse_indices
461 .sort_by_cached_key(|_, b| b.iter().map(|o| o.entry).min().map(|v| -(v as i64)));
462
463 let mut exposed_modules_imported: FxHashSet<ResolvedVc<Box<dyn Module>>> =
467 FxHashSet::with_capacity_and_hasher(module_merged_groups.len(), Default::default());
468 let mut exposed_modules_namespace: FxHashSet<ResolvedVc<Box<dyn Module>>> =
469 FxHashSet::with_capacity_and_hasher(module_merged_groups.len(), Default::default());
470
471 module_graph.traverse_edges_from_entries_dfs(
472 entries,
473 &mut (),
474 |_, _, _| Ok(GraphTraversalAction::Continue),
475 |parent_info, node, _| {
476 let module = node;
477
478 if let Some((parent, _)) = parent_info {
479 let same_bitmap = module_merged_groups
480 .get(&parent)
481 .context("every module should have a bitmap")?
482 == module_merged_groups
483 .get(&module)
484 .context("every module should have a bitmap")?;
485
486 if same_bitmap {
487 intra_group_references
488 .entry(parent)
489 .or_default()
490 .insert(module);
491 }
492 }
493
494 if match parent_info {
495 None => true,
496 Some((parent, _)) => {
497 module_merged_groups
498 .get(&parent)
499 .context("every module should have a bitmap")?
500 != module_merged_groups
501 .get(&module)
502 .context("every module should have a bitmap")?
503 }
504 } {
505 exposed_modules_imported.insert(module);
510 }
511 if parent_info
512 .is_some_and(|(_, r)| matches!(r.binding_usage.export, ExportUsage::All))
513 {
514 exposed_modules_namespace.insert(module);
517 }
518 Ok(())
519 },
520 )?;
521
522 drop(inner_span);
523 let inner_span = tracing::info_span!("reconciliation").entered();
524 while let Some((_, common_occurrences)) = lists_reverse_indices.pop() {
525 if common_occurrences.len() < 2 {
526 continue;
528 }
529 let first_occurrence = &common_occurrences[0];
533
534 let mut common_length = 2;
536 loop {
537 let m = lists[first_occurrence.chunk_group][first_occurrence.list]
538 .get(first_occurrence.entry + common_length - 1);
539 if m.is_some()
540 && common_occurrences.iter().skip(1).all(
541 |ListOccurrence {
542 chunk_group,
543 list,
544 entry,
545 }| {
546 lists[*chunk_group][*list].get(*entry + common_length - 1) == m
547 },
548 )
549 {
550 common_length += 1;
551 continue;
552 }
553
554 common_length -= 1;
556 break;
557 }
558
559 let common_list = lists[first_occurrence.chunk_group][first_occurrence.list]
564 [first_occurrence.entry..first_occurrence.entry + common_length]
565 .to_vec();
566
567 let common_list_index = lists[LISTS_COMMON_IDX].len();
568 lists[LISTS_COMMON_IDX].push(common_list.clone());
569
570 for (i, &m) in common_list.iter().enumerate().skip(1) {
573 let occurrences = lists_reverse_indices
574 .get_mut(&m)
575 .context("every module should have occurrences")?;
576 for common_occurrence in &common_occurrences {
577 let removed = occurrences.swap_remove(&ListOccurrence {
578 chunk_group: common_occurrence.chunk_group,
579 list: common_occurrence.list,
580 entry: common_occurrence.entry + i,
581 });
582 debug_assert!(removed);
583 }
584 occurrences.insert(ListOccurrence {
585 chunk_group: LISTS_COMMON_IDX,
586 list: common_list_index,
587 entry: i,
588 });
589 }
590
591 for common_occurrence in &common_occurrences {
592 let list = &mut lists[common_occurrence.chunk_group][common_occurrence.list];
593 let after_list = list.split_off(common_occurrence.entry + common_length);
594 list.truncate(common_occurrence.entry);
595 let before_list = &*list;
596
597 {
604 let before_list =
605 FxHashSet::from_iter(before_list.iter().map(|m| ResolvedVc::upcast(*m)));
606 let common_list =
607 FxHashSet::from_iter(common_list.iter().map(|m| ResolvedVc::upcast(*m)));
608 let after_list =
609 FxHashSet::from_iter(after_list.iter().map(|m| ResolvedVc::upcast(*m)));
610
611 let references_from_before = before_list
612 .iter()
613 .filter_map(|m| intra_group_references.get(m))
614 .flatten()
615 .copied()
616 .filter(|m| common_list.contains(m) || after_list.contains(m))
617 .collect::<FxHashSet<_>>();
618 let references_from_common = common_list
619 .iter()
620 .filter_map(|m| intra_group_references.get(m))
621 .flatten()
622 .filter(|m| before_list.contains(m) || after_list.contains(m))
623 .collect::<FxHashSet<_>>();
624 let references_from_after = after_list
625 .iter()
626 .filter_map(|m| intra_group_references.get(m))
627 .flatten()
628 .copied()
629 .filter(|m| before_list.contains(m) || common_list.contains(m))
630 .collect::<FxHashSet<_>>();
631
632 let modules_to_expose = before_list
633 .iter()
634 .chain(common_list.iter())
635 .chain(after_list.iter())
636 .copied()
637 .filter(|m| {
638 references_from_before.contains(m)
639 || references_from_common.contains(m)
640 || references_from_after.contains(m)
641 });
642
643 exposed_modules_imported.extend(modules_to_expose);
644 }
645
646 if !after_list.is_empty() {
649 let after_index = lists[LISTS_COMMON_IDX].len();
650 lists[LISTS_COMMON_IDX].push(after_list.clone());
651 for (i, &m) in after_list.iter().enumerate() {
652 let Some(occurrences) = lists_reverse_indices.get_mut(&m) else {
653 bail!("Couldn't find module in reverse list");
654 };
655
656 let removed = occurrences.swap_remove(&ListOccurrence {
657 chunk_group: common_occurrence.chunk_group,
658 list: common_occurrence.list,
659 entry: common_occurrence.entry + common_length + i,
660 });
661 debug_assert!(removed);
662
663 occurrences.insert(ListOccurrence {
664 chunk_group: LISTS_COMMON_IDX,
665 list: after_index,
666 entry: i,
667 });
668 }
669 }
670 }
671 }
672
673 let lists = lists.into_iter().flatten().collect::<FxHashSet<_>>();
675
676 drop(inner_span);
677 let inner_span = tracing::info_span!("merging");
678 let result = lists
680 .into_iter()
681 .map(async |list| {
682 if list.len() < 2 {
683 return Ok(None);
685 }
686
687 let list_set = list
688 .iter()
689 .map(|&m| ResolvedVc::upcast::<Box<dyn Module>>(m))
690 .collect::<FxIndexSet<_>>();
691
692 let entry_points = list
693 .iter()
694 .filter(|m| {
695 intra_group_references_rev
696 .get(&ResolvedVc::upcast(**m))
697 .is_none_or(|refs| refs.is_disjoint(&list_set))
698 })
699 .map(|m| **m)
700 .collect::<Vec<_>>();
701 debug_assert_ne!(entry_points.len(), 0);
702
703 let list_exposed = list
704 .iter()
705 .map(|&m| {
706 (
707 m,
708 if exposed_modules_imported.contains(&ResolvedVc::upcast(m)) {
709 MergeableModuleExposure::External
710 } else if exposed_modules_namespace.contains(&ResolvedVc::upcast(m)) {
711 MergeableModuleExposure::Internal
712 } else {
713 MergeableModuleExposure::None
714 },
715 )
716 })
717 .collect::<Vec<_>>();
718
719 let entry = *list.last().unwrap();
720 let result = entry
721 .merge(
722 MergeableModulesExposed::interned(list_exposed),
723 MergeableModules::interned(entry_points),
724 )
725 .to_resolved()
726 .await?;
727
728 let list_len = list.len();
729 Ok(Some((
730 ResolvedVc::upcast::<Box<dyn Module>>(entry),
731 result,
732 list.into_iter()
733 .take(list_len - 1)
734 .map(ResolvedVc::upcast::<Box<dyn Module>>)
735 .collect::<Vec<_>>(),
736 )))
737 })
738 .try_join()
739 .instrument(inner_span)
740 .await?;
741
742 #[allow(clippy::type_complexity)]
743 let mut replacements: FxHashMap<
744 ResolvedVc<Box<dyn Module>>,
745 ResolvedVc<Box<dyn ChunkableModule>>,
746 > = Default::default();
747 #[allow(clippy::type_complexity)]
748 let mut replacements_to_original: FxHashMap<
749 ResolvedVc<Box<dyn Module>>,
750 ResolvedVc<Box<dyn Module>>,
751 > = Default::default();
752 let mut included: FxHashSet<ResolvedVc<Box<dyn Module>>> = FxHashSet::default();
753
754 for (original, replacement, replacement_included) in result.into_iter().flatten() {
755 replacements.insert(original, replacement);
756 replacements_to_original.insert(ResolvedVc::upcast(replacement), original);
757 included.extend(replacement_included);
758 }
759
760 span.record("merged_groups", replacements.len());
761 span.record("included_modules", included.len());
762
763 Ok(MergedModuleInfo {
764 replacements,
765 replacements_to_original,
766 included,
767 }
768 .cell())
769 }
770 .instrument(span_outer)
771 .await
772}