diff --git a/turbopack/crates/turbopack-core/src/chunk/chunk_group.rs b/turbopack/crates/turbopack-core/src/chunk/chunk_group.rs index f5d9e669d7e060..01fa0293654839 100644 --- a/turbopack/crates/turbopack-core/src/chunk/chunk_group.rs +++ b/turbopack/crates/turbopack-core/src/chunk/chunk_group.rs @@ -1,4 +1,4 @@ -use std::{cell::RefCell, collections::HashSet, sync::atomic::AtomicBool}; +use std::{collections::HashSet, sync::atomic::AtomicBool}; use anyhow::{Context, Result}; use rustc_hash::FxHashMap; @@ -386,13 +386,16 @@ pub async fn chunk_group_content( .into_iter() .map(async |chunkable_module| match chunkable_module { ChunkableModuleOrBatch::Module(module) => { - if !merged_modules_ref.should_create_chunk_item_for(ResolvedVc::upcast(module)) + if !merged_modules_ref + .should_create_chunk_item_for(ResolvedVc::upcast(module)) + .await? { return Ok(None); } - let module = if let Some(replacement) = - merged_modules_ref.should_replace_module(ResolvedVc::upcast(module)) + let module = if let Some(replacement) = merged_modules_ref + .should_replace_module(ResolvedVc::upcast(module)) + .await? { replacement } else { @@ -448,28 +451,34 @@ async fn map_module_batch( let merged_modules = merged_modules.await?; let batch_ref = batch.await?; - let modified = RefCell::new(false); + let modified = AtomicBool::new(false); let modules = batch_ref .modules .iter() - .flat_map(|&module| { - if !merged_modules.should_create_chunk_item_for(ResolvedVc::upcast(module)) { - *modified.borrow_mut() = true; - return None; + .copied() + .map(async |module| { + if !merged_modules + .should_create_chunk_item_for(ResolvedVc::upcast(module)) + .await? + { + modified.store(true, std::sync::atomic::Ordering::Relaxed); + return Ok(None); } - let module = if let Some(replacement) = - merged_modules.should_replace_module(ResolvedVc::upcast(module)) + let module = if let Some(replacement) = merged_modules + .should_replace_module(ResolvedVc::upcast(module)) + .await? { - *modified.borrow_mut() = true; + modified.store(true, std::sync::atomic::Ordering::Relaxed); replacement } else { module }; - Some(module) + Ok(Some(module)) }) - .collect::>(); + .try_flat_join() + .await?; if modified.into_inner() { Ok(ModuleBatch::new( @@ -496,18 +505,22 @@ async fn map_module_batch_group( .copied() .map(async |chunkable_module| match chunkable_module { ModuleOrBatch::Module(module) => { - if !merged_modules_ref.should_create_chunk_item_for(module) { + if !merged_modules_ref + .should_create_chunk_item_for(module) + .await? + { modified.store(true, std::sync::atomic::Ordering::Relaxed); return Ok(None); } - let module = - if let Some(replacement) = merged_modules_ref.should_replace_module(module) { - modified.store(true, std::sync::atomic::Ordering::Relaxed); - ResolvedVc::upcast(replacement) - } else { - module - }; + let module = if let Some(replacement) = + merged_modules_ref.should_replace_module(module).await? + { + modified.store(true, std::sync::atomic::Ordering::Relaxed); + ResolvedVc::upcast(replacement) + } else { + module + }; Ok(Some(ModuleOrBatch::Module(module))) } diff --git a/turbopack/crates/turbopack-core/src/chunk/chunking/production.rs b/turbopack/crates/turbopack-core/src/chunk/chunking/production.rs index 1da33069ad583b..29669c0d3a09d9 100644 --- a/turbopack/crates/turbopack-core/src/chunk/chunking/production.rs +++ b/turbopack/crates/turbopack-core/src/chunk/chunking/production.rs @@ -76,6 +76,7 @@ pub async fn make_production_chunks( // lookup using the original module. let original_module = merged_modules .get_original_module(ResolvedVc::upcast(module)) + .await? .context("every module should have a chunk group")?; module_chunk_groups .get(&original_module) diff --git a/turbopack/crates/turbopack-core/src/module_graph/merged_modules.rs b/turbopack/crates/turbopack-core/src/module_graph/merged_modules.rs index 9c8d25e489bb26..69e5c541f5b8a3 100644 --- a/turbopack/crates/turbopack-core/src/module_graph/merged_modules.rs +++ b/turbopack/crates/turbopack-core/src/module_graph/merged_modules.rs @@ -17,42 +17,63 @@ use crate::{ resolve::ExportUsage, }; +#[turbo_tasks::value(transparent, cell = "keyed")] +#[allow(clippy::type_complexity)] +pub struct MergedModulesReplacements( + FxHashMap>, ResolvedVc>>, +); + +#[turbo_tasks::value(transparent, cell = "keyed")] +#[allow(clippy::type_complexity)] +pub struct MergedModulesOriginalModules( + FxHashMap>, ResolvedVc>>, +); + +#[turbo_tasks::value(transparent, cell = "keyed")] +#[allow(clippy::type_complexity)] +pub struct MergedModulesIncluded(FxHashSet>>); + #[turbo_tasks::value] pub struct MergedModuleInfo { /// A map of modules to the merged module containing the module plus additional modules. - #[allow(clippy::type_complexity)] - pub replacements: FxHashMap>, ResolvedVc>>, + pub replacements: ResolvedVc, /// A map of replacement modules to their corresponding chunk group info (which is the same as /// the chunk group info of the original module it replaced). - #[allow(clippy::type_complexity)] - pub replacements_to_original: - FxHashMap>, ResolvedVc>>, + pub replacements_to_original: ResolvedVc, /// A map of modules that are already contained as values in replacements. - pub included: FxHashSet>>, + pub included: ResolvedVc, } impl MergedModuleInfo { /// Whether the given module should be replaced with a merged module. - pub fn should_replace_module( + pub async fn should_replace_module( &self, module: ResolvedVc>, - ) -> Option>> { - self.replacements.get(&module).copied() + ) -> Result>>> { + Ok(self.replacements.get(&module).await?.as_deref().copied()) } /// Returns the original module for the given replacement module (useful for retrieving the /// chunk group info). - pub fn get_original_module( + pub async fn get_original_module( &self, module: ResolvedVc>, - ) -> Option>> { - self.replacements_to_original.get(&module).copied() + ) -> Result>>> { + Ok(self + .replacements_to_original + .get(&module) + .await? + .as_deref() + .copied()) } // Whether the given module should be skipped during chunking, as it is already included in a // module returned by some `should_replace_module` call. - pub fn should_create_chunk_item_for(&self, module: ResolvedVc>) -> bool { - !self.included.contains(&module) + pub async fn should_create_chunk_item_for( + &self, + module: ResolvedVc>, + ) -> Result { + Ok(!self.included.contains_key(&module).await?) } } @@ -758,9 +779,9 @@ pub async fn compute_merged_modules(module_graph: Vc) -> Result