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