From 8107a2721b6abb07c3120ac90e03c39f2a44ab0c Mon Sep 17 00:00:00 2001
From: Bjorn Lu <bjornlu.dev@gmail.com>
Date: Fri, 8 Mar 2024 23:03:02 +0800
Subject: [PATCH] Treeshake unused Astro scoped styles (#10291)

---
 .changeset/nine-trains-drop.md                |  5 ++
 .../src/core/build/plugins/plugin-css.ts      | 53 ++++++++++++++++++-
 packages/astro/src/core/compile/compile.ts    | 17 ++----
 packages/astro/src/core/compile/style.ts      | 15 +++---
 packages/astro/src/core/compile/types.ts      | 12 +++++
 packages/astro/src/vite-plugin-astro/index.ts | 23 ++++++--
 packages/astro/src/vite-plugin-astro/types.ts | 23 +++++++-
 packages/astro/test/0-css.test.js             | 10 ++++
 .../pages/barrel-styles/_components/A.astro   |  7 +++
 .../pages/barrel-styles/_components/B.astro   |  7 +++
 .../pages/barrel-styles/_components/C.astro   |  7 +++
 .../pages/barrel-styles/_components/index.js  |  3 ++
 .../0-css/src/pages/barrel-styles/index.astro |  5 ++
 .../src/components/Three.astro                |  6 ---
 .../src/components/Three.js                   |  1 +
 .../src/pages/one.astro                       |  2 +-
 .../src/styles/Three.css                      |  1 +
 17 files changed, 166 insertions(+), 31 deletions(-)
 create mode 100644 .changeset/nine-trains-drop.md
 create mode 100644 packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/A.astro
 create mode 100644 packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/B.astro
 create mode 100644 packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/C.astro
 create mode 100644 packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/index.js
 create mode 100644 packages/astro/test/fixtures/0-css/src/pages/barrel-styles/index.astro
 delete mode 100644 packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.astro
 create mode 100644 packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.js
 create mode 100644 packages/astro/test/fixtures/css-order-dynamic-import/src/styles/Three.css

diff --git a/.changeset/nine-trains-drop.md b/.changeset/nine-trains-drop.md
new file mode 100644
index 000000000000..d7ef4c5e18e8
--- /dev/null
+++ b/.changeset/nine-trains-drop.md
@@ -0,0 +1,5 @@
+---
+"astro": patch
+---
+
+Treeshakes unused Astro component scoped styles
diff --git a/packages/astro/src/core/build/plugins/plugin-css.ts b/packages/astro/src/core/build/plugins/plugin-css.ts
index 12d84fd05918..a84ce37d8a21 100644
--- a/packages/astro/src/core/build/plugins/plugin-css.ts
+++ b/packages/astro/src/core/build/plugins/plugin-css.ts
@@ -1,11 +1,12 @@
 import type { GetModuleInfo } from 'rollup';
-import type { BuildOptions, Plugin as VitePlugin, ResolvedConfig } from 'vite';
+import type { BuildOptions, Plugin as VitePlugin, ResolvedConfig, Rollup } from 'vite';
 import { isBuildableCSSRequest } from '../../../vite-plugin-astro-server/util.js';
 import type { BuildInternals } from '../internal.js';
 import type { AstroBuildPlugin, BuildTarget } from '../plugin.js';
 import type { PageBuildData, StaticBuildOptions, StylesheetAsset } from '../types.js';
 
 import { PROPAGATED_ASSET_FLAG } from '../../../content/consts.js';
+import type { AstroPluginCssMetadata } from '../../../vite-plugin-astro/index.js';
 import * as assetName from '../css-asset-name.js';
 import { moduleIsTopLevelPage, walkParentInfos } from '../graph.js';
 import {
@@ -180,6 +181,32 @@ function rollupPluginAstroBuildCSS(options: PluginOptions): VitePlugin[] {
 		},
 	};
 
+	/**
+	 * This plugin is a port of https://github.com/vitejs/vite/pull/16058. It enables removing unused
+	 * scoped CSS from the bundle if the scoped target (e.g. Astro files) were not bundled.
+	 * Once/If that PR is merged, we can refactor this away, renaming `meta.astroCss` to `meta.vite`.
+	 */
+	const cssScopeToPlugin: VitePlugin = {
+		name: 'astro:rollup-plugin-css-scope-to',
+		renderChunk(_, chunk, __, meta) {
+			for (const id in chunk.modules) {
+				// If this CSS is scoped to its importers exports, check if those importers exports
+				// are rendered in the chunks. If they are not, we can skip bundling this CSS.
+				const modMeta = this.getModuleInfo(id)?.meta as AstroPluginCssMetadata | undefined;
+				const cssScopeTo = modMeta?.astroCss?.cssScopeTo;
+				if (cssScopeTo && !isCssScopeToRendered(cssScopeTo, Object.values(meta.chunks))) {
+					// If this CSS is not used, delete it from the chunk modules so that Vite is unable
+					// to trace that it's used
+					delete chunk.modules[id];
+					const moduleIdsIndex = chunk.moduleIds.indexOf(id);
+					if (moduleIdsIndex > -1) {
+						chunk.moduleIds.splice(moduleIdsIndex, 1);
+					}
+				}
+			}
+		},
+	};
+
 	const singleCssPlugin: VitePlugin = {
 		name: 'astro:rollup-plugin-single-css',
 		enforce: 'post',
@@ -283,7 +310,7 @@ function rollupPluginAstroBuildCSS(options: PluginOptions): VitePlugin[] {
 		},
 	};
 
-	return [cssBuildPlugin, singleCssPlugin, inlineStylesheetsPlugin];
+	return [cssBuildPlugin, cssScopeToPlugin, singleCssPlugin, inlineStylesheetsPlugin];
 }
 
 /***** UTILITY FUNCTIONS *****/
