Repository: vercel/next.js
Pull request: Turbopack: compute ordered entries in module batches

### What?

When computing the module batches, also compute the correct postorder of isolated and shared chunk groups

Closes PACK-4189

Changed files:
filename: turbopack/crates/turbopack-core/src/chunk/mod.rs
patch: @@ -205,25 +205,34 @@ pub enum ChunkingType {
 
 impl ChunkingType {
     pub fn is_inherit_async(&self) -> bool {
-        match self {
-            ChunkingType::Parallel => false,
-            ChunkingType::ParallelInheritAsync => true,
-            ChunkingType::Async => false,
-            ChunkingType::Isolated { .. } => false,
-            ChunkingType::Shared { inherit_async, .. } => *inherit_async,
-            ChunkingType::Traced => false,
-        }
+        matches!(
+            self,
+            ChunkingType::ParallelInheritAsync
+                | ChunkingType::Shared {
+                    inherit_async: true,
+                    ..
+                }
+        )
     }
 
     pub fn is_parallel(&self) -> bool {
-        match self {
-            ChunkingType::Parallel => true,
-            ChunkingType::ParallelInheritAsync => true,
-            ChunkingType::Async => false,
-            ChunkingType::Isolated { .. } => false,
-            ChunkingType::Shared { .. } => false,
-            ChunkingType::Traced => false,
-        }
+        matches!(
+            self,
+            ChunkingType::Parallel | ChunkingType::ParallelInheritAsync
+        )
+    }
+
+    pub fn is_merged(&self) -> bool {
+        matches!(
+            self,
+            ChunkingType::Isolated {
+                merge_tag: Some(_),
+                ..
+            } | ChunkingType::Shared {
+                merge_tag: Some(_),
+                ..
+            }
+        )
     }
 
     pub fn without_inherit_async(&self) -> Self {

filename: turbopack/crates/turbopack-core/src/module_graph/chunk_group_info.rs
patch: @@ -88,6 +88,8 @@ pub struct ChunkGroupInfo {
     pub module_chunk_groups: FxHashMap<ResolvedVc<Box<dyn Module>>, RoaringBitmapWrapper>,
     #[turbo_tasks(trace_ignore)]
     pub chunk_groups: FxIndexSet<ChunkGroup>,
+    #[turbo_tasks(trace_ignore)]
+    pub chunk_group_keys: FxIndexSet<ChunkGroupKey>,
 }
 
 #[turbo_tasks::value_impl]
@@ -174,6 +176,19 @@ pub enum ChunkGroup {
 }
 
 impl ChunkGroup {
+    /// Returns the parent group when this chunk group is a merged group. In that case `entries()`
+    /// are in unspecified order.
+    pub fn get_merged_parent(&self) -> Option<usize> {
+        match self {
+            ChunkGroup::IsolatedMerged { parent, .. } | ChunkGroup::SharedMerged { parent, .. } => {
+                Some(*parent)
+            }
+            _ => None,
+        }
+    }
+
+    /// Iterates over the entries of the chunk group. When `get_merged_parent` is Some, the order is
+    /// unspecified.
     pub fn entries(&self) -> impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + Clone + '_ {
         match self {
             ChunkGroup::Async(e) | ChunkGroup::Isolated(e) | ChunkGroup::Shared(e) => {
@@ -185,6 +200,15 @@ impl ChunkGroup {
         }
     }
 
+    pub fn entries_count(&self) -> usize {
+        match self {
+            ChunkGroup::Async(_) | ChunkGroup::Isolated(_) | ChunkGroup::Shared(_) => 1,
+            ChunkGroup::Entry(entries)
+            | ChunkGroup::IsolatedMerged { entries, .. }
+            | ChunkGroup::SharedMerged { entries, .. } => entries.len(),
+        }
+    }
+
     pub async fn debug_str(&self, chunk_group_info: &ChunkGroupInfo) -> Result<String> {
         Ok(match self {
             ChunkGroup::Entry(entries) => format!(
@@ -245,8 +269,8 @@ impl ChunkGroup {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-enum ChunkGroupKey {
+#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
+pub enum ChunkGroupKey {
     /// e.g. a page
     Entry(Vec<ResolvedVc<Box<dyn Module>>>),
     /// a module with an incoming async edge
@@ -270,6 +294,12 @@ enum ChunkGroupKey {
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct ChunkGroupId(u32);
 
+impl From<usize> for ChunkGroupId {
+    fn from(id: usize) -> Self {
+        Self(id as u32)
+    }
+}
+
 impl Deref for ChunkGroupId {
     type Target = u32;
     fn deref(&self) -> &Self::Target {
@@ -621,6 +651,7 @@ pub async fn compute_chunk_group_info(graph: &ModuleGraph) -> Result<Vc<ChunkGro
 
         Ok(ChunkGroupInfo {
             module_chunk_groups,
+            chunk_group_keys: chunk_groups_map.keys().cloned().collect(),
             chunk_groups: chunk_groups_map
                 .into_iter()
                 .map(|(k, (_, merged_entries))| match k {

filename: turbopack/crates/turbopack-core/src/module_graph/module_batches.rs
patch: @@ -20,7 +20,7 @@ use crate::{
     chunk::{ChunkableModule, ChunkingType},
     module::Module,
     module_graph::{
-        chunk_group_info::{ChunkGroupInfo, RoaringBitmapWrapper},
+        chunk_group_info::{ChunkGroupInfo, ChunkGroupKey, RoaringBitmapWrapper},
         module_batch::{ModuleBatch, ModuleBatchGroup, ModuleOrBatch},
         traced_di_graph::{iter_neighbors_rev, TracedDiGraph},
         GraphTraversalAction, ModuleGraph,
@@ -48,6 +48,8 @@ pub struct ModuleBatchesGraphEdge {
     pub module: Option<ResolvedVc<Box<dyn Module>>>,
 }
 
+type EntriesList = FxIndexSet<ResolvedVc<Box<dyn Module>>>;
+
 #[turbo_tasks::value(cell = "new", eq = "manual", into = "new")]
 pub struct ModuleBatchesGraph {
     graph: TracedDiGraph<ModuleOrBatch, ModuleBatchesGraphEdge>,
@@ -61,6 +63,12 @@ pub struct ModuleBatchesGraph {
     #[turbo_tasks(trace_ignore)]
     entries: FxHashMap<ResolvedVc<Box<dyn Module>>, NodeIndex>,
     batch_groups: FxHashMap<ModuleOrBatch, ResolvedVc<ModuleBatchGroup>>,
+
+    /// For chunk groups where the postorder of entries is different than the order of the
+    /// `ChunkGroup::entries()` this contains Some with the postorder list of entries of that chunk
+    /// group. The index in this list corresponds to the index in the
+    /// chunk_group_info.chunk_groups.
+    ordered_entries: Vec<Option<EntriesList>>,
 }
 
 impl ModuleBatchesGraph {
@@ -79,6 +87,28 @@ impl ModuleBatchesGraph {
         Ok(*entry)
     }
 
+    pub fn get_ordered_entries<'l>(
+        &'l self,
+        chunk_group_info: &'l ChunkGroupInfo,
+        idx: usize,
+    ) -> impl Iterator<Item = ResolvedVc<Box<dyn Module>>> + 'l {
+        if let Some(ordered_entries) = self
+            .ordered_entries
+            .get(idx)
+            .as_ref()
+            .and_then(|o| o.as_ref())
+        {
+            if let Some(chunk_group) = chunk_group_info.chunk_groups.get_index(idx) {
+                debug_assert_eq!(ordered_entries.len(), chunk_group.entries_count());
+            }
+            Either::Left(Either::Left(ordered_entries.iter().copied()))
+        } else if let Some(chunk_group) = chunk_group_info.chunk_groups.get_index(idx) {
+            Either::Right(chunk_group.entries())
+        } else {
+            Either::Left(Either::Right(std::iter::empty()))
+        }
+    }
+
     pub fn get_batch_group(
         &self,
         module_or_batch: &ModuleOrBatch,
@@ -323,6 +353,8 @@ pub async fn compute_module_batches(
         let mut pre_batches = PreBatches::new();
         let mut queue: VecDeque<(ResolvedVc<Box<dyn Module>>, PreBatchIndex)> = VecDeque::new();
 
+        let mut chunk_group_indicies_with_merged_children = FxHashSet::default();
+
         // Start with the entries
         for chunk_group in &chunk_group_info.chunk_groups {
             for entry in chunk_group.entries() {
@@ -340,6 +372,9 @@ pub async fn compute_module_batches(
                     pre_batches.single_module_entries.insert(entry);
                 }
             }
+            if let Some(parent) = chunk_group.get_merged_parent() {
+                chunk_group_indicies_with_merged_children.insert(parent);
+            }
         }
 
         let mut initial_pre_batch_items = 0;
@@ -360,6 +395,83 @@ pub async fn compute_module_batches(
         span.record("initial_pre_batch_items", initial_pre_batch_items);
         span.record("initial_pre_batches", pre_batches.batches.len());
 
+        // Figure out the order of all merged groups
+        let mut ordered_entries: Vec<Option<EntriesList>> =
+            vec![None; chunk_group_info.chunk_groups.len()];
+        for (i, chunk_group) in chunk_group_info.chunk_groups.iter().enumerate() {
+            if !chunk_group_indicies_with_merged_children.contains(&i) {
+                continue;
+            }
+            let mut merged_modules: FxHashMap<ChunkingType, FxIndexSet<_>> = FxHashMap::default();
+            let mut stack = ordered_entries[i]
+                .as_ref()
+                .map_or_else(
+                    || Either::Left(chunk_group.entries()),
+                    |v| Either::Right(v.iter().copied()),
+                )
+                .filter_map(|module| {
+                    if let Some(chunkable_module) = ResolvedVc::try_downcast(module) {
+                        let idx = *pre_batches.entries.get(&chunkable_module).unwrap();
+                        Some((idx, 0))
+                    } else {
+                        None
+                    }
+                })
+                .collect::<Vec<_>>();
+            stack.reverse();
+            let mut visited = FxHashSet::default();
+            while let Some((idx, mut pos)) = stack.pop() {
+                let batch = &pre_batches.batches[idx];
+                while let Some(item) = batch.items.get_index(pos) {
+                    match item {
+                        PreBatchItem::ParallelModule(_) => {}
+                        PreBatchItem::ParallelReference(other_idx) => {
+                            if visited.insert(*other_idx) {
+                                stack.push((idx, pos + 1));
+                                stack.push((*other_idx, 0));
+                                break;
+                            }
+                        }
+                        PreBatchItem::NonParallelEdge(chunking_type, module) => {
+                            if chunking_type.is_merged() {
+                                merged_modules
+                                    .entry(chunking_type.clone())
+                                    .or_default()
+                                    .insert(*module);
+                            }
+                        }
+                    }
+                    pos += 1;
+                }
+            }
+            if !merged_modules.is_empty() {
+                for (ty, merged_modules) in merged_modules {
+                    let chunk_group_key = match ty {
+                        ChunkingType::Isolated {
+                            merge_tag: Some(merge_tag),
+                            ..
+                        } => ChunkGroupKey::IsolatedMerged {
+                            parent: i.into(),
+                            merge_tag: merge_tag.clone(),
+                        },
+                        ChunkingType::Shared {
+                            merge_tag: Some(merge_tag),
+                            ..
+                        } => ChunkGroupKey::SharedMerged {
+                            parent: i.into(),
+                            merge_tag: merge_tag.clone(),
+                        },
+                        _ => unreachable!(),
+                    };
+                    let idx = chunk_group_info
+                        .chunk_group_keys
+                        .get_index_of(&chunk_group_key)
+                        .unwrap();
+                    ordered_entries[idx] = Some(merged_modules);
+                }
+            }
+        }
+
         // Create a map of parallel module to the batches they are contained in.
         let mut parallel_module_to_pre_batch: FxIndexMap<_, Vec<PreBatchIndex>> =
             FxIndexMap::default();
@@ -748,6 +860,7 @@ pub async fn compute_module_batches(
             graph: TracedDiGraph(graph),
             entries,
             batch_groups,
+            ordered_entries,
         }
         .cell())
     }