Treeshake exported client components that are not imported (#6527)

* Treeshake exported client components that are not imported

* Fix plugin name

* Fix mdx test

.changeset/cyan-camels-swim.md

@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Treeshake exported client components that are not imported

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

@@ -48,15 +48,25 @@ export interface BuildInternals {
 	pagesByClientOnly: Map<string, Set<PageBuildData>>;
 
 	/**
-	 * A list of hydrated components that are discovered during the SSR build
+	 * A map of hydrated components to export names that are discovered during the SSR build.
 	 * These will be used as the top-level entrypoints for the client build.
+	 *
+	 * @example
+	 * '/project/Component1.jsx' => ['default']
+	 * '/project/Component2.jsx' => ['Counter', 'Timer']
+	 * '/project/Component3.jsx' => ['*']
 	 */
-	discoveredHydratedComponents: Set<string>;
+	discoveredHydratedComponents: Map<string, string[]>;
 	/**
-	 * A list of client:only components that are discovered during the SSR build
+	 * A list of client:only components to export names that are discovered during the SSR build.
 	 * These will be used as the top-level entrypoints for the client build.
+	 *
+	 * @example
+	 * '/project/Component1.jsx' => ['default']
+	 * '/project/Component2.jsx' => ['Counter', 'Timer']
+	 * '/project/Component3.jsx' => ['*']
 	 */
-	discoveredClientOnlyComponents: Set<string>;
+	discoveredClientOnlyComponents: Map<string, string[]>;
 	/**
 	 * A list of hoisted scripts that are discovered during the SSR build
 	 * These will be used as the top-level entrypoints for the client build.
@@ -93,8 +103,8 @@ export function createBuildInternals(): BuildInternals {
 		pagesByViteID: new Map(),
 		pagesByClientOnly: new Map(),
 
-		discoveredHydratedComponents: new Set(),
-		discoveredClientOnlyComponents: new Set(),
+		discoveredHydratedComponents: new Map(),
+		discoveredClientOnlyComponents: new Map(),
 		discoveredScripts: new Set(),
 		staticFiles: new Set(),
 		propagation: new Map(),

packages/astro/src/core/build/plugins/index.ts

@@ -3,6 +3,7 @@ import { astroHeadPropagationBuildPlugin } from '../../../vite-plugin-head-propa
 import type { AstroBuildPluginContainer } from '../plugin';
 import { pluginAliasResolve } from './plugin-alias-resolve.js';
 import { pluginAnalyzer } from './plugin-analyzer.js';
+import { pluginComponentEntry } from './plugin-component-entry.js';
 import { pluginCSS } from './plugin-css.js';
 import { pluginHoistedScripts } from './plugin-hoisted-scripts.js';
 import { pluginInternals } from './plugin-internals.js';
@@ -11,6 +12,7 @@ import { pluginPrerender } from './plugin-prerender.js';
 import { pluginSSR } from './plugin-ssr.js';
 
 export function registerAllPlugins({ internals, options, register }: AstroBuildPluginContainer) {
+	register(pluginComponentEntry(internals));
 	register(pluginAliasResolve(internals));
 	register(pluginAnalyzer(internals));
 	register(pluginInternals(internals));

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

@@ -137,7 +137,12 @@ export function vitePluginAnalyzer(internals: BuildInternals): VitePlugin {
 
 				for (const c of astro.hydratedComponents) {
 					const rid = c.resolvedPath ? decodeURI(c.resolvedPath) : c.specifier;
-					internals.discoveredHydratedComponents.add(rid);
+					if (internals.discoveredHydratedComponents.has(rid)) {
+						const exportNames = internals.discoveredHydratedComponents.get(rid);
+						exportNames?.push(c.exportName)
+					} else {
+						internals.discoveredHydratedComponents.set(rid, [c.exportName]);
+					}
 				}
 
 				// Scan hoisted scripts
@@ -148,7 +153,12 @@ export function vitePluginAnalyzer(internals: BuildInternals): VitePlugin {
 
 					for (const c of astro.clientOnlyComponents) {
 						const cid = c.resolvedPath ? decodeURI(c.resolvedPath) : c.specifier;
-						internals.discoveredClientOnlyComponents.add(cid);
+						if (internals.discoveredClientOnlyComponents.has(cid)) {
+							const exportNames = internals.discoveredClientOnlyComponents.get(cid);
+							exportNames?.push(c.exportName)
+						} else {
+							internals.discoveredClientOnlyComponents.set(cid, [c.exportName]);
+						}
 						clientOnlys.push(cid);
 
 						const resolvedId = await this.resolve(c.specifier, id);

packages/astro/src/core/build/plugins/plugin-component-entry.ts

@@ -0,0 +1,89 @@
+import type { Plugin as VitePlugin } from 'vite';
+import type { BuildInternals } from '../internal.js';
+import type { AstroBuildPlugin } from '../plugin.js';
+
+const astroEntryPrefix = '\0astro-entry:';
+
+/**
+ * When adding hydrated or client:only components as Rollup inputs, sometimes we're not using all
+ * of the export names, e.g. `import { Counter } from './ManyComponents.jsx'`. This plugin proxies
+ * entries to re-export only the names the user is using.
+ */
+export function vitePluginComponentEntry(internals: BuildInternals): VitePlugin {
+	const componentToExportNames: Map<string, string[]> = new Map();
+
+	mergeComponentExportNames(internals.discoveredHydratedComponents);
+	mergeComponentExportNames(internals.discoveredClientOnlyComponents);
+
+	for (const [componentId, exportNames] of componentToExportNames) {
+		// If one of the imports has a dot, it's a namespaced import, e.g. `import * as foo from 'foo'`
+		// and `<foo.Counter />`, in which case we re-export `foo` entirely and we don't need to handle
+		// it in this plugin as it's default behaviour from Rollup.
+		if (exportNames.some((name) => name.includes('.') || name === '*')) {
+			componentToExportNames.delete(componentId);
+		} else {
+			componentToExportNames.set(componentId, Array.from(new Set(exportNames)));
+		}
+	}
+
+	function mergeComponentExportNames(components: Map<string, string[]>) {
+		for (const [componentId, exportNames] of components) {
+			if (componentToExportNames.has(componentId)) {
+				componentToExportNames.get(componentId)?.push(...exportNames);
+			} else {
+				componentToExportNames.set(componentId, exportNames);
+			}
+		}
+	}
+
+	return {
+		name: '@astro/plugin-component-entry',
+		enforce: 'pre',
+		config(config) {
+			const rollupInput = config.build?.rollupOptions?.input;
+			// Astro passes an array of inputs by default. Even though other Vite plugins could
+			// change this to an object, it shouldn't happen in practice as our plugin runs first.
+			if (Array.isArray(rollupInput)) {
+				// @ts-expect-error input is definitely defined here, but typescript thinks it doesn't
+				config.build.rollupOptions.input = rollupInput.map((id) => {
+					if (componentToExportNames.has(id)) {
+						return astroEntryPrefix + id;
+					} else {
+						return id;
+					}
+				});
+			}
+		},
+		async resolveId(id) {
+			if (id.startsWith(astroEntryPrefix)) {
+				return id;
+			}
+		},
+		async load(id) {
+			if (id.startsWith(astroEntryPrefix)) {
+				const componentId = id.slice(astroEntryPrefix.length);
+				const exportNames = componentToExportNames.get(componentId);
+				if (exportNames) {
+					return `export { ${exportNames.join(', ')} } from ${JSON.stringify(componentId)}`;
+				}
+			}
+		},
+	};
+}
+
+export function normalizeEntryId(id: string): string {
+	return id.startsWith(astroEntryPrefix) ? id.slice(astroEntryPrefix.length) : id;
+}
+
+export function pluginComponentEntry(internals: BuildInternals): AstroBuildPlugin {
+	return {
+		build: 'client',
+		hooks: {
+			'build:before': () => {
+				return {
+					vitePlugin: vitePluginComponentEntry(internals),
+				};
+			},
+		},
+	};
+}

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

@@ -1,6 +1,7 @@
 import type { Plugin as VitePlugin, UserConfig } from 'vite';
 import type { BuildInternals } from '../internal.js';
 import type { AstroBuildPlugin } from '../plugin';
+import { normalizeEntryId } from './plugin-component-entry.js';
 
 export function vitePluginInternals(input: Set<string>, internals: BuildInternals): VitePlugin {
 	return {
@@ -52,7 +53,7 @@ export function vitePluginInternals(input: Set<string>, internals: BuildInternal
 				if (chunk.type === 'chunk' && chunk.facadeModuleId) {
 					const specifiers = mapping.get(chunk.facadeModuleId) || new Set([chunk.facadeModuleId]);
 					for (const specifier of specifiers) {
-						internals.entrySpecifierToBundleMap.set(specifier, chunk.fileName);
+						internals.entrySpecifierToBundleMap.set(normalizeEntryId(specifier), chunk.fileName);
 					}
 				} else if (chunk.type === 'chunk') {
 					for (const id of Object.keys(chunk.modules)) {

packages/astro/src/core/build/static-build.ts

@@ -82,8 +82,8 @@ export async function viteBuild(opts: StaticBuildOptions) {
 		.filter((a) => typeof a === 'string') as string[];
 
 	const clientInput = new Set([
-		...internals.discoveredHydratedComponents,
-		...internals.discoveredClientOnlyComponents,
+		...internals.discoveredHydratedComponents.keys(),
+		...internals.discoveredClientOnlyComponents.keys(),
 		...rendererClientEntrypoints,
 		...internals.discoveredScripts,
 	]);

packages/astro/test/astro-component-bundling.test.js

@@ -0,0 +1,20 @@
+import { expect } from 'chai';
+import { loadFixture } from './test-utils.js';
+
+describe('Component bundling', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({ root: './fixtures/astro-component-bundling/' });
+		await fixture.build();
+	});
+
+	it('should treeshake FooComponent', async () => {
+		const astroChunkDir = await fixture.readdir('/_astro');
+		const manyComponentsChunkName = astroChunkDir.find((chunk) =>
+			chunk.startsWith('ManyComponents')
+		);
+		const manyComponentsChunkContent = await fixture.readFile(`/_astro/${manyComponentsChunkName}`);
+		expect(manyComponentsChunkContent).to.not.include('FooComponent');
+	});
+});

packages/astro/test/fixtures/astro-component-bundling/astro.config.mjs

@@ -0,0 +1,7 @@
+import { defineConfig } from 'astro/config';
+import react from '@astrojs/react';
+
+// https://astro.build/config
+export default defineConfig({
+	integrations: [react()],
+});

packages/astro/test/fixtures/astro-component-bundling/package.json

@@ -0,0 +1,11 @@
+{
+  "name": "@test/astro-component-bundling",
+  "version": "0.0.0",
+  "private": true,
+  "dependencies": {
+    "@astrojs/react": "workspace:*",
+    "astro": "workspace:*",
+    "react": "^18.2.0",
+    "react-dom": "^18.2.0"
+  }
+}

packages/astro/test/fixtures/astro-component-bundling/src/components/ManyComponents.jsx

@@ -0,0 +1,3 @@
+export const FooComponent = () => <div>Foo</div>;
+export const BarComponent = () => <div>Bar</div>;
+export const BazComponent = () => <div>Baz</div>;

packages/astro/test/fixtures/astro-component-bundling/src/pages/index.astro

@@ -0,0 +1,10 @@
+---
+import { BarComponent, BazComponent } from '../components/ManyComponents.jsx'
+---
+<html>
+<head><title>Component bundling</title></head>
+<body>
+  <BarComponent client:idle />
+  <BazComponent client:only="react" />
+</body>
+</html>