@@ -331,3 +358,25 @@ function appendCSSToPage(
 		}
 	}
 }
+
+/**
+ * `cssScopeTo` is a map of `importer`s to its `export`s. This function iterate each `cssScopeTo` entries
+ * and check if the `importer` and its `export`s exists in the final chunks. If at least one matches,
+ * `cssScopeTo` is considered "rendered" by Rollup and we return true.
+ */
+function isCssScopeToRendered(
+	cssScopeTo: Record<string, string[]>,
+	chunks: Rollup.RenderedChunk[]
+) {
+	for (const moduleId in cssScopeTo) {
+		const exports = cssScopeTo[moduleId];
+		// Find the chunk that renders this `moduleId` and get the rendered module
+		const renderedModule = chunks.find((c) => c.moduleIds.includes(moduleId))?.modules[moduleId];
+		// Return true if `renderedModule` exists and one of its exports is rendered
+		if (renderedModule?.renderedExports.some((e) => exports.includes(e))) {
+			return true;
+		}
+	}
+
+	return false;
+}
diff --git a/packages/astro/src/core/compile/compile.ts b/packages/astro/src/core/compile/compile.ts
index 9b95e2c89aab..52e86b2b07a1 100644
--- a/packages/astro/src/core/compile/compile.ts
+++ b/packages/astro/src/core/compile/compile.ts
@@ -10,7 +10,8 @@ import type { AstroError } from '../errors/errors.js';
 import { AggregateError, CompilerError } from '../errors/errors.js';
 import { AstroErrorData } from '../errors/index.js';
 import { resolvePath } from '../util.js';
