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