Turbopack: use selective reads for merged modules

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::<Vec<_>>();
+        .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)))
             }

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)

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<Box<dyn Module>>, ResolvedVc<Box<dyn ChunkableModule>>>,
+);
+
+#[turbo_tasks::value(transparent, cell = "keyed")]
+#[allow(clippy::type_complexity)]
+pub struct MergedModulesOriginalModules(
+    FxHashMap<ResolvedVc<Box<dyn Module>>, ResolvedVc<Box<dyn Module>>>,
+);
+
+#[turbo_tasks::value(transparent, cell = "keyed")]
+#[allow(clippy::type_complexity)]
+pub struct MergedModulesIncluded(FxHashSet<ResolvedVc<Box<dyn Module>>>);
+
 #[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<Box<dyn Module>>, ResolvedVc<Box<dyn ChunkableModule>>>,
+    pub replacements: ResolvedVc<MergedModulesReplacements>,
     /// 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<Box<dyn Module>>, ResolvedVc<Box<dyn Module>>>,
+    pub replacements_to_original: ResolvedVc<MergedModulesOriginalModules>,
     /// A map of modules that are already contained as values in replacements.
-    pub included: FxHashSet<ResolvedVc<Box<dyn Module>>>,
+    pub included: ResolvedVc<MergedModulesIncluded>,
 }
 
 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<Box<dyn Module>>,
-    ) -> Option<ResolvedVc<Box<dyn ChunkableModule>>> {
-        self.replacements.get(&module).copied()
+    ) -> Result<Option<ResolvedVc<Box<dyn ChunkableModule>>>> {
+        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<Box<dyn Module>>,
-    ) -> Option<ResolvedVc<Box<dyn Module>>> {
-        self.replacements_to_original.get(&module).copied()
+    ) -> Result<Option<ResolvedVc<Box<dyn Module>>>> {
+        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<Box<dyn Module>>) -> bool {
-        !self.included.contains(&module)
+    pub async fn should_create_chunk_item_for(
+        &self,
+        module: ResolvedVc<Box<dyn Module>>,
+    ) -> Result<bool> {
+        Ok(!self.included.contains_key(&module).await?)
     }
 }
 
@@ -758,9 +779,9 @@ pub async fn compute_merged_modules(module_graph: Vc<ModuleGraph>) -> Result<Vc<
         span.record("included_modules", included.len());
 
         Ok(MergedModuleInfo {
-            replacements,
-            replacements_to_original,
-            included,
+            replacements: ResolvedVc::cell(replacements),
+            replacements_to_original: ResolvedVc::cell(replacements_to_original),
+            included: ResolvedVc::cell(included),
         }
         .cell())
     }