Chunking Refactor Step 1 (vercel/turborepo#6104)

### Description

This is the first step of refactoring the chunking. In the end we want
to avoid creating chunks from modules directly, but enforce everything
going through the ChunkingContext to be chunked. This allows us to
replace the existing chunking algorithm with a much more efficient one
that avoid duplication between chunks in first place and doesn't require
a post-chunking optimization.

This change moves a few methods around.

* `Module::as_chunk` is moved to `ChunkItem::as_chunk` as temporary
intermediate state.
* `EcmascriptPlaceable::as_chunk_item` is moved to
`ChunkableModule::as_chunk_item`. This generalizes the concept of
converting a Module into a ChunkItem

### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->


Closes WEB-1714

---------

Co-authored-by: Justin Ridgewell <justin@ridgewell.name>

crates/turbopack-build/src/chunking_context.rs

@@ -8,7 +8,7 @@ use turbo_tasks::{
 };
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
-    chunk::{Chunk, ChunkableModule, ChunkingContext, Chunks, EvaluatableAssets},
+    chunk::{Chunk, ChunkableModuleExt, ChunkingContext, Chunks, EvaluatableAssets},
     environment::Environment,
     ident::AssetIdent,
     module::Module,

crates/turbopack-build/src/ecmascript/node/content.rs

@@ -6,6 +6,7 @@ use turbo_tasks::{TryJoinIterExt, Value, Vc};
 use turbo_tasks_fs::File;
 use turbopack_core::{
     asset::AssetContent,
+    chunk::ChunkItemExt,
     code_builder::{Code, CodeBuilder},
     output::OutputAsset,
     source_map::{GenerateSourceMap, OptionSourceMap},

crates/turbopack-build/src/ecmascript/node/entry/chunk.rs

@@ -6,16 +6,13 @@ use turbo_tasks::{ValueToString, Vc};
 use turbo_tasks_fs::{File, FileSystemPath};
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    chunk::{ChunkingContext, EvaluatableAssets},
+    chunk::{ChunkItemExt, ChunkableModule, ChunkingContext, EvaluatableAssets},
     code_builder::{Code, CodeBuilder},
     ident::AssetIdent,
     output::{OutputAsset, OutputAssets},
     source_map::{GenerateSourceMap, OptionSourceMap, SourceMapAsset},
 };
-use turbopack_ecmascript::{
-    chunk::{EcmascriptChunkItemExt, EcmascriptChunkPlaceable},
-    utils::StringifyJs,
-};
+use turbopack_ecmascript::{chunk::EcmascriptChunkPlaceable, utils::StringifyJs};
 
 use super::runtime::EcmascriptBuildNodeRuntimeChunk;
 use crate::BuildChunkingContext;

crates/turbopack-cli/src/build/mod.rs

@@ -14,7 +14,7 @@ use turbopack_build::{BuildChunkingContext, MinifyType};
 use turbopack_cli_utils::issue::{ConsoleUi, LogOptions};
 use turbopack_core::{
     asset::Asset,
-    chunk::{ChunkableModule, ChunkingContext, EvaluatableAssets},
+    chunk::{ChunkableModule, ChunkableModuleExt, ChunkingContext, EvaluatableAssets},
     environment::{BrowserEnvironment, Environment, ExecutionEnvironment},
     issue::{handle_issues, IssueReporter, IssueSeverity},
     module::Module,

crates/turbopack-core/src/chunk/chunking_context.rs

@@ -1,9 +1,10 @@
 use anyhow::Result;
-use turbo_tasks::Vc;
+use turbo_tasks::{ValueToString, Vc};
 use turbo_tasks_fs::FileSystemPath;
 
 use super::{Chunk, EvaluatableAssets};
 use crate::{
+    chunk::{ChunkItem, ModuleId},
     environment::Environment,
     ident::AssetIdent,
     module::Module,
@@ -63,4 +64,16 @@ pub trait ChunkingContext {
         entry: Vc<Box<dyn Chunk>>,
         evaluatable_assets: Vc<EvaluatableAssets>,
     ) -> Vc<OutputAssets>;
+
+    async fn chunk_item_id(
+        self: Vc<Self>,
+        chunk_item: Vc<Box<dyn ChunkItem>>,
+    ) -> Result<Vc<ModuleId>> {
+        let layer = self.layer();
+        let mut ident = chunk_item.asset_ident();
+        if !layer.await?.is_empty() {
+            ident = ident.with_modifier(layer)
+        }
+        Ok(ModuleId::String(ident.to_string().await?.clone_value()).cell())
+    }
 }

crates/turbopack-core/src/chunk/mod.rs

@@ -22,7 +22,7 @@ use turbo_tasks::{
     debug::ValueDebugFormat,
     graph::{AdjacencyMap, GraphTraversal, GraphTraversalResult, Visit, VisitControlFlow},
     trace::TraceRawVcs,
-    ReadRef, TryJoinIterExt, Value, ValueToString, Vc,
+    ReadRef, TryJoinIterExt, Upcast, Value, ValueToString, Vc,
 };
 use turbo_tasks_fs::FileSystemPath;
 use turbo_tasks_hash::DeterministicHash;
@@ -85,22 +85,46 @@ pub struct ModuleIds(Vec<Vc<ModuleId>>);
 /// A [Module] that can be converted into a [Chunk].
 #[turbo_tasks::value_trait]
 pub trait ChunkableModule: Module + Asset {
+    fn as_chunk_item(
+        self: Vc<Self>,
+        chunking_context: Vc<Box<dyn ChunkingContext>>,
+    ) -> Vc<Box<dyn ChunkItem>>;
+}
+
+pub trait ChunkableModuleExt {
+    fn as_chunk(
+        self: Vc<Self>,
+        chunking_context: Vc<Box<dyn ChunkingContext>>,
+        availability_info: Value<AvailabilityInfo>,
+    ) -> Vc<Box<dyn Chunk>>
+    where
+        Self: Send;
+    fn as_root_chunk(
+        self: Vc<Self>,
+        chunking_context: Vc<Box<dyn ChunkingContext>>,
+    ) -> Vc<Box<dyn Chunk>>
+    where
+        Self: Send;
+}
+
+impl<T: ChunkableModule + Send + Upcast<Box<dyn Module>>> ChunkableModuleExt for T {
     fn as_chunk(
         self: Vc<Self>,
         chunking_context: Vc<Box<dyn ChunkingContext>>,
         availability_info: Value<AvailabilityInfo>,
-    ) -> Vc<Box<dyn Chunk>>;
+    ) -> Vc<Box<dyn Chunk>> {
+        let chunk_item = self.as_chunk_item(chunking_context);
+        chunk_item.as_chunk(availability_info)
+    }
 
     fn as_root_chunk(
         self: Vc<Self>,
         chunking_context: Vc<Box<dyn ChunkingContext>>,
     ) -> Vc<Box<dyn Chunk>> {
-        self.as_chunk(
-            chunking_context,
-            Value::new(AvailabilityInfo::Root {
-                current_availability_root: Vc::upcast(self),
-            }),
-        )
+        let chunk_item = self.as_chunk_item(chunking_context);
+        chunk_item.as_chunk(Value::new(AvailabilityInfo::Root {
+            current_availability_root: Vc::upcast(self),
+        }))
     }
 }
 
@@ -672,6 +696,8 @@ 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.
@@ -681,7 +707,25 @@ pub trait ChunkItem {
     /// TODO(alexkirsz) This should have a default impl that returns empty
     /// references.
     fn references(self: Vc<Self>) -> Vc<ModuleReferences>;
+
+    fn chunking_context(self: Vc<Self>) -> Vc<Box<dyn ChunkingContext>>;
 }
 
 #[turbo_tasks::value(transparent)]
 pub struct ChunkItems(Vec<Vc<Box<dyn ChunkItem>>>);
+
+pub trait ChunkItemExt: Send {
+    /// Returns the module id of this chunk item.
+    fn id(self: Vc<Self>) -> Vc<ModuleId>;
+}
+
+impl<T> ChunkItemExt for T
+where
+    T: Upcast<Box<dyn ChunkItem>>,
+{
+    /// Returns the module id of this chunk item.
+    fn id(self: Vc<Self>) -> Vc<ModuleId> {
+        let chunk_item = Vc::upcast(self);
+        chunk_item.chunking_context().chunk_item_id(chunk_item)
+    }
+}

crates/turbopack-css/src/asset.rs

@@ -111,34 +111,21 @@ impl Asset for CssModuleAsset {
 
 #[turbo_tasks::value_impl]
 impl ChunkableModule for CssModuleAsset {
-    #[turbo_tasks::function]
-    fn as_chunk(
-        self: Vc<Self>,
-        chunking_context: Vc<Box<dyn ChunkingContext>>,
-        availability_info: Value<AvailabilityInfo>,
-    ) -> Vc<Box<dyn Chunk>> {
-        Vc::upcast(CssChunk::new(
-            chunking_context,
-            Vc::upcast(self),
-            availability_info,
-        ))
-    }
-}
-
-#[turbo_tasks::value_impl]
-impl CssChunkPlaceable for CssModuleAsset {
     #[turbo_tasks::function]
     fn as_chunk_item(
         self: Vc<Self>,
         chunking_context: Vc<Box<dyn ChunkingContext>>,
-    ) -> Vc<Box<dyn CssChunkItem>> {
+    ) -> Vc<Box<dyn turbopack_core::chunk::ChunkItem>> {
         Vc::upcast(CssModuleChunkItem::cell(CssModuleChunkItem {
             module: self,
             chunking_context,
         }))
     }
 }
 
+#[turbo_tasks::value_impl]
+impl CssChunkPlaceable for CssModuleAsset {}
+
 #[turbo_tasks::value_impl]
 impl ResolveOrigin for CssModuleAsset {
     #[turbo_tasks::function]
@@ -169,6 +156,20 @@ impl ChunkItem for CssModuleChunkItem {
     fn references(&self) -> Vc<ModuleReferences> {
         self.module.references()
     }
+
+    #[turbo_tasks::function]
+    async fn chunking_context(&self) -> Vc<Box<dyn ChunkingContext>> {
+        Vc::upcast(self.chunking_context)
+    }
+
+    #[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,
+        ))
+    }
 }
 
 #[turbo_tasks::value_impl]
