pass inner modules to chunking to match chunking types

Author: sokra

crates/next-api/src/app.rs

@@ -53,8 +53,8 @@ use turbopack::{
 use turbopack_core::{
     asset::AssetContent,
     chunk::{
-        availability_info::AvailabilityInfo, ChunkableModule, ChunkableModules, ChunkingContext,
-        ChunkingContextExt, EvaluatableAsset, EvaluatableAssets,
+        availability_info::AvailabilityInfo, ChunkableModules, ChunkingContext, ChunkingContextExt,
+        EvaluatableAsset, EvaluatableAssets,
     },
     file_source::FileSource,
     ident::AssetIdent,
@@ -833,7 +833,7 @@ impl AppProject {
                     let graph = SingleModuleGraph::new_with_entries_visited(
                         server_utils
                             .iter()
-                            .map(|m| **m)
+                            .map(|m| Vc::upcast(**m))
                             .chain(extra_entries)
                             .collect(),
                         VisitedModules::empty(),
@@ -1663,10 +1663,7 @@ impl AppEndpoint {
                         let server_utils = client_references
                             .server_utils
                             .iter()
-                            .map(|m| async move {
-                                Ok(*ResolvedVc::try_downcast::<Box<dyn ChunkableModule>>(*m)
-                                    .context("Expected server utils to be chunkable")?)
-                            })
+                            .map(async |m| Ok(Vc::upcast(*m.await?.module)))
                             .try_join()
                             .await?;
                         let chunk_group = chunking_context
@@ -1708,7 +1705,7 @@ impl AppEndpoint {
                             let chunk_group = chunking_context
                                 .chunk_group(
                                     server_component.ident(),
-                                    *ResolvedVc::upcast(server_component),
+                                    Vc::upcast(*server_component.await?.module),
                                     module_graph,
                                     Value::new(current_availability_info),
                                 )

crates/next-core/src/next_client_reference/visit_client_reference.rs

@@ -81,7 +81,7 @@ pub struct ClientReferenceGraphResult {
     pub client_references_by_server_component:
         FxIndexMap<Option<ResolvedVc<NextServerComponentModule>>, Vec<ResolvedVc<Box<dyn Module>>>>,
     pub server_component_entries: Vec<ResolvedVc<NextServerComponentModule>>,
-    pub server_utils: Vec<ResolvedVc<Box<dyn Module>>>,
+    pub server_utils: Vec<ResolvedVc<NextServerUtilityModule>>,
     pub visited_nodes: ResolvedVc<VisitedClientReferenceGraphNodes>,
 }
 
@@ -147,7 +147,7 @@ impl ClientReferenceGraphResult {
 #[derive(Clone, Debug)]
 pub struct ServerEntries {
     pub server_component_entries: Vec<ResolvedVc<NextServerComponentModule>>,
-    pub server_utils: Vec<ResolvedVc<Box<dyn Module>>>,
+    pub server_utils: Vec<ResolvedVc<NextServerUtilityModule>>,
 }
 
 #[turbo_tasks::function]
@@ -268,7 +268,7 @@ impl VisitClientReferenceNodeState {
 enum VisitClientReferenceNodeType {
     ClientReference(ClientReference, ReadRef<RcStr>),
     ServerComponentEntry(ResolvedVc<NextServerComponentModule>, ReadRef<RcStr>),
-    ServerUtilEntry(ResolvedVc<Box<dyn Module>>, ReadRef<RcStr>),
+    ServerUtilEntry(ResolvedVc<NextServerUtilityModule>, ReadRef<RcStr>),
     Internal(ResolvedVc<Box<dyn Module>>, ReadRef<RcStr>),
 }
 
@@ -306,7 +306,9 @@ impl Visit<VisitClientReferenceNode> for VisitClientReference {
                 // nodes' edges.
                 VisitClientReferenceNodeType::ClientReference(..) => return Ok(vec![]),
                 VisitClientReferenceNodeType::Internal(module, _) => module,
-                VisitClientReferenceNodeType::ServerUtilEntry(module, _) => module,
+                VisitClientReferenceNodeType::ServerUtilEntry(module, _) => {
+                    ResolvedVc::upcast(module)
+                }
                 VisitClientReferenceNodeType::ServerComponentEntry(module, _) => {
                     ResolvedVc::upcast(module)
                 }
@@ -370,11 +372,13 @@ impl Visit<VisitClientReferenceNode> for VisitClientReference {
                         });
                     }
 
-                    if ResolvedVc::try_downcast_type::<NextServerUtilityModule>(*module).is_some() {
+                    if let Some(server_util_module) =
+                        ResolvedVc::try_downcast_type::<NextServerUtilityModule>(*module)
+                    {
                         return Ok(VisitClientReferenceNode {
                             state: VisitClientReferenceNodeState::InServerUtil,
                             ty: VisitClientReferenceNodeType::ServerUtilEntry(
-                                *module,
+                                server_util_module,
                                 module.ident().to_string().await?,
                             ),
                         });

crates/next-core/src/next_server_component/server_component_module.rs

@@ -32,7 +32,7 @@ fn modifier() -> Vc<RcStr> {
 
 #[turbo_tasks::value(shared)]
 pub struct NextServerComponentModule {
-    module: ResolvedVc<Box<dyn EcmascriptChunkPlaceable>>,
+    pub module: ResolvedVc<Box<dyn EcmascriptChunkPlaceable>>,
 }
 
 #[turbo_tasks::value_impl]

crates/next-core/src/next_server_utility/server_utility_module.rs

@@ -32,7 +32,7 @@ fn modifier() -> Vc<RcStr> {
 
 #[turbo_tasks::value(shared)]
 pub struct NextServerUtilityModule {
-    module: ResolvedVc<Box<dyn EcmascriptChunkPlaceable>>,
+    pub module: ResolvedVc<Box<dyn EcmascriptChunkPlaceable>>,
 }
 
 #[turbo_tasks::value_impl]