Refactor prerendering chunk handling

.changeset/three-boxes-sniff.md

@@ -0,0 +1,5 @@
+---
+"astro": patch
+---
+
+Refactors prerendering chunk handling to correctly remove unused code during the SSR runtime

packages/astro/src/core/build/internal.ts

@@ -113,6 +113,11 @@ export interface BuildInternals {
 	ssrSplitEntryChunks: Map<string, Rollup.OutputChunk>;
 	componentMetadata: SSRResult['componentMetadata'];
 	middlewareEntryPoint?: URL;
+
+	/**
+	 * Chunks in the bundle that are only used in prerendering that we can delete later
+	 */
+	prerenderOnlyChunks: Rollup.OutputChunk[];
 }
 
 /**
@@ -151,6 +156,7 @@ export function createBuildInternals(): BuildInternals {
 		ssrSplitEntryChunks: new Map(),
 		entryPoints: new Map(),
 		cacheManifestUsed: false,
+		prerenderOnlyChunks: [],
 	};
 }
 

packages/astro/src/core/build/page-data.ts

@@ -57,7 +57,6 @@ export async function collectPagesData(
 				moduleSpecifier: '',
 				styles: [],
 				hoistedScript: undefined,
-				hasSharedModules: false,
 			};
 
 			clearInterval(routeCollectionLogTimeout);
@@ -80,7 +79,6 @@ export async function collectPagesData(
 			moduleSpecifier: '',
 			styles: [],
 			hoistedScript: undefined,
-			hasSharedModules: false,
 		};
 	}
 

packages/astro/src/core/build/plugins/plugin-chunks.ts

@@ -12,6 +12,15 @@ export function vitePluginChunks(): VitePlugin {
 					if (id.includes('astro/dist/runtime/server/')) {
 						return 'astro/server';
 					}
+					// Split the Astro runtime into a separate chunk for readability
+					if (id.includes('astro/dist/runtime')) {
+						return 'astro';
+					}
+					// Place `astro/env/setup` import in its own chunk to prevent Rollup's TLA bug
+					// https://github.com/rollup/rollup/issues/4708
+					if (id.includes('astro/dist/env/setup')) {
+						return 'astro/env-setup';
+					}
 				},
 			});
 		},

packages/astro/src/core/build/plugins/plugin-prerender.ts

@@ -1,87 +1,111 @@
-import path from 'node:path';
-import type { Plugin as VitePlugin } from 'vite';
+import type { Rollup, Plugin as VitePlugin } from 'vite';
 import { getPrerenderMetadata } from '../../../prerender/metadata.js';
 import type { BuildInternals } from '../internal.js';
 import type { AstroBuildPlugin } from '../plugin.js';
 import type { StaticBuildOptions } from '../types.js';
-import { extendManualChunks } from './util.js';
+import { ASTRO_PAGE_RESOLVED_MODULE_ID } from './plugin-pages.js';
+import { getPagesFromVirtualModulePageName } from './util.js';
 
-function vitePluginPrerender(opts: StaticBuildOptions, internals: BuildInternals): VitePlugin {
+function vitePluginPrerender(internals: BuildInternals): VitePlugin {
 	return {
 		name: 'astro:rollup-plugin-prerender',
 
-		outputOptions(outputOptions) {
-			extendManualChunks(outputOptions, {
-				after(id, meta) {
-					// Split the Astro runtime into a separate chunk for readability
-					if (id.includes('astro/dist/runtime')) {
-						return 'astro';
-					}
-					const pageInfo = internals.pagesByViteID.get(id);
-					let hasSharedModules = false;
-					if (pageInfo) {
-						// prerendered pages should be split into their own chunk
-						// Important: this can't be in the `pages/` directory!
-						if (getPrerenderMetadata(meta.getModuleInfo(id)!)) {
-							const infoMeta = meta.getModuleInfo(id)!;
+		generateBundle(_, bundle) {
+			const moduleIds = this.getModuleIds();
+			for (const id of moduleIds) {
+				const pageInfo = internals.pagesByViteID.get(id);
+				if (!pageInfo) continue;
+				const moduleInfo = this.getModuleInfo(id);
+				if (!moduleInfo) continue;
 
-							// Here, we check if this page is importing modules that are shared among other modules e.g. middleware, other SSR pages, etc.
-							// we loop the modules that the current page imports
-							for (const moduleId of infoMeta.importedIds) {
-								// we retrieve the metadata of the module
-								const moduleMeta = meta.getModuleInfo(moduleId)!;
-								if (
-									// a shared modules should be inside the `src/` folder, at least
-									moduleMeta.id.startsWith(opts.settings.config.srcDir.pathname) &&
-									// and has at least two importers: the current page and something else
-									moduleMeta.importers.length > 1
-								) {
-									// Now, we have to trace back the modules imported and analyze them;
-									// understanding if a module is eventually shared between two pages isn't easy, because a module could
-									// be imported by a page and a component that is eventually imported by a page.
-									//
-									// Given the previous statement, we only check if
-									// - the module is a page, and it's not pre-rendered
-									// - the module is the middleware
-									// If one of these conditions is met, we need a separate chunk
-									for (const importer of moduleMeta.importedIds) {
-										// we don't want to analyze the same module again, so we skip it
-										if (importer !== id) {
-											const importerModuleMeta = meta.getModuleInfo(importer);
-											if (importerModuleMeta) {
-												// if the module is inside the pages
-												if (importerModuleMeta.id.includes('/pages')) {
-													// we check if it's not pre-rendered
-													if (getPrerenderMetadata(importerModuleMeta) === false) {
-														hasSharedModules = true;
-														break;
-													}
-												}
-												// module isn't an Astro route/page, it could be a middleware
-												else if (importerModuleMeta.id.includes('/middleware')) {
-													hasSharedModules = true;
-													break;
-												}
-											}
-										}
-									}
-								}
-							}
+				const prerender = !!getPrerenderMetadata(moduleInfo);
+				pageInfo.route.prerender = prerender;
+			}
 
-							pageInfo.hasSharedModules = hasSharedModules;
-							pageInfo.route.prerender = true;
-							return 'prerender';
-						}
-						pageInfo.route.prerender = false;
-						// dynamic pages should all go in their own chunk in the pages/* directory
-						return `pages/${path.basename(pageInfo.component)}`;
-					}
-				},
-			});
+			// Find all chunks used in the SSR runtime (that aren't used for prerendering only), then use
+			// the Set to find the inverse, where chunks that are only used for prerendering. It's faster
+			// to compute `internals.prerenderOnlyChunks` this way. The prerendered chunks will be deleted
+			// after we finish prerendering.
+			const nonPrerenderOnlyChunks = getNonPrerenderOnlyChunks(bundle, internals);
+			internals.prerenderOnlyChunks = Object.values(bundle).filter((chunk) => {
+				return chunk.type === 'chunk' && !nonPrerenderOnlyChunks.has(chunk);
+			}) as Rollup.OutputChunk[];
 		},
 	};
 }
 
+function getNonPrerenderOnlyChunks(bundle: Rollup.OutputBundle, internals: BuildInternals) {
+	const chunks = Object.values(bundle);
+
+	const prerenderOnlyEntryChunks = new Set<Rollup.OutputChunk>();
+	const nonPrerenderOnlyEntryChunks = new Set<Rollup.OutputChunk>();
+	for (const chunk of chunks) {
+		if (chunk.type === 'chunk' && (chunk.isEntry || chunk.isDynamicEntry)) {
+			// See if this entry chunk is prerendered, if so, skip it
+			if (chunk.facadeModuleId?.startsWith(ASTRO_PAGE_RESOLVED_MODULE_ID)) {
+				const pageDatas = getPagesFromVirtualModulePageName(
+					internals,
+					ASTRO_PAGE_RESOLVED_MODULE_ID,
+					chunk.facadeModuleId
+				);
+				const prerender = pageDatas.every((pageData) => pageData.route.prerender);
+				if (prerender) {
+					prerenderOnlyEntryChunks.add(chunk);
+					continue;
+				}
+			}
+			// Ideally we should record entries when `functionPerRoute` is enabled, but this breaks some tests
+			// that expect the entrypoint to still exist even if it should be unused.
+			// TODO: Revisit this so we can delete additional unused chunks
+			// else if (chunk.facadeModuleId?.startsWith(RESOLVED_SPLIT_MODULE_ID)) {
+			// 	const pageDatas = getPagesFromVirtualModulePageName(
+			// 		internals,
+			// 		RESOLVED_SPLIT_MODULE_ID,
+			// 		chunk.facadeModuleId
+			// 	);
+			// 	const prerender = pageDatas.every((pageData) => pageData.route.prerender);
+			// 	if (prerender) {
+			// 		prerenderOnlyEntryChunks.add(chunk);
+			// 		continue;
+			// 	}
+			// }
+
+			nonPrerenderOnlyEntryChunks.add(chunk);
+		}
+	}
+
+	// From the `nonPrerenderedEntryChunks`, we crawl all the imports/dynamicImports to find all
+	// other chunks that are use by the non-prerendered runtime
+	const nonPrerenderOnlyChunks = new Set();
+	for (const entryChunk of nonPrerenderOnlyEntryChunks) {
+		crawlChunk(entryChunk);
+	}
+
+	function crawlChunk(chunk: Rollup.OutputChunk) {
+		if (nonPrerenderOnlyChunks.has(chunk)) return;
+		nonPrerenderOnlyChunks.add(chunk);
+		for (const importFileName of chunk.imports) {
+			const importChunk = bundle[importFileName];
+			if (importChunk?.type === 'chunk') {
+				crawlChunk(importChunk);
+			}
+		}
+		for (const dynamicImportFileName of chunk.dynamicImports) {
+			const dynamicImportChunk = bundle[dynamicImportFileName];
+			// The main server entry (entry.mjs) may import a prerender-only entry chunk, we skip in this case
+			// to prevent incorrectly marking it as non-prerendered.
+			if (
+				dynamicImportChunk?.type === 'chunk' &&
+				!prerenderOnlyEntryChunks.has(dynamicImportChunk)
+			) {
+				crawlChunk(dynamicImportChunk);
+			}
+		}
+	}
+
+	return nonPrerenderOnlyChunks;
+}
+
 export function pluginPrerender(
 	opts: StaticBuildOptions,
 	internals: BuildInternals
@@ -96,7 +120,7 @@ export function pluginPrerender(
 		hooks: {
 			'build:before': () => {
 				return {
-					vitePlugin: vitePluginPrerender(opts, internals),
+					vitePlugin: vitePluginPrerender(internals),
 				};
 			},
 		},

packages/astro/src/core/build/plugins/plugin-ssr.ts

@@ -27,7 +27,17 @@ function vitePluginSSR(
 		name: '@astrojs/vite-plugin-astro-ssr-server',
 		enforce: 'post',
 		options(opts) {
-			return addRollupInput(opts, [SSR_VIRTUAL_MODULE_ID]);
+			const inputs = new Set<string>();
+
+			for (const pageData of Object.values(options.allPages)) {
+				if (routeIsRedirect(pageData.route)) {
+					continue;
+				}
+				inputs.add(getVirtualModulePageName(ASTRO_PAGE_MODULE_ID, pageData.component));
+			}
+
+			inputs.add(SSR_VIRTUAL_MODULE_ID);
+			return addRollupInput(opts, Array.from(inputs));
 		},
 		resolveId(id) {
 			if (id === SSR_VIRTUAL_MODULE_ID) {
@@ -72,7 +82,6 @@ function vitePluginSSR(
 				contents.push(...ssrCode.contents);
 				return [...imports, ...contents, ...exports].join('\n');
 			}
-			return void 0;
 		},
 		async generateBundle(_opts, bundle) {
 			// Add assets from this SSR chunk as well.
@@ -141,23 +150,20 @@ function vitePluginSSRSplit(
 	adapter: AstroAdapter,
 	options: StaticBuildOptions
 ): VitePlugin {
-	const functionPerRouteEnabled = isFunctionPerRouteEnabled(options.settings.adapter);
 	return {
 		name: '@astrojs/vite-plugin-astro-ssr-split',
 		enforce: 'post',
 		options(opts) {
-			if (functionPerRouteEnabled) {
-				const inputs = new Set<string>();
+			const inputs = new Set<string>();
 
-				for (const pageData of Object.values(options.allPages)) {
-					if (routeIsRedirect(pageData.route)) {
-						continue;
-					}
-					inputs.add(getVirtualModulePageName(SPLIT_MODULE_ID, pageData.component));
+			for (const pageData of Object.values(options.allPages)) {
+				if (routeIsRedirect(pageData.route)) {
+					continue;
 				}
-
-				return addRollupInput(opts, Array.from(inputs));
+				inputs.add(getVirtualModulePageName(SPLIT_MODULE_ID, pageData.component));
 			}
+
+			return addRollupInput(opts, Array.from(inputs));
 		},
 		resolveId(id) {
 			if (id.startsWith(SPLIT_MODULE_ID)) {
@@ -185,7 +191,6 @@ function vitePluginSSRSplit(
 
 				return [...imports, ...contents, ...exports].join('\n');
 			}
-			return void 0;
 		},
 		async generateBundle(_opts, bundle) {
 			// Add assets from this SSR chunk as well.