-import { createStylePreprocessor } from './style.js';
+import { type PartialCompileCssResult, createStylePreprocessor } from './style.js';
+import type { CompileCssResult } from './types.js';
 
 export interface CompileProps {
 	astroConfig: AstroConfig;
@@ -20,14 +21,6 @@ export interface CompileProps {
 	source: string;
 }
 
-export interface CompileCssResult {
-	code: string;
-	/**
-	 * The dependencies of the transformed CSS (Normalized paths)
-	 */
-	dependencies?: string[];
-}
-
 export interface CompileResult extends Omit<TransformResult, 'css'> {
 	css: CompileCssResult[];
 }
@@ -42,7 +35,7 @@ export async function compile({
 	// Because `@astrojs/compiler` can't return the dependencies for each style transformed,
 	// we need to use an external array to track the dependencies whenever preprocessing is called,
 	// and we'll rebuild the final `css` result after transformation.
-	const cssDeps: CompileCssResult['dependencies'][] = [];
+	const cssPartialCompileResults: PartialCompileCssResult[] = [];
 	const cssTransformErrors: AstroError[] = [];
 	let transformResult: TransformResult;
 
@@ -71,7 +64,7 @@ export async function compile({
 			preprocessStyle: createStylePreprocessor({
 				filename,
 				viteConfig,
-				cssDeps,
+				cssPartialCompileResults,
 				cssTransformErrors,
 			}),
 			async resolvePath(specifier) {
@@ -96,8 +89,8 @@ export async function compile({
 	return {
 		...transformResult,
 		css: transformResult.css.map((code, i) => ({
+			...cssPartialCompileResults[i],
 			code,
-			dependencies: cssDeps[i],
 		})),
 	};
 }
diff --git a/packages/astro/src/core/compile/style.ts b/packages/astro/src/core/compile/style.ts
index 45d45c99e7c1..5d517a5146f5 100644
--- a/packages/astro/src/core/compile/style.ts
+++ b/packages/astro/src/core/compile/style.ts
@@ -2,17 +2,19 @@ import fs from 'node:fs';
 import type { TransformOptions } from '@astrojs/compiler';
 import { type ResolvedConfig, normalizePath, preprocessCSS } from 'vite';
 import { AstroErrorData, CSSError, positionAt } from '../errors/index.js';
-import type { CompileCssResult } from './compile.js';
+import type { CompileCssResult } from './types.js';
+
+export type PartialCompileCssResult = Pick<CompileCssResult, 'isGlobal' | 'dependencies'>;
 
 export function createStylePreprocessor({
 	filename,
 	viteConfig,
-	cssDeps,
+	cssPartialCompileResults,
 	cssTransformErrors,
 }: {
 	filename: string;
 	viteConfig: ResolvedConfig;
-	cssDeps: CompileCssResult['dependencies'][];
+	cssPartialCompileResults: Partial<CompileCssResult>[];
 	cssTransformErrors: Error[];
 }): TransformOptions['preprocessStyle'] {
 	let processedStylesCount = 0;
@@ -24,9 +26,10 @@ export function createStylePreprocessor({
 		try {
 			const result = await preprocessCSS(content, id, viteConfig);
 
-			if (result.deps) {
-				cssDeps[index] = [...result.deps].map((dep) => normalizePath(dep));
-			}
+			cssPartialCompileResults[index] = {
+				isGlobal: !!attrs['is:global'],
+				dependencies: result.deps ? [...result.deps].map((dep) => normalizePath(dep)) : [],
+			};
 
 			let map: string | undefined;
 			if (result.map) {
diff --git a/packages/astro/src/core/compile/types.ts b/packages/astro/src/core/compile/types.ts
index 1ef4bdfdcd52..9d1c653cb481 100644
--- a/packages/astro/src/core/compile/types.ts
+++ b/packages/astro/src/core/compile/types.ts
@@ -10,3 +10,15 @@ export type TransformStyle = (
 	source: string,
 	lang: string
 ) => TransformStyleResult | Promise<TransformStyleResult>;
+
+export interface CompileCssResult {
+	code: string;
+	/**
+	 * Whether this is `<style is:global>`
+	 */
+	isGlobal: boolean;
+	/**
+	 * The dependencies of the transformed CSS (Normalized/forward-slash-only absolute paths)
+	 */
+	dependencies: string[];
+}
diff --git a/packages/astro/src/vite-plugin-astro/index.ts b/packages/astro/src/vite-plugin-astro/index.ts
index 7f46069aeeaa..bc9dbd5aa449 100644
--- a/packages/astro/src/vite-plugin-astro/index.ts
+++ b/packages/astro/src/vite-plugin-astro/index.ts
@@ -2,7 +2,11 @@ import type { SourceDescription } from 'rollup';
 import type * as vite from 'vite';
 import type { AstroConfig, AstroSettings } from '../@types/astro.js';
 import type { Logger } from '../core/logger/core.js';
-import type { CompileMetadata, PluginMetadata as AstroPluginMetadata } from './types.js';
+import type {
+	CompileMetadata,
+	PluginCssMetadata as AstroPluginCssMetadata,
+	PluginMetadata as AstroPluginMetadata,
+} from './types.js';
 
 import { normalizePath } from 'vite';
 import { normalizeFilename } from '../vite-plugin-utils/index.js';
@@ -11,7 +15,7 @@ import { handleHotUpdate } from './hmr.js';
 import { parseAstroRequest } from './query.js';
 import { loadId } from './utils.js';
 export { getAstroMetadata } from './metadata.js';
-export type { AstroPluginMetadata };
+export type { AstroPluginMetadata, AstroPluginCssMetadata };
 
 interface AstroPluginOptions {
 	settings: AstroSettings;
@@ -116,7 +120,20 @@ export default function astro({ settings, logger }: AstroPluginOptions): vite.Pl
 					// Register dependencies from preprocessing this style
 					result.dependencies?.forEach((dep) => this.addWatchFile(dep));
 
-					return { code: result.code };
+					return {
+						code: result.code,
+						// This metadata is used by `cssScopeToPlugin` to remove this module from the bundle
+						// if the `filename` default export (the Astro component) is unused.
+						meta: result.isGlobal
+							? undefined
+							: ({
+									astroCss: {
+										cssScopeTo: {
+											[filename]: ['default'],
+										},
+									},
+								} satisfies AstroPluginCssMetadata),
+					};
 				}
 				case 'script': {
 					if (typeof query.index === 'undefined') {
diff --git a/packages/astro/src/vite-plugin-astro/types.ts b/packages/astro/src/vite-plugin-astro/types.ts
index a954ac109ea2..8e82165f5ebd 100644
--- a/packages/astro/src/vite-plugin-astro/types.ts
+++ b/packages/astro/src/vite-plugin-astro/types.ts
@@ -1,6 +1,6 @@
 import type { HoistedScript, TransformResult } from '@astrojs/compiler';
 import type { PropagationHint } from '../@types/astro.js';
-import type { CompileCssResult } from '../core/compile/compile.js';
+import type { CompileCssResult } from '../core/compile/types.js';
 
 export interface PageOptions {
 	prerender?: boolean;
@@ -17,6 +17,27 @@ export interface PluginMetadata {
 	};
 }
 
+export interface PluginCssMetadata {
+	astroCss: {
+		/**
+		 * For Astro CSS virtual modules, it can scope to the main Astro module's default export
+		 * so that if those exports are treeshaken away, the CSS module will also be treeshaken.
+		 *
+		 * Example config if the CSS id is `/src/Foo.astro?astro&type=style&lang.css`:
+		 * ```js
+		 * cssScopeTo: {
+		 *   '/src/Foo.astro': ['default']
+		 * }
+		 * ```
+		 *
+		 * The above is the only config we use today, but we're exposing as a `Record` to follow the
+		 * upstream Vite implementation: https://github.com/vitejs/vite/pull/16058. When/If that lands,
+		 * we can also remove our custom implementation.
+		 */
+		cssScopeTo: Record<string, string[]>;
+	};
+}
+
 export interface CompileMetadata {
 	/** Used for HMR to compare code changes */
 	originalCode: string;
diff --git a/packages/astro/test/0-css.test.js b/packages/astro/test/0-css.test.js
index e3f182b161e9..71dc29db13ce 100644
--- a/packages/astro/test/0-css.test.js
+++ b/packages/astro/test/0-css.test.js
@@ -95,6 +95,16 @@ describe('CSS', function () {
 			it('<style lang="scss">', async () => {
 				assert.match(bundledCSS, /h1\[data-astro-cid-[^{]*\{color:#ff69b4\}/);
 			});
+
+			it('Styles through barrel files should only include used Astro scoped styles', async () => {
+				const barrelHtml = await fixture.readFile('/barrel-styles/index.html');
+				const barrel$ = cheerio.load(barrelHtml);
+				const barrelBundledCssHref = barrel$('link[rel=stylesheet][href^=/_astro/]').attr('href');
+				const style = await fixture.readFile(barrelBundledCssHref.replace(/^\/?/, '/'));
+				assert.match(style, /\.comp-a\[data-astro-cid/);
+				assert.match(style, /\.comp-c\{/);
+				assert.doesNotMatch(style, /\.comp-b/);
+			});
 		});
 
 		describe('Styles in src/', () => {
diff --git a/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/A.astro b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/A.astro
new file mode 100644
index 000000000000..96f96d7ece1a
--- /dev/null
+++ b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/A.astro
@@ -0,0 +1,7 @@
+<p class="comp-a">A</p>
+
+<style>
+  .comp-a {
+    color: red;
+  }
+</style>
diff --git a/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/B.astro b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/B.astro
new file mode 100644
index 000000000000..e5723da09d13
--- /dev/null
+++ b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/B.astro
@@ -0,0 +1,7 @@
+<p class="comp-b">B</p>
+
+<style>
+  .comp-b {
+    color: red;
+  }
+</style>
diff --git a/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/C.astro b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/C.astro
new file mode 100644
index 000000000000..856ae398a813
--- /dev/null
+++ b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/C.astro
@@ -0,0 +1,7 @@
+<p class="comp-c">C</p>
+
+<style is:global>
+  .comp-c {
+    color: red;
+  }
+</style>
diff --git a/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/index.js b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/index.js
new file mode 100644
index 000000000000..151125f620e9
--- /dev/null
+++ b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/_components/index.js
@@ -0,0 +1,3 @@
+export { default as A } from './A.astro';
+export { default as B } from './B.astro';
+export { default as C } from './C.astro';
diff --git a/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/index.astro b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/index.astro
new file mode 100644
index 000000000000..99b3f0b3a29b
--- /dev/null
+++ b/packages/astro/test/fixtures/0-css/src/pages/barrel-styles/index.astro
@@ -0,0 +1,5 @@
+---
+import { A } from './_components';
+---
+
+<A />
diff --git a/packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.astro b/packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.astro
deleted file mode 100644
index 8918fdc78fef..000000000000
--- a/packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.astro
+++ /dev/null
@@ -1,6 +0,0 @@
----
----
-<style>
-	body { background: yellow;}
-</style>
-<div>testing</div>
diff --git a/packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.js b/packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.js
new file mode 100644
index 000000000000..e3aa682ff448
--- /dev/null
+++ b/packages/astro/test/fixtures/css-order-dynamic-import/src/components/Three.js
@@ -0,0 +1 @@
+import "../styles/Three.css"
diff --git a/packages/astro/test/fixtures/css-order-dynamic-import/src/pages/one.astro b/packages/astro/test/fixtures/css-order-dynamic-import/src/pages/one.astro
index abd194abc72b..d69fade463cb 100644
--- a/packages/astro/test/fixtures/css-order-dynamic-import/src/pages/one.astro
+++ b/packages/astro/test/fixtures/css-order-dynamic-import/src/pages/one.astro
@@ -1,7 +1,7 @@
 ---
 import '../components/One.astro';
 import '../components/Two.astro';
-await import('../components/Three.astro');
+await import('../components/Three.js');
 ---
 <html>
 <head>
diff --git a/packages/astro/test/fixtures/css-order-dynamic-import/src/styles/Three.css b/packages/astro/test/fixtures/css-order-dynamic-import/src/styles/Three.css
new file mode 100644
index 000000000000..a9f2b8f49eb9
--- /dev/null
+++ b/packages/astro/test/fixtures/css-order-dynamic-import/src/styles/Three.css
@@ -0,0 +1 @@
+body { background: yellow;}
\ No newline at end of file