Chunking Refactor Step 2 (vercel/turborepo#6120)

### Description

The second step in our chunking refactoring, this removes our use of
`Module::as_chunk` and `Module::as_root_chunk`. Instead, the only way to
generate a chunk is directly from a root `ChunkItem`.

Re: vercel/turborepo#6104

### Testing Instructions

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


Closes WEB-1720

---------

Co-authored-by: Tobias Koppers <[email protected]>

Author: jridgewell

crates/turbopack-build/src/chunking_context.rs

@@ -8,7 +8,10 @@ use turbo_tasks::{
 };
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
-    chunk::{Chunk, ChunkableModuleExt, ChunkingContext, Chunks, EvaluatableAssets},
+    chunk::{
+        availability_info::AvailabilityInfo, Chunk, ChunkItem, ChunkableModule, ChunkingContext,
+        Chunks, EvaluatableAssets,
+    },
     environment::Environment,
     ident::AssetIdent,
     module::Module,
@@ -163,7 +166,12 @@ impl BuildChunkingContext {
         module: Vc<Box<dyn EcmascriptChunkPlaceable>>,
         evaluatable_assets: Vc<EvaluatableAssets>,
     ) -> Result<Vc<Box<dyn OutputAsset>>> {
-        let entry_chunk = module.as_root_chunk(Vc::upcast(self));
+        let entry_chunk =
+            module
+                .as_chunk_item(Vc::upcast(self))
+                .as_chunk(Value::new(AvailabilityInfo::Root {
+                    current_availability_root: Vc::upcast(module),
+                }));
 
         let other_chunks = self
             .get_chunk_assets(entry_chunk, evaluatable_assets)
@@ -214,7 +222,10 @@ impl BuildChunkingContext {
             .map({
                 move |evaluatable_asset| async move {
                     evaluatable_asset
-                        .as_root_chunk(Vc::upcast(self))
+                        .as_chunk_item(Vc::upcast(self))
+                        .as_chunk(Value::new(AvailabilityInfo::Root {
+                            current_availability_root: Vc::upcast(*evaluatable_asset),
+                        }))
                         .resolve()
                         .await
                 }
@@ -350,8 +361,12 @@ impl ChunkingContext for BuildChunkingContext {
     #[turbo_tasks::function]
     async fn chunk_group(
         self: Vc<Self>,
-        entry_chunk: Vc<Box<dyn Chunk>>,
+        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 parallel_chunks = get_parallel_chunks([entry_chunk]).await?;
 
         let optimized_chunks = get_optimized_chunks(parallel_chunks).await?;
@@ -368,7 +383,7 @@ impl ChunkingContext for BuildChunkingContext {
     #[turbo_tasks::function]
     async fn evaluated_chunk_group(
         self: Vc<Self>,
-        _entry_chunk: Vc<Box<dyn Chunk>>,
+        _ident: Vc<AssetIdent>,
         _evaluatable_assets: Vc<EvaluatableAssets>,
     ) -> Result<Vc<OutputAssets>> {
         // TODO(alexkirsz) This method should be part of a separate trait that is

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, ChunkableModuleExt, ChunkingContext, EvaluatableAssets},
+    chunk::{ChunkableModule, ChunkingContextExt, EvaluatableAssets},
     environment::{BrowserEnvironment, Environment, ExecutionEnvironment},
     issue::{handle_issues, IssueReporter, IssueSeverity},
     module::Module,
@@ -269,7 +269,7 @@ async fn build_internal(
                 } else if let Some(chunkable) =
                     Vc::try_resolve_sidecast::<Box<dyn ChunkableModule>>(entry_module).await?
                 {
-                    chunking_context.chunk_group(chunkable.as_root_chunk(chunking_context))
+                    chunking_context.root_chunk_group(chunkable)
                 } else {
                     // TODO convert into a serve-able asset
                     bail!(

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

@@ -1,8 +1,8 @@
 use anyhow::Result;
-use turbo_tasks::{ValueToString, Vc};
+use turbo_tasks::{Upcast, Value, ValueToString, Vc};
 use turbo_tasks_fs::FileSystemPath;
 
-use super::{Chunk, EvaluatableAssets};
+use super::{availability_info::AvailabilityInfo, ChunkableModule, EvaluatableAssets};
 use crate::{
     chunk::{ChunkItem, ModuleId},
     environment::Environment,
@@ -57,11 +57,15 @@ pub trait ChunkingContext {
 
     fn with_layer(self: Vc<Self>, layer: String) -> Vc<Self>;
 
-    fn chunk_group(self: Vc<Self>, entry: Vc<Box<dyn Chunk>>) -> Vc<OutputAssets>;
+    fn chunk_group(
+        self: Vc<Self>,
+        module: Vc<Box<dyn ChunkableModule>>,
+        availability_info: Value<AvailabilityInfo>,
+    ) -> Vc<OutputAssets>;
 
     fn evaluated_chunk_group(
         self: Vc<Self>,
-        entry: Vc<Box<dyn Chunk>>,
+        ident: Vc<AssetIdent>,
         evaluatable_assets: Vc<EvaluatableAssets>,
     ) -> Vc<OutputAssets>;
 
@@ -77,3 +81,20 @@ pub trait ChunkingContext {
         Ok(ModuleId::String(ident.to_string().await?.clone_value()).cell())
     }
 }
+
+pub trait ChunkingContextExt {
+    fn root_chunk_group(self: Vc<Self>, module: Vc<Box<dyn ChunkableModule>>) -> Vc<OutputAssets>
+    where
+        Self: Send;
+}
+
+impl<T: ChunkingContext + Send + Upcast<Box<dyn ChunkingContext>>> ChunkingContextExt for T {
+    fn root_chunk_group(self: Vc<Self>, module: Vc<Box<dyn ChunkableModule>>) -> Vc<OutputAssets> {
+        self.chunk_group(
+            module,
+            Value::new(AvailabilityInfo::Root {
+                current_availability_root: Vc::upcast(module),
+            }),
+        )
+    }
+}

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

@@ -29,7 +29,7 @@ use turbo_tasks_hash::DeterministicHash;
 
 use self::availability_info::AvailabilityInfo;
 pub use self::{
-    chunking_context::ChunkingContext,
+    chunking_context::{ChunkingContext, ChunkingContextExt},
     data::{ChunkData, ChunkDataOption, ChunksData},
     evaluate::{EvaluatableAsset, EvaluatableAssetExt, EvaluatableAssets},
     passthrough_asset::PassthroughModule,
@@ -40,7 +40,6 @@ use crate::{
     module::{Module, Modules},
     output::OutputAssets,
     reference::{ModuleReference, ModuleReferences},
-    resolve::ModuleResolveResult,
 };
 
 /// A module id, which can be a number or string
@@ -91,43 +90,6 @@ pub trait ChunkableModule: Module + Asset {
     ) -> 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>> {
-        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>> {
-        let chunk_item = self.as_chunk_item(chunking_context);
-        chunk_item.as_chunk(Value::new(AvailabilityInfo::Root {
-            current_availability_root: Vc::upcast(self),
-        }))
-    }
-}
-
 #[turbo_tasks::value(transparent)]
 pub struct Chunks(Vec<Vc<Box<dyn Chunk>>>);
 
@@ -224,53 +186,6 @@ pub trait ChunkableModuleReference: ModuleReference + ValueToString {
     }
 }
 
-/// A reference to multiple chunks from a [ChunkGroup]
-#[turbo_tasks::value]
-pub struct ChunkGroupReference {
-    chunking_context: Vc<Box<dyn ChunkingContext>>,
-    entry: Vc<Box<dyn Chunk>>,
-}
-
-#[turbo_tasks::value_impl]
-impl ChunkGroupReference {
-    #[turbo_tasks::function]
-    pub fn new(
-        chunking_context: Vc<Box<dyn ChunkingContext>>,
-        entry: Vc<Box<dyn Chunk>>,
-    ) -> Vc<Self> {
-        Self::cell(ChunkGroupReference {
-            chunking_context,
-            entry,
-        })
-    }
-
-    #[turbo_tasks::function]
-    async fn chunks(self: Vc<Self>) -> Result<Vc<OutputAssets>> {
-        let this = self.await?;
-        Ok(this.chunking_context.chunk_group(this.entry))
-    }
-}
-
-#[turbo_tasks::value_impl]
-impl ModuleReference for ChunkGroupReference {
-    #[turbo_tasks::function]
-    async fn resolve_reference(self: Vc<Self>) -> Result<Vc<ModuleResolveResult>> {
-        let set = self.chunks().await?.clone_value();
-        Ok(ModuleResolveResult::output_assets(set).cell())
-    }
-}
-
-#[turbo_tasks::value_impl]
-impl ValueToString for ChunkGroupReference {
-    #[turbo_tasks::function]
-    async fn to_string(&self) -> Result<Vc<String>> {
-        Ok(Vc::cell(format!(
-            "chunk group ({})",
-            self.entry.ident().to_string().await?
-        )))
-    }
-}
-
 pub struct ChunkContentResult<I> {
     pub chunk_items: Vec<I>,
     pub chunks: Vec<Vc<Box<dyn Chunk>>>,
@@ -440,17 +355,20 @@ where
                 }
             }
             ChunkingType::Parallel => {
-                let chunk = chunkable_module.as_chunk(
-                    chunk_content_context.chunking_context,
-                    chunk_content_context.availability_info,
-                );
+                let chunk = chunkable_module
+                    .as_chunk_item(chunk_content_context.chunking_context)
+                    .as_chunk(chunk_content_context.availability_info);
                 graph_nodes.push((
                     Some((module, chunking_type)),
                     ChunkContentGraphNode::Chunk(chunk),
                 ));
             }
             ChunkingType::IsolatedParallel => {
-                let chunk = chunkable_module.as_root_chunk(chunk_content_context.chunking_context);
+                let chunk = chunkable_module
+                    .as_chunk_item(chunk_content_context.chunking_context)
+                    .as_chunk(Value::new(AvailabilityInfo::Root {
+                        current_availability_root: Vc::upcast(chunkable_module),
+                    }));
                 graph_nodes.push((
                     Some((module, chunking_type)),
                     ChunkContentGraphNode::Chunk(chunk),
@@ -479,10 +397,9 @@ where
                     }
                 }
 
-                let chunk = chunkable_module.as_chunk(
-                    chunk_content_context.chunking_context,
-                    chunk_content_context.availability_info,
-                );
+                let chunk = chunkable_module
+                    .as_chunk_item(chunk_content_context.chunking_context)
+                    .as_chunk(chunk_content_context.availability_info);
                 graph_nodes.push((
                     Some((module, chunking_type)),
                     ChunkContentGraphNode::Chunk(chunk),

crates/turbopack-dev-server/src/html.rs

@@ -5,8 +5,9 @@ use turbo_tasks_fs::{File, FileSystemPath};
 use turbo_tasks_hash::{encode_hex, Xxh3Hash64Hasher};
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    chunk::{ChunkableModule, ChunkableModuleExt, ChunkingContext, EvaluatableAssets},
+    chunk::{ChunkableModule, ChunkingContext, ChunkingContextExt, EvaluatableAssets},
     ident::AssetIdent,
+    module::Module,
     output::{OutputAsset, OutputAssets},
     version::{Version, VersionedContent},
 };
@@ -131,11 +132,18 @@ impl DevHtmlAsset {
             .map(|entry| async move {
                 let (chunkable_module, chunking_context, runtime_entries) = entry;
 
-                let chunk = chunkable_module.as_root_chunk(*chunking_context);
                 let assets = if let Some(runtime_entries) = runtime_entries {
-                    chunking_context.evaluated_chunk_group(chunk, *runtime_entries)
+                    let runtime_entries = if let Some(evaluatable) =
+                        Vc::try_resolve_downcast(*chunkable_module).await?
+                    {
+                        runtime_entries.with_entry(evaluatable)
+                    } else {
+                        *runtime_entries
+                    };
+                    chunking_context
+                        .evaluated_chunk_group(chunkable_module.ident(), runtime_entries)
                 } else {
-                    chunking_context.chunk_group(chunk)
+                    chunking_context.root_chunk_group(Vc::upcast(*chunkable_module))
                 };
 
                 assets.await