@@ -192,10 +193,12 @@ impl CssChunkItem for CssModuleChunkItem {
                     if let Some(placeable) =
                         Vc::try_resolve_downcast::<Box<dyn CssChunkPlaceable>>(module).await?
                     {
-                        imports.push(CssImport::Internal(
-                            import_ref,
-                            placeable.as_chunk_item(chunking_context),
-                        ));
+                        let item = placeable.as_chunk_item(chunking_context);
+                        if let Some(css_item) =
+                            Vc::try_resolve_downcast::<Box<dyn CssChunkItem>>(item).await?
+                        {
+                            imports.push(CssImport::Internal(import_ref, css_item));
+                        }
                     }
                 }
             } else if let Some(compose_ref) =
@@ -210,9 +213,12 @@ impl CssChunkItem for CssModuleChunkItem {
                     if let Some(placeable) =
                         Vc::try_resolve_downcast::<Box<dyn CssChunkPlaceable>>(module).await?
                     {
-                        imports.push(CssImport::Composes(
-                            placeable.as_chunk_item(chunking_context),
-                        ));
+                        let item = placeable.as_chunk_item(chunking_context);
+                        if let Some(css_item) =
+                            Vc::try_resolve_downcast::<Box<dyn CssChunkItem>>(item).await?
+                        {
+                            imports.push(CssImport::Composes(css_item));
+                        }
                     }
                 }
             }

