Extract as_chunk into shared ChunkType trait (#6123)

### Description Step 3 in our chunking refactors extracts `ChunkItem::as_chunk` into a new `ChunkType` trait. This new trait isn't useful yet, but will eventually allow us to collect `ChunkItems` of a similar type into a lists, and perform chunking over all similar items at once. Re: #6120 ### Testing Instructions Closes WEB-1722
vercel · Oct 6, 2023 · 8439948 · 8439948
1 parent 624521f
commit 8439948
Show file tree

Hide file tree

Showing 19 changed files with 284 additions and 182 deletions.
diff --git a/crates/turbopack-build/src/chunking_context.rs b/crates/turbopack-build/src/chunking_context.rs
@@ -9,8 +9,8 @@ use turbo_tasks::{
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
     chunk::{
-        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkableModule, ChunkingContext,
-        Chunks, EvaluatableAssets,
+        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkType, ChunkableModule,
+        ChunkingContext, Chunks, EvaluatableAssets,
     },
     environment::Environment,
     ident::AssetIdent,
@@ -166,12 +166,13 @@ impl BuildChunkingContext {
         module: Vc<Box<dyn EcmascriptChunkPlaceable>>,
         evaluatable_assets: Vc<EvaluatableAssets>,
     ) -> Result<Vc<Box<dyn OutputAsset>>> {
-        let entry_chunk =
-            module
-                .as_chunk_item(Vc::upcast(self))
-                .as_chunk(Value::new(AvailabilityInfo::Root {
-                    current_availability_root: Vc::upcast(module),
-                }));
+        let entry_chunk_item = module.as_chunk_item(Vc::upcast(self));
+        let entry_chunk = entry_chunk_item.ty().as_chunk(
+            entry_chunk_item,
+            Value::new(AvailabilityInfo::Root {
+                current_availability_root: Vc::upcast(module),
+            }),
+        );
 
         let other_chunks = self
             .get_chunk_assets(entry_chunk, evaluatable_assets)
@@ -221,11 +222,15 @@ impl BuildChunkingContext {
             .iter()
             .map({
                 move |evaluatable_asset| async move {
-                    evaluatable_asset
-                        .as_chunk_item(Vc::upcast(self))
-                        .as_chunk(Value::new(AvailabilityInfo::Root {
-                            current_availability_root: Vc::upcast(*evaluatable_asset),
-                        }))
+                    let chunk_item = evaluatable_asset.as_chunk_item(Vc::upcast(self));
+                    chunk_item
+                        .ty()
+                        .as_chunk(
+                            chunk_item,
+                            Value::new(AvailabilityInfo::Root {
+                                current_availability_root: Vc::upcast(*evaluatable_asset),
+                            }),
+                        )
                         .resolve()
                         .await
                 }
@@ -364,9 +369,10 @@ impl ChunkingContext for BuildChunkingContext {
         module: Vc<Box<dyn ChunkableModule>>,
         availability_info: Value<AvailabilityInfo>,
     ) -> Result<Vc<OutputAssets>> {
-        let entry_chunk = module
-            .as_chunk_item(Vc::upcast(self))
-            .as_chunk(availability_info);
+        let entry_chunk_item = module.as_chunk_item(Vc::upcast(self));
+        let entry_chunk = entry_chunk_item
+            .ty()
+            .as_chunk(entry_chunk_item, availability_info);
         let parallel_chunks = get_parallel_chunks([entry_chunk]).await?;
 
         let optimized_chunks = get_optimized_chunks(parallel_chunks).await?;

diff --git a/crates/turbopack-core/src/chunk/mod.rs b/crates/turbopack-core/src/chunk/mod.rs
@@ -355,20 +355,25 @@ where
                 }
             }
             ChunkingType::Parallel => {
-                let chunk = chunkable_module
-                    .as_chunk_item(chunk_content_context.chunking_context)
-                    .as_chunk(chunk_content_context.availability_info);
+                let chunk_item =
+                    chunkable_module.as_chunk_item(chunk_content_context.chunking_context);
+                let chunk = chunk_item
+                    .ty()
+                    .as_chunk(chunk_item, chunk_content_context.availability_info);
                 graph_nodes.push((
                     Some((module, chunking_type)),
                     ChunkContentGraphNode::Chunk(chunk),
                 ));
             }
             ChunkingType::IsolatedParallel => {
-                let chunk = chunkable_module
-                    .as_chunk_item(chunk_content_context.chunking_context)
-                    .as_chunk(Value::new(AvailabilityInfo::Root {
+                let chunk_item =
+                    chunkable_module.as_chunk_item(chunk_content_context.chunking_context);
+                let chunk = chunk_item.ty().as_chunk(
+                    chunk_item,
+                    Value::new(AvailabilityInfo::Root {
                         current_availability_root: Vc::upcast(chunkable_module),
-                    }));
+                    }),
+                );
                 graph_nodes.push((
                     Some((module, chunking_type)),
                     ChunkContentGraphNode::Chunk(chunk),
@@ -397,9 +402,11 @@ where
                     }
                 }
 
-                let chunk = chunkable_module
-                    .as_chunk_item(chunk_content_context.chunking_context)
-                    .as_chunk(chunk_content_context.availability_info);
+                let chunk_item =
+                    chunkable_module.as_chunk_item(chunk_content_context.chunking_context);
+                let chunk = chunk_item
+                    .ty()
+                    .as_chunk(chunk_item, chunk_content_context.availability_info);
                 graph_nodes.push((
                     Some((module, chunking_type)),
                     ChunkContentGraphNode::Chunk(chunk),
@@ -613,8 +620,6 @@ where
 
 #[turbo_tasks::value_trait]
 pub trait ChunkItem {
-    fn as_chunk(self: Vc<Self>, availability_info: Value<AvailabilityInfo>) -> Vc<Box<dyn Chunk>>;
-
     /// The [AssetIdent] of the [Module] that this [ChunkItem] was created from.
     /// For most chunk types this must uniquely identify the asset as it's the
     /// source of the module id used at runtime.
@@ -625,9 +630,26 @@ pub trait ChunkItem {
     /// references.
     fn references(self: Vc<Self>) -> Vc<ModuleReferences>;
 
+    /// The type of chunk this item should be assembled into.
+    fn ty(self: Vc<Self>) -> Vc<Box<dyn ChunkType>>;
+
+    /// A temporary method to retrieve the module associated with this
+    /// ChunkItem. TODO: Remove this as part of the chunk refactoring.
+    fn module(self: Vc<Self>) -> Vc<Box<dyn Module>>;
+
     fn chunking_context(self: Vc<Self>) -> Vc<Box<dyn ChunkingContext>>;
 }
 
+#[turbo_tasks::value_trait]
+pub trait ChunkType {
+    /// Create a new chunk for the given subgraph.
+    fn as_chunk(
+        &self,
+        chunk_item: Vc<Box<dyn ChunkItem>>,
+        availability_info: Value<AvailabilityInfo>,
+    ) -> Vc<Box<dyn Chunk>>;
+}
+
 #[turbo_tasks::value(transparent)]
 pub struct ChunkItems(Vec<Vc<Box<dyn ChunkItem>>>);
 

diff --git a/crates/turbopack-css/src/asset.rs b/crates/turbopack-css/src/asset.rs
@@ -7,13 +7,11 @@ use swc_core::{
         visit::{VisitMutWith, VisitMutWithPath},
     },
 };
-use turbo_tasks::{TryJoinIterExt, Value, ValueToString, Vc};
+use turbo_tasks::{TryJoinIterExt, ValueToString, Vc};
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    chunk::{
-        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkableModule, ChunkingContext,
-    },
+    chunk::{ChunkItem, ChunkType, ChunkableModule, ChunkingContext},
     context::AssetContext,
     ident::AssetIdent,
     module::Module,
@@ -23,7 +21,7 @@ use turbopack_core::{
 };
 
 use crate::{
-    chunk::{CssChunk, CssChunkItem, CssChunkItemContent, CssChunkPlaceable, CssImport},
+    chunk::{CssChunkItem, CssChunkItemContent, CssChunkPlaceable, CssChunkType, CssImport},
     code_gen::CodeGenerateable,
     parse::{parse_css, ParseCss, ParseCssResult, ParseCssResultSourceMap},
     path_visitor::ApplyVisitors,
@@ -163,12 +161,13 @@ impl ChunkItem for CssModuleChunkItem {
     }
 
     #[turbo_tasks::function]
-    fn as_chunk(&self, availability_info: Value<AvailabilityInfo>) -> Vc<Box<dyn Chunk>> {
-        Vc::upcast(CssChunk::new(
-            self.chunking_context,
-            Vc::upcast(self.module),
-            availability_info,
-        ))
+    fn ty(&self) -> Vc<Box<dyn ChunkType>> {
+        Vc::upcast(Vc::<CssChunkType>::default())
+    }
+
+    #[turbo_tasks::function]
+    fn module(&self) -> Vc<Box<dyn Module>> {
+        Vc::upcast(self.module)
     }
 }
 

diff --git a/crates/turbopack-css/src/chunk/mod.rs b/crates/turbopack-css/src/chunk/mod.rs
@@ -4,16 +4,16 @@ pub(crate) mod writer;
 
 use std::fmt::Write;
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 use indexmap::IndexSet;
-use turbo_tasks::{TryJoinIterExt, Value, ValueToString, Vc};
+use turbo_tasks::{TryJoinIterExt, Value, ValueDefault, ValueToString, Vc};
 use turbo_tasks_fs::{rope::Rope, File, FileSystemPathOption};
 use turbopack_core::{
     asset::{Asset, AssetContent},
     chunk::{
         availability_info::AvailabilityInfo, chunk_content, chunk_content_split, Chunk,
-        ChunkContentResult, ChunkItem, ChunkItemExt, ChunkableModule, ChunkingContext, Chunks,
-        FromChunkableModule, ModuleId, OutputChunk, OutputChunkRuntimeInfo,
+        ChunkContentResult, ChunkItem, ChunkItemExt, ChunkType, ChunkableModule, ChunkingContext,
+        Chunks, FromChunkableModule, ModuleId, OutputChunk, OutputChunkRuntimeInfo,
     },
     code_builder::{Code, CodeBuilder},
     ident::AssetIdent,
@@ -614,3 +614,34 @@ impl Introspectable for CssChunk {
         Ok(Vc::cell(children))
     }
 }
+
+#[derive(Default)]
+#[turbo_tasks::value]
+pub struct CssChunkType {}
+
+#[turbo_tasks::value_impl]
+impl ChunkType for CssChunkType {
+    #[turbo_tasks::function]
+    async fn as_chunk(
+        &self,
+        chunk_item: Vc<Box<dyn ChunkItem>>,
+        availability_info: Value<AvailabilityInfo>,
+    ) -> Result<Vc<Box<dyn Chunk>>> {
+        let placeable = Vc::try_resolve_sidecast::<Box<dyn CssChunkPlaceable>>(chunk_item.module())
+            .await?
+            .context("Module must implmement CssChunkPlaceable to be used as a CSS Chunk")?;
+        Ok(Vc::upcast(CssChunk::new(
+            chunk_item.chunking_context(),
+            placeable,
+            availability_info,
+        )))
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ValueDefault for CssChunkType {
+    #[turbo_tasks::function]
+    fn value_default() -> Vc<Self> {
+        Self::default().cell()
+    }
+}
diff --git a/crates/turbopack-css/src/module_asset.rs b/crates/turbopack-css/src/module_asset.rs
@@ -11,10 +11,7 @@ use turbo_tasks::{Value, ValueToString, Vc};
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    chunk::{
-        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkItemExt, ChunkableModule,
-        ChunkingContext,
-    },
+    chunk::{ChunkItem, ChunkItemExt, ChunkType, ChunkableModule, ChunkingContext},
     context::AssetContext,
     ident::AssetIdent,
     issue::{Issue, IssueExt, IssueSeverity},
@@ -26,8 +23,8 @@ use turbopack_core::{
 };
 use turbopack_ecmascript::{
     chunk::{
-        EcmascriptChunk, EcmascriptChunkItem, EcmascriptChunkItemContent, EcmascriptChunkPlaceable,
-        EcmascriptChunkingContext, EcmascriptExports,
+        EcmascriptChunkItem, EcmascriptChunkItemContent, EcmascriptChunkPlaceable,
+        EcmascriptChunkType, EcmascriptChunkingContext, EcmascriptExports,
     },
     utils::StringifyJs,
     ParseResultSourceMap,
@@ -277,12 +274,13 @@ impl ChunkItem for ModuleChunkItem {
     }
 
     #[turbo_tasks::function]
-    fn as_chunk(&self, availability_info: Value<AvailabilityInfo>) -> Vc<Box<dyn Chunk>> {
-        Vc::upcast(EcmascriptChunk::new(
-            Vc::upcast(self.chunking_context),
-            Vc::upcast(self.module),
-            availability_info,
-        ))
+    fn ty(&self) -> Vc<Box<dyn ChunkType>> {
+        Vc::upcast(Vc::<EcmascriptChunkType>::default())
+    }
+
+    #[turbo_tasks::function]
+    fn module(&self) -> Vc<Box<dyn Module>> {
+        Vc::upcast(self.module)
     }
 }
 

diff --git a/crates/turbopack-dev/src/chunking_context.rs b/crates/turbopack-dev/src/chunking_context.rs
@@ -7,8 +7,8 @@ use turbo_tasks::{
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
     chunk::{
-        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkableModule, ChunkingContext,
-        Chunks, EvaluatableAssets,
+        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkType, ChunkableModule,
+        ChunkingContext, Chunks, EvaluatableAssets,
     },
     environment::Environment,
     ident::AssetIdent,
@@ -346,9 +346,10 @@ impl ChunkingContext for DevChunkingContext {
         module: Vc<Box<dyn ChunkableModule>>,
         availability_info: Value<AvailabilityInfo>,
     ) -> Result<Vc<OutputAssets>> {
-        let entry_chunk = module
-            .as_chunk_item(Vc::upcast(self))
-            .as_chunk(availability_info);
+        let entry_chunk_item = module.as_chunk_item(Vc::upcast(self));
+        let entry_chunk = entry_chunk_item
+            .ty()
+            .as_chunk(entry_chunk_item, availability_info);
         let parallel_chunks = get_parallel_chunks([entry_chunk]).await?;
 
         let optimized_chunks = get_optimized_chunks(parallel_chunks).await?;
@@ -381,11 +382,15 @@ impl ChunkingContext for DevChunkingContext {
             .iter()
             .map({
                 move |evaluatable_asset| async move {
-                    evaluatable_asset
-                        .as_chunk_item(Vc::upcast(self))
-                        .as_chunk(Value::new(AvailabilityInfo::Root {
-                            current_availability_root: Vc::upcast(*evaluatable_asset),
-                        }))
+                    let chunk_item = evaluatable_asset.as_chunk_item(Vc::upcast(self));
+                    chunk_item
+                        .ty()
+                        .as_chunk(
+                            chunk_item,
+                            Value::new(AvailabilityInfo::Root {
+                                current_availability_root: Vc::upcast(*evaluatable_asset),
+                            }),
+                        )
                         .resolve()
                         .await
                 }

diff --git a/crates/turbopack-ecmascript/src/chunk/chunk_type.rs b/crates/turbopack-ecmascript/src/chunk/chunk_type.rs
@@ -0,0 +1,40 @@
+use anyhow::{Context, Result};
+use turbo_tasks::{Value, ValueDefault, Vc};
+use turbopack_core::chunk::{availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkType};
+
+use super::{EcmascriptChunk, EcmascriptChunkPlaceable};
+
+#[derive(Default)]
+#[turbo_tasks::value]
+pub struct EcmascriptChunkType {}
+
+#[turbo_tasks::value_impl]
+impl ChunkType for EcmascriptChunkType {
+    #[turbo_tasks::function]
+    async fn as_chunk(
+        &self,
+        chunk_item: Vc<Box<dyn ChunkItem>>,
+        availability_info: Value<AvailabilityInfo>,
+    ) -> Result<Vc<Box<dyn Chunk>>> {
+        let placeable =
+            Vc::try_resolve_sidecast::<Box<dyn EcmascriptChunkPlaceable>>(chunk_item.module())
+                .await?
+                .context(
+                    "Module must implmement EcmascriptChunkPlaceable to be used as a EcmaScript \
+                     Chunk",
+                )?;
+        Ok(Vc::upcast(EcmascriptChunk::new(
+            chunk_item.chunking_context(),
+            placeable,
+            availability_info,
+        )))
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ValueDefault for EcmascriptChunkType {
+    #[turbo_tasks::function]
+    fn value_default() -> Vc<Self> {
+        Self::default().cell()
+    }
+}
diff --git a/crates/turbopack-ecmascript/src/chunk/mod.rs b/crates/turbopack-ecmascript/src/chunk/mod.rs
@@ -1,3 +1,4 @@
+pub(crate) mod chunk_type;
 pub(crate) mod content;
 pub(crate) mod context;
 pub(crate) mod data;
@@ -30,6 +31,7 @@ use turbopack_core::{
 
 use self::content::ecmascript_chunk_content;
 pub use self::{
+    chunk_type::EcmascriptChunkType,
     content::EcmascriptChunkContent,
     context::EcmascriptChunkingContext,
     data::EcmascriptChunkData,