refactor: Ensure module preload polyfill is inlined into main bundle (#147)

* refactor: Ensure module preload polyfill is inlined into main bundle

* docs: Update patch comments

* refactor: Bail out on multiple outputs

Author: rschristian

src/prerender.ts

@@ -10,7 +10,12 @@ import type { Plugin, ResolvedConfig } from "vite";
 
 // Vite re-exports Rollup's type defs in newer versions,
 // merge into above type import when we bump the Vite devDep
-import type { InputOption, OutputAsset, OutputChunk } from "rollup";
+import type {
+	InputOption,
+	OutputAsset,
+	OutputChunk,
+	OutputOptions,
+} from "rollup";
 
 interface HeadElement {
 	type: string;
@@ -74,6 +79,7 @@ export function PrerenderPlugin({
 	additionalPrerenderRoutes,
 }: PrerenderPluginOptions = {}): Plugin {
 	const preloadHelperId = "vite/preload-helper";
+	const preloadPolyfillId = "vite/modulepreload-polyfill";
 	let viteConfig = {} as ResolvedConfig;
 	let userEnabledSourceMaps: boolean | undefined;
 
@@ -123,6 +129,34 @@ export function PrerenderPlugin({
 			config.build.sourcemap = true;
 
 			viteConfig = config;
+
+			// With this plugin adding an additional input, Rollup/Vite tries to be smart
+			// and extract our prerender script (which is often their main bundle) to a separate
+			// chunk that the entry & prerender chunks can depend on. Unfortunately, this means the
+			// first script the browser loads is the module preload polyfill & a sync import of the main
+			// bundle. This is obviously less than ideal as the main bundle should be directly referenced
+			// by the user's HTML to speed up loading a bit.
+
+			// We're only going to alter the chunking behavior in the default cases, where the user and/or
+			// other plugins haven't already configured this. It'd be impossible to avoid breakages otherwise.
+			if (
+				Array.isArray(config.build.rollupOptions.output) ||
+				(config.build.rollupOptions.output as OutputOptions)?.manualChunks
+			) {
+				return;
+			}
+
+			config.build.rollupOptions.output ??= {};
+			(config.build.rollupOptions.output as OutputOptions).manualChunks = (
+				id: string,
+			) => {
+				if (
+					id.includes(prerenderScript as string) ||
+					id.includes(preloadPolyfillId)
+				) {
+					return "index";
+				}
+			};
 		},
 		async options(opts) {
 			if (!opts.input) return;
@@ -139,15 +173,15 @@ export function PrerenderPlugin({
 						: { ...opts.input, prerenderEntry: prerenderScript };
 			opts.preserveEntrySignatures = "allow-extension";
 		},
-		// Injects a window check into Vite's preload helper, instantly resolving
-		// the module rather than attempting to add a <link> to the document.
+		// Injects window checks into Vite's preload helper & modulepreload polyfill
 		transform(code, id) {
-			// Vite keeps changing up the ID, best we can do for cross-version
-			// compat is an `includes`
 			if (id.includes(preloadHelperId)) {
+				// Injects a window check into Vite's preload helper, instantly resolving
+				// the module rather than attempting to add a <link> to the document.
+				const s = new MagicString(code);
+
 				// Through v5.0.4
 				// https://github.com/vitejs/vite/blob/b93dfe3e08f56cafe2e549efd80285a12a3dc2f0/packages/vite/src/node/plugins/importAnalysisBuild.ts#L95-L98
-				const s = new MagicString(code);
 				s.replace(
 					`if (!__VITE_IS_MODERN__ || !deps || deps.length === 0) {`,
 					`if (!__VITE_IS_MODERN__ || !deps || deps.length === 0 || typeof window === 'undefined') {`,
@@ -162,6 +196,23 @@ export function PrerenderPlugin({
 					code: s.toString(),
 					map: s.generateMap({ hires: true }),
 				};
+			} else if (id.includes(preloadPolyfillId)) {
+				const s = new MagicString(code);
+				// Replacement for `'link'` && `"link"` as the output from their tooling has
+				// differed over the years. Should be better than switching to regex.
+				// https://github.com/vitejs/vite/blob/20fdf210ee0ac0824b2db74876527cb7f378a9e8/packages/vite/src/node/plugins/modulePreloadPolyfill.ts#L62
+				s.replace(
+					`const relList = document.createElement('link').relList;`,
+					`if (typeof window === "undefined") return;\n  const relList = document.createElement('link').relList;`,
+				);
+				s.replace(
+					`const relList = document.createElement("link").relList;`,
+					`if (typeof window === "undefined") return;\n  const relList = document.createElement("link").relList;`,
+				);
+				return {
+					code: s.toString(),
+					map: s.generateMap({ hires: true }),
+				};
 			}
 		},
 		async generateBundle(_opts, bundle) {

test/build.test.mjs

@@ -2,6 +2,7 @@ import { execFile } from "node:child_process";
 import { test } from "node:test";
 import { promisify } from "node:util";
 import { promises as fs } from "node:fs";
+import path from "node:path";
 import assert from "node:assert";
 import { dir } from "./util.mjs";
 
@@ -34,4 +35,8 @@ test("builds demo successfully", async () => {
 
 	// `additionalPrerenderRoutes` config option
 	assert.doesNotThrow(async () => await fs.access(dir("demo/dist/404/index.html")));
+
+	const outputFiles = await fs.readdir(path.join(dir("demo/dist"), 'assets'));
+	const outputIndexJS = outputFiles.filter(f => /^index\..+\.js$/.test(f));
+	assert.strictEqual(outputIndexJS.length, 1);
 });