crates/turbopack-css/src/chunk/mod.rs

@@ -12,7 +12,7 @@ use turbopack_core::{
     asset::{Asset, AssetContent},
     chunk::{
         availability_info::AvailabilityInfo, chunk_content, chunk_content_split, Chunk,
-        ChunkContentResult, ChunkItem, ChunkableModule, ChunkingContext, Chunks,
+        ChunkContentResult, ChunkItem, ChunkItemExt, ChunkableModule, ChunkingContext, Chunks,
         FromChunkableModule, ModuleId, OutputChunk, OutputChunkRuntimeInfo,
     },
     code_builder::{Code, CodeBuilder},
@@ -144,9 +144,13 @@ impl CssChunkContent {
         for entry in this.main_entries.await?.iter() {
             let entry_item = entry.as_chunk_item(this.chunking_context);
 
-            // TODO(WEB-1261)
-            for external_import in expand_imports(&mut body, entry_item).await? {
-                external_imports.insert(external_import.await?.to_owned());
+            if let Some(css_item) =
+                Vc::try_resolve_downcast::<Box<dyn CssChunkItem>>(entry_item).await?
+            {
+                // TODO(WEB-1261)
+                for external_import in expand_imports(&mut body, css_item).await? {
+                    external_imports.insert(external_import.await?.to_owned());
+                }
             }
         }
 
@@ -332,7 +336,12 @@ impl OutputChunk for CssChunk {
         let imports_chunk_items: Vec<_> = entries_chunk_items
             .iter()
             .map(|&chunk_item| async move {
-                Ok(chunk_item
+                let Some(css_item) =
+                    Vc::try_resolve_downcast::<Box<dyn CssChunkItem>>(chunk_item).await?
+                else {
+                    return Ok(vec![]);
+                };
+                Ok(css_item
                     .content()
                     .await?
                     .imports
@@ -492,13 +501,9 @@ impl CssChunkContext {
     }
 }
 
+// TODO: remove
 #[turbo_tasks::value_trait]
-pub trait CssChunkPlaceable: ChunkableModule + Module + Asset {
-    fn as_chunk_item(
-        self: Vc<Self>,
-        chunking_context: Vc<Box<dyn ChunkingContext>>,
-    ) -> Vc<Box<dyn CssChunkItem>>;
-}
+pub trait CssChunkPlaceable: ChunkableModule + Module + Asset {}
 
 #[turbo_tasks::value(transparent)]
 pub struct CssChunkPlaceables(Vec<Vc<Box<dyn CssChunkPlaceable>>>);
@@ -523,7 +528,7 @@ pub trait CssChunkItem: ChunkItem {
     fn content(self: Vc<Self>) -> Vc<CssChunkItemContent>;
     fn chunking_context(self: Vc<Self>) -> Vc<Box<dyn ChunkingContext>>;
     fn id(self: Vc<Self>) -> Vc<ModuleId> {
-        CssChunkContext::of(self.chunking_context()).chunk_item_id(self)
+        CssChunkContext::of(CssChunkItem::chunking_context(self)).chunk_item_id(self)
     }
 }
 
@@ -536,7 +541,10 @@ impl FromChunkableModule for Box<dyn CssChunkItem> {
         if let Some(placeable) =
             Vc::try_resolve_downcast::<Box<dyn CssChunkPlaceable>>(asset).await?
         {
-            return Ok(Some(placeable.as_chunk_item(chunking_context)));
+            let item = placeable.as_chunk_item(chunking_context);
+            if let Some(css_item) = Vc::try_resolve_downcast::<Box<dyn CssChunkItem>>(item).await? {
+                return Ok(Some(css_item));
+            }
         }
         Ok(None)
     }