Walk chunkable references instead

vercel · Dec 3, 2024 · ed32c69 · ed32c69
1 parent 30ec84b
commit ed32c69
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 11 deletions.
diff --git a/crates/next-api/src/module_graph.rs b/crates/next-api/src/module_graph.rs
@@ -14,10 +14,11 @@ use turbo_tasks::{
     CollectiblesSource, FxIndexMap, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, Vc,
 };
 use turbopack_core::{
+    chunk::ChunkingType,
     context::AssetContext,
     issue::Issue,
     module::{Module, Modules},
-    reference::primary_referenced_modules,
+    reference::primary_chunkable_referenced_modules,
 };
 
 use crate::{
@@ -83,11 +84,12 @@ impl SingleModuleGraph {
                 graph.add_edge(parent_idx, idx, ());
             }
 
-            for reference in primary_referenced_modules(*module).await?.iter() {
-                if reference.ident().path().await?.extension_ref() == Some("map") {
-                    continue;
+            for (ty, references) in primary_chunkable_referenced_modules(*module).await?.iter() {
+                if !matches!(ty, ChunkingType::Traced) {
+                    for reference in references {
+                        stack.push((Some(idx), *reference));
+                    }
                 }
-                stack.push((Some(idx), *reference));
             }
         }
 

diff --git a/...c/next_client_reference/ecmascript_client_reference/ecmascript_client_reference_module.rs b/...c/next_client_reference/ecmascript_client_reference/ecmascript_client_reference_module.rs
@@ -97,7 +97,7 @@ impl Asset for EcmascriptClientReferenceModule {
 }
 
 #[turbo_tasks::value]
-struct EcmascriptClientReference {
+pub(crate) struct EcmascriptClientReference {
     module: ResolvedVc<Box<dyn Module>>,
     ty: ChunkGroupType,
     description: ResolvedVc<RcStr>,

diff --git a/..._client_reference/ecmascript_client_reference/ecmascript_client_reference_proxy_module.rs b/..._client_reference/ecmascript_client_reference/ecmascript_client_reference_proxy_module.rs
@@ -7,12 +7,14 @@ use turbo_tasks::{ResolvedVc, Value, ValueToString, Vc};
 use turbo_tasks_fs::File;
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    chunk::{AsyncModuleInfo, ChunkItem, ChunkType, ChunkableModule, ChunkingContext},
+    chunk::{
+        AsyncModuleInfo, ChunkGroupType, ChunkItem, ChunkType, ChunkableModule, ChunkingContext,
+    },
     code_builder::CodeBuilder,
     context::AssetContext,
     ident::AssetIdent,
     module::Module,
-    reference::{ModuleReferences, SingleModuleReference},
+    reference::ModuleReferences,
     reference_type::ReferenceType,
     virtual_source::VirtualSource,
 };
@@ -24,7 +26,10 @@ use turbopack_ecmascript::{
     utils::StringifyJs,
 };
 
-use super::ecmascript_client_reference_module::EcmascriptClientReferenceModule;
+use crate::next_client_reference::{
+    ecmascript_client_reference::ecmascript_client_reference_module::EcmascriptClientReference,
+    EcmascriptClientReferenceModule,
+};
 
 /// A [`EcmascriptClientReferenceProxyModule`] is used in RSC to represent
 /// a client or SSR asset.
@@ -186,12 +191,17 @@ impl Module for EcmascriptClientReferenceProxyModule {
             .await?
             .iter()
             .copied()
-            .chain(once(Vc::upcast(SingleModuleReference::new(
+            // TODO this will break once ChunkingType::Isolated is properly implemented.
+            //
+            // We should instead merge EcmascriptClientReferenceProxyModule and
+            // EcmascriptClientReferenceModule into a single module
+            .chain(once(Vc::upcast(EcmascriptClientReference::new(
                 Vc::upcast(EcmascriptClientReferenceModule::new(
                     **server_module_ident,
                     **client_module,
                     **ssr_module,
                 )),
+                ChunkGroupType::Entry,
                 client_reference_description(),
             ))))
             .collect();

diff --git a/turbopack/crates/turbopack-core/src/reference/mod.rs b/turbopack/crates/turbopack-core/src/reference/mod.rs
@@ -4,7 +4,7 @@ use anyhow::Result;
 use turbo_rcstr::RcStr;
 use turbo_tasks::{
     graph::{AdjacencyMap, GraphTraversal},
-    FxIndexSet, ResolvedVc, TryJoinIterExt, ValueToString, Vc,
+    FxIndexSet, ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, ValueToString, Vc,
 };
 
 use crate::{
@@ -267,6 +267,45 @@ pub async fn primary_referenced_modules(module: Vc<Box<dyn Module>>) -> Result<V
     Ok(Vc::cell(modules))
 }
 
+type ModulesVec = Vec<ResolvedVc<Box<dyn Module>>>;
+#[turbo_tasks::value(transparent)]
+pub struct ModulesWithChunkingType(Vec<(ChunkingType, ModulesVec)>);
+
+/// Aggregates all primary [Module]s referenced by an [Module] via [ChunkableModuleReference]s.
+/// This does not include transitively references [Module]s, only includes
+/// primary [Module]s referenced.
+///
+/// [Module]: crate::module::Module
+#[turbo_tasks::function]
+pub async fn primary_chunkable_referenced_modules(
+    module: Vc<Box<dyn Module>>,
+) -> Result<Vc<ModulesWithChunkingType>> {
+    let modules = module
+        .references()
+        .await?
+        .iter()
+        .map(|reference| async {
+            if let Some(reference) =
+                Vc::try_resolve_downcast::<Box<dyn ChunkableModuleReference>>(*reference).await?
+            {
+                if let Some(chunking_type) = &*reference.chunking_type().await? {
+                    let resolved = reference
+                        .resolve_reference()
+                        .resolve()
+                        .await?
+                        .primary_modules()
+                        .await?
+                        .clone_value();
+                    return Ok(Some((chunking_type.clone(), resolved)));
+                }
+            }
+            Ok(None)
+        })
+        .try_flat_join()
+        .await?;
+    Ok(Vc::cell(modules))
+}
+
 /// Walks the asset graph from multiple assets and collect all referenced
 /// assets.
 #[turbo_tasks::function]