Turbopack: Refactor output assets to allow to lazy compute output assets

Author: sokra

crates/next-api/src/app.rs

@@ -31,12 +31,12 @@ use turbo_tasks::{
 };
 use turbopack_core::{
     chunk::{
-        ChunkItem, ChunkableModule, ChunkingContext, ModuleChunkItemIdExt, ModuleId,
+        ChunkableModule, ChunkingContext, ModuleChunkItemIdExt, ModuleId,
         availability_info::AvailabilityInfo,
     },
     module::Module,
     module_graph::{ModuleGraph, SingleModuleGraph, SingleModuleGraphModuleNode},
-    output::OutputAssets,
+    output::{OutputAssetsReference, OutputAssetsWithReferenced},
 };
 
 use crate::module_graph::DynamicImportEntriesWithImporter;
@@ -72,6 +72,7 @@ pub(crate) async fn collect_next_dynamic_chunks(
                             )?,
                         )
                         .context("Client reference chunk group not found for next/dynamic import")?
+                        .await?
                         .availability_info
                 }
                 NextDynamicChunkAvailability::AvailabilityInfo(availability_info) => {
@@ -103,7 +104,7 @@ pub(crate) async fn collect_next_dynamic_chunks(
 pub struct DynamicImportedChunks(
     pub  FxIndexMap<
         ResolvedVc<NextDynamicEntryModule>,
-        (ResolvedVc<ModuleId>, ResolvedVc<OutputAssets>),
+        (ResolvedVc<ModuleId>, ResolvedVc<OutputAssetsWithReferenced>),
     >,
 );
 

crates/next-api/src/dynamic_imports.rs

@@ -1,11 +1,12 @@
 use anyhow::Result;
-use next_core::{all_assets_from_entries, next_manifests::NextFontManifest};
+use next_core::next_manifests::NextFontManifest;
 use turbo_rcstr::RcStr;
 use turbo_tasks::{ResolvedVc, Vc};
 use turbo_tasks_fs::{File, FileSystemPath};
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    output::{OutputAsset, OutputAssets},
+    output::{OutputAsset, OutputAssets, OutputAssetsReference},
+    reference::all_assets_from_entries,
 };
 
 use crate::paths::get_font_paths_from_root;
@@ -22,6 +23,9 @@ pub struct FontManifest {
     pub app_dir: bool,
 }
 
+#[turbo_tasks::value_impl]
+impl OutputAssetsReference for FontManifest {}
+
 #[turbo_tasks::value_impl]
 impl OutputAsset for FontManifest {
     #[turbo_tasks::function]

crates/next-api/src/font.rs

@@ -1,6 +1,5 @@
 use anyhow::{Result, bail};
 use next_core::{
-    all_assets_from_entries,
     next_edge::entry::wrap_edge_entry,
     next_manifests::{InstrumentationDefinition, MiddlewaresManifestV2},
 };
@@ -21,6 +20,7 @@ use turbopack_core::{
         chunk_group_info::{ChunkGroup, ChunkGroupEntry},
     },
     output::{OutputAsset, OutputAssets, OutputAssetsWithReferenced},
+    reference::all_assets_from_entries,
     reference_type::{EntryReferenceSubType, ReferenceType},
     source::Source,
     virtual_output::VirtualOutputAsset,

crates/next-api/src/instrumentation.rs

@@ -23,7 +23,7 @@ pub async fn create_react_loadable_manifest(
     let mut loadable_manifest: FxIndexMap<String, LoadableManifest> = FxIndexMap::default();
 
     for (_, (module_id, chunk_output)) in dynamic_import_entries.into_iter() {
-        let chunk_output = chunk_output.await?;
+        let chunk_output = chunk_output.assets().await?;
 
         let id = &*module_id.await?;
 

crates/next-api/src/loadable_manifest.rs

@@ -2,7 +2,6 @@ use std::future::IntoFuture;
 
 use anyhow::{Result, bail};
 use next_core::{
-    all_assets_from_entries,
     middleware::get_middleware_module,
     next_edge::entry::wrap_edge_entry,
     next_manifests::{EdgeFunctionDefinition, MiddlewaresManifestV2, ProxyMatcher, Regions},
@@ -26,6 +25,7 @@ use turbopack_core::{
         chunk_group_info::{ChunkGroup, ChunkGroupEntry},
     },
     output::{OutputAsset, OutputAssets, OutputAssetsWithReferenced},
+    reference::all_assets_from_entries,
     reference_type::{EntryReferenceSubType, ReferenceType},
     source::Source,
     virtual_output::VirtualOutputAsset,

crates/next-api/src/middleware.rs

@@ -16,7 +16,7 @@ use turbopack_core::{
     asset::{Asset, AssetContent},
     context::AssetContext,
     file_source::FileSource,
-    output::{OutputAsset, OutputAssets},
+    output::{OutputAsset, OutputAssets, OutputAssetsReference},
     reference_type::{CommonJsReferenceSubType, ReferenceType},
     resolve::{origin::PlainResolveOrigin, parse::Request},
     traced_asset::TracedAsset,
@@ -76,6 +76,9 @@ impl ServerNftJsonAsset {
     }
 }
 
+#[turbo_tasks::value_impl]
+impl OutputAssetsReference for ServerNftJsonAsset {}
+
 #[turbo_tasks::value_impl]
 impl OutputAsset for ServerNftJsonAsset {
     #[turbo_tasks::function]

crates/next-api/src/next_server_nft.rs

@@ -14,7 +14,7 @@ use turbo_tasks_fs::{
 };
 use turbopack_core::{
     asset::{Asset, AssetContent},
-    output::{OutputAsset, OutputAssets},
+    output::{OutputAsset, OutputAssets, OutputAssetsReference},
 };
 
 use crate::project::Project;
@@ -59,6 +59,9 @@ impl NftJsonAsset {
     }
 }
 
+#[turbo_tasks::value_impl]
+impl OutputAssetsReference for NftJsonAsset {}
+
 #[turbo_tasks::value_impl]
 impl OutputAsset for NftJsonAsset {
     #[turbo_tasks::function]
@@ -454,7 +457,7 @@ impl Visit<(ResolvedVc<Box<dyn OutputAsset>>, Option<ReadRef<RcStr>>)>
         node: &(ResolvedVc<Box<dyn OutputAsset>>, Option<ReadRef<RcStr>>),
     ) -> Self::EdgesFuture {
         let client_root = self.client_root.clone();
-        let exclude_glob = self.exclude_glob.clone();
+        let exclude_glob: Option<ReadRef<Glob>> = self.exclude_glob.clone();
         get_referenced_server_assets(self.emit_spans, node.0, client_root, exclude_glob)
     }
 
@@ -478,7 +481,7 @@ async fn get_referenced_server_assets(
     client_root: Option<FileSystemPath>,
     exclude_glob: Option<ReadRef<Glob>>,
 ) -> Result<Vec<(ResolvedVc<Box<dyn OutputAsset>>, Option<ReadRef<RcStr>>)>> {
-    let refs = asset.references().await?;
+    let refs = asset.references().direct_assets().await?;
 
     refs.iter()
         .map(async |asset| {

crates/next-api/src/nft_json.rs

@@ -1,8 +1,8 @@
 use anyhow::{Context, Result, bail};
 use futures::future::BoxFuture;
 use next_core::{
-    PageLoaderAsset, all_assets_from_entries, create_page_loader_entry_module,
-    get_asset_path_from_pathname, get_edge_resolve_options_context,
+    PageLoaderAsset, create_page_loader_entry_module, get_asset_path_from_pathname,
+    get_edge_resolve_options_context,
     hmr_entry::HmrEntryModule,
     mode::NextMode,
     next_client::{
@@ -56,7 +56,8 @@ use turbopack_core::{
         GraphEntries, ModuleGraph, SingleModuleGraph, VisitedModules,
         chunk_group_info::{ChunkGroup, ChunkGroupEntry},
     },
-    output::{OptionOutputAsset, OutputAsset, OutputAssets, OutputAssetsWithReferenced},
+    output::{OptionOutputAsset, OutputAsset, OutputAssets},
+    reference::all_assets_from_entries,
     reference_type::{EcmaScriptModulesReferenceSubType, EntryReferenceSubType, ReferenceType},
     resolve::{origin::PlainResolveOrigin, parse::Request, pattern::Pattern},
     source::Source,
@@ -1002,34 +1003,24 @@ impl PageEndpoint {
                 NextRuntime::Edge => edge_chunking_context,
             };
 
-            let mut current_chunks = OutputAssets::empty();
-            let mut current_referenced_assets = OutputAssets::empty();
-            let mut current_availability_info = AvailabilityInfo::Root;
+            let mut current_chunk_group = ChunkGroupResult::empty_resolved();
             for layout in [document_module, app_module].iter().flatten().copied() {
                 let span = tracing::trace_span!(
                     "layout segment",
                     name = display(layout.ident().to_string().await?)
                 );
                 async {
-                    let ChunkGroupResult {
-                        assets,
-                        referenced_assets,
-                        availability_info,
-                    } = *chunking_context
-                        .chunk_group(
-                            layout.ident(),
-                            ChunkGroup::Shared(layout),
-                            ssr_module_graph,
-                            current_availability_info,
-                        )
-                        .await?;
+                    let chunk_group = chunking_context.chunk_group(
+                        layout.ident(),
+                        ChunkGroup::Shared(layout),
+                        ssr_module_graph,
+                        current_chunk_group.await?.availability_info,
+                    );
 
-                    current_chunks = current_chunks.concatenate(*assets).resolve().await?;
-                    current_referenced_assets = current_referenced_assets
-                        .concatenate(*referenced_assets)
-                        .resolve()
+                    current_chunk_group = current_chunk_group
+                        .concatenate(chunk_group)
+                        .to_resolved()
                         .await?;
-                    current_availability_info = availability_info;
 
                     anyhow::Ok(())
                 }
@@ -1041,27 +1032,22 @@ impl PageEndpoint {
                 .context("could not process page loader entry module")?;
             let is_edge = matches!(runtime, NextRuntime::Edge);
             if is_edge {
-                let OutputAssetsWithReferenced {
-                    assets: edge_assets,
-                    referenced_assets: edge_referenced_assets,
-                } = *edge_chunking_context
-                    .evaluated_chunk_group_assets(
-                        ssr_module.ident(),
-                        ChunkGroup::Entry(vec![ResolvedVc::upcast(ssr_module_evaluatable)]),
-                        ssr_module_graph,
-                        current_availability_info,
-                    )
+                let chunk_assets = edge_chunking_context.evaluated_chunk_group_assets(
+                    ssr_module.ident(),
+                    ChunkGroup::Entry(vec![ResolvedVc::upcast(ssr_module_evaluatable)]),
+                    ssr_module_graph,
+                    current_chunk_group.await?.availability_info,
+                );
+
+                let chunk_assets = current_chunk_group
+                    .output_assets_with_referenced()
+                    .concatenate(chunk_assets)
+                    .to_resolved()
                     .await?;
 
                 Ok(SsrChunk::Edge {
-                    assets: current_chunks
-                        .concatenate(*edge_assets)
-                        .to_resolved()
-                        .await?,
-                    referenced_assets: current_referenced_assets
-                        .concatenate(*edge_referenced_assets)
-                        .to_resolved()
-                        .await?,
+                    assets: chunk_assets.assets().to_resolved().await?,
+                    referenced_assets: chunk_assets.referenced_assets().to_resolved().await?,
                     dynamic_import_entries,
                     regions: regions.clone(),
                 }
@@ -1078,9 +1064,9 @@ impl PageEndpoint {
                         ssr_entry_chunk_path,
                         EvaluatableAssets::empty().with_entry(*ssr_module_evaluatable),
                         ssr_module_graph,
-                        current_chunks,
-                        current_referenced_assets,
-                        current_availability_info,
+                        current_chunk_group.assets(),
+                        current_chunk_group.referenced_assets(),
+                        current_chunk_group.await?.availability_info,
                     )
                     .to_resolved()
                     .await?;
@@ -1320,21 +1306,14 @@ impl PageEndpoint {
 
         let ssr_chunk = match this.ty {
             PageEndpointType::Html => {
-                let client_chunk_group = self.client_chunk_group().await?;
-                let client_chunks = *client_chunk_group.assets;
-                client_assets.extend(client_chunks.await?.iter().map(|asset| **asset));
-                client_assets.extend(
-                    client_chunk_group
-                        .referenced_assets
-                        .await?
-                        .iter()
-                        .map(|asset| **asset),
-                );
+                let client_chunk_group = self.client_chunk_group();
+                client_assets.extend(client_chunk_group.direct_assets().await?.iter().copied());
+                let client_chunks = *client_chunk_group.await?.assets;
 
                 let build_manifest = self.build_manifest(client_chunks).to_resolved().await?;
-                let page_loader = self.page_loader(client_chunks);
+                let page_loader = self.page_loader(client_chunks).to_resolved().await?;
                 let client_build_manifest = self
-                    .client_build_manifest(page_loader)
+                    .client_build_manifest(*page_loader)
                     .to_resolved()
                     .await?;
                 client_assets.push(page_loader);
@@ -1349,7 +1328,7 @@ impl PageEndpoint {
             PageEndpointType::SsrOnly => self.ssr_chunk(emit_manifests),
         };
 
-        let client_assets = OutputAssets::new(client_assets).to_resolved().await?;
+        let client_assets: ResolvedVc<OutputAssets> = ResolvedVc::cell(client_assets);
 
         let manifest_path_prefix = get_asset_prefix_from_pathname(&this.pathname);
         let node_root = this.pages_project.project().node_root().owned().await?;
@@ -1380,7 +1359,7 @@ impl PageEndpoint {
             let webpack_stats = generate_webpack_stats(
                 self.client_module_graph(),
                 this.original_name.clone(),
-                client_assets.await?.iter().copied(),
+                client_assets.await?.into_iter().copied(),
             )
             .await?;
             let stats_output = VirtualOutputAsset::new(

crates/next-api/src/pages.rs

@@ -1,12 +1,13 @@
 use anyhow::Result;
-use next_core::{all_assets_from_entries, next_manifests::AssetBinding};
+use next_core::next_manifests::AssetBinding;
 use tracing::Instrument;
 use turbo_rcstr::RcStr;
 use turbo_tasks::{ResolvedVc, TryFlatJoinIterExt, TryJoinIterExt, Vc};
 use turbo_tasks_fs::FileSystemPath;
 use turbopack_core::{
     asset::Asset,
     output::{OutputAsset, OutputAssets},
+    reference::all_assets_from_entries,
 };
 use turbopack_wasm::wasm_edge_var_name;