Add MDX optimize option

.changeset/dirty-singers-enjoy.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/mdx': patch
+---
+
+Add `optimize` option for faster builds and rendering

packages/integrations/mdx/README.md

@@ -183,6 +183,42 @@ These are plugins that modify the output [estree](https://github.com/estree/estr
 
 We suggest [using AST Explorer](https://astexplorer.net/) to play with estree outputs, and trying [`estree-util-visit`](https://unifiedjs.com/explore/package/estree-util-visit/) for searching across JavaScript nodes.
 
+### `optimize`
+
+Whether to optimize the MDX output for faster builds and rendering. An internal rehype plugin will be injected last to transform static parts of the `hast` into [`set:html` properties](https://docs.astro.build/en/reference/directives-reference/#sethtml). As this modifies the intermediate `estree` and leaves some generated HTML to be un-escaped, the option is disabled by default.
+
+As the optimizer works within individual MDX files, if you're [passing custom components to an imported MDX content](https://docs.astro.build/en/guides/markdown-content/#custom-components-with-imported-mdx), you'll need to configure `customComponentNames` to prevent the optimizer from handling the custom components:
+
+```astro
+---
+import { Content, components } from '../content.mdx';
+import Heading from '../Heading.astro';
+---
+
+<Content components={{...components, h1: Heading }} />
+```
+
+__`astro.config.mjs`__
+
+```js
+import { defineConfig } from 'astro/config';
+import mdx from '@astrojs/mdx';
+
+export default defineConfig({
+  integrations: [
+    mdx({
+      optimize: {
+        // Prevent the optimizer from handling `h1` elements
+        customComponentNames: ['h1'],
+      },
+    })
+  ]
+});
+```
+
+> **Note**
+> If you're using `export const components = { ... }` in your MDX file, the optimizer will automatically detect and add them to `customComponentNames` so you don't have to manually configure it.
+
 ## Examples
 
 *   The [Astro MDX starter template](https://github.com/withastro/astro/tree/latest/examples/with-mdx) shows how to use MDX files in your Astro project.

packages/integrations/mdx/package.json

@@ -43,6 +43,7 @@
     "estree-util-visit": "^1.2.0",
     "github-slugger": "^1.4.0",
     "gray-matter": "^4.0.3",
+    "hast-util-to-html": "^8.0.4",
     "kleur": "^4.1.4",
     "rehype-raw": "^6.1.1",
     "remark-frontmatter": "^4.0.1",

packages/integrations/mdx/src/index.ts

@@ -12,6 +12,7 @@ import { SourceMapGenerator } from 'source-map';
 import { VFile } from 'vfile';
 import type { Plugin as VitePlugin } from 'vite';
 import { getRehypePlugins, getRemarkPlugins, recmaInjectImportMetaEnvPlugin } from './plugins.js';
+import type { OptimizeOptions } from './rehype-optimize-static.js';
 import { getFileInfo, ignoreStringPlugins, parseFrontmatter } from './utils.js';
 
 export type MdxOptions = Omit<typeof markdownConfigDefaults, 'remarkPlugins' | 'rehypePlugins'> & {
@@ -22,6 +23,7 @@ export type MdxOptions = Omit<typeof markdownConfigDefaults, 'remarkPlugins' | '
 	remarkPlugins: PluggableList;
 	rehypePlugins: PluggableList;
 	remarkRehype: RemarkRehypeOptions;
+	optimize: boolean | OptimizeOptions;
 };
 
 type SetupHookParams = HookParameters<'astro:config:setup'> & {
@@ -195,6 +197,7 @@ function markdownConfigToMdxOptions(markdownConfig: typeof markdownConfigDefault
 		remarkPlugins: ignoreStringPlugins(markdownConfig.remarkPlugins),
 		rehypePlugins: ignoreStringPlugins(markdownConfig.rehypePlugins),
 		remarkRehype: (markdownConfig.remarkRehype as any) ?? {},
+		optimize: false,
 	};
 }
 
@@ -215,6 +218,7 @@ function applyDefaultOptions({
 		remarkPlugins: options.remarkPlugins ?? defaults.remarkPlugins,
 		rehypePlugins: options.rehypePlugins ?? defaults.rehypePlugins,
 		shikiConfig: options.shikiConfig ?? defaults.shikiConfig,
+		optimize: options.optimize ?? defaults.optimize,
 	};
 }
 

packages/integrations/mdx/src/plugins.ts

@@ -16,6 +16,7 @@ import type { VFile } from 'vfile';
 import type { MdxOptions } from './index.js';
 import { rehypeInjectHeadingsExport } from './rehype-collect-headings.js';
 import rehypeMetaString from './rehype-meta-string.js';
+import { rehypeOptimizeStatic } from './rehype-optimize-static.js';
 import { remarkImageToComponent } from './remark-images-to-component.js';
 import remarkPrism from './remark-prism.js';
 import remarkShiki from './remark-shiki.js';
@@ -145,6 +146,12 @@ export function getRehypePlugins(mdxOptions: MdxOptions): MdxRollupPluginOptions
 		// computed from `astro.data.frontmatter` in VFile data
 		rehypeApplyFrontmatterExport,
 	];
+
+	if (mdxOptions.optimize) {
+		const options = mdxOptions.optimize === true ? undefined : mdxOptions.optimize;
+		rehypePlugins.push([rehypeOptimizeStatic, options]);
+	}
+
 	return rehypePlugins;
 }
 

packages/integrations/mdx/src/rehype-optimize-static.ts

@@ -0,0 +1,97 @@
+import { visit } from 'estree-util-visit';
+import { toHtml } from 'hast-util-to-html';
+
+// accessing untyped hast and mdx types
+type Node = any;
+
+export interface OptimizeOptions {
+	customComponentNames?: string[];
+}
+
+const exportConstComponentsRe = /export\s+const\s+components\s*=/;
+
+/**
+ * For MDX only, collapse static subtrees of the hast into `set:html`. Subtrees
+ * do not include any MDX elements.
+ *
+ * This optimization reduces the JS output as more content are represented as a
+ * string instead, which also reduces the AST size that Rollup holds in memory.
+ */
+export function rehypeOptimizeStatic(options?: OptimizeOptions) {
+	return (tree: any) => {
+		// Read `export const components` to make sure the components are not optimized.
+		const customComponentNames = new Set<string>(options?.customComponentNames);
+		for (const child of tree.children) {
+			if (child.type === 'mdxjsEsm' && exportConstComponentsRe.test(child.value)) {
+				const objectPropertyNodes = child.data.estree.body[0]?.declarations?.[0]?.init?.properties;
+				if (objectPropertyNodes) {
+					for (const objectPropertyNode of objectPropertyNodes) {
+						const componentName = objectPropertyNode.key?.name ?? objectPropertyNode.key?.value;
+						if (componentName) {
+							customComponentNames.add(componentName);
+						}
+					}
+				}
+			}
+		}
+
+		// All possible elements that could be the root of a subtree
+		const allPossibleElements = new Set<Node>();
+		// The current collapsible element stack while traversing the tree
+		const elementStack: Node[] = [];
+
+		visit(tree, {
+			enter(node) {
+				// @ts-expect-error read tagName naively
+				const isCustomComponent = node.tagName && customComponentNames.has(node.tagName);
+				// For nodes that can't be optimized, eliminate all elements in the
+				// `elementStack` from the `allPossibleElements` set.
+				if (node.type.startsWith('mdx') || isCustomComponent) {
+					for (const el of elementStack) {
+						allPossibleElements.delete(el);
+					}
+					// Micro-optimization: While this destroys the meaning of an element
+					// stack for this node, things will still work but we won't repeatedly
+					// run the above for other nodes anymore. If this is confusing, you can
+					// comment out the code below when reading.
+					elementStack.length = 0;
+				}
+				// For possible subtree root nodes, record them
+				if (node.type === 'element' || node.type === 'mdxJsxFlowElement') {
+					elementStack.push(node);
+					allPossibleElements.add(node);
+				}
+			},
+			leave(node, _, __, parents) {
+				// Similar as above, but pop the `elementStack`
+				if (node.type === 'element' || node.type === 'mdxJsxFlowElement') {
+					elementStack.pop();
+					// Many possible elements could be part of a subtree, in order to find
+					// the root, we check the parent of the element we're popping. If the
+					// parent exists in `allPossibleElements`, then we're definitely not
+					// the root, so remove ourselves. This will work retroactively as we
+					// climb back up the tree.
+					const parent = parents[parents.length - 1];
+					if (allPossibleElements.has(parent)) {
+						allPossibleElements.delete(node);
+					}
+				}
+			},
+		});
+
+		// For all possible subtree roots, collapse them into `set:html` and
+		// strip of their children
+		for (const el of allPossibleElements) {
+			if (el.type === 'mdxJsxFlowElement') {
+				el.attributes.push({
+					type: 'mdxJsxAttribute',
+					name: 'set:html',
+					value: toHtml(el.children),
+				});
+			} else {
+				el.properties['set:html'] = toHtml(el.children);
+			}
+			el.children = [];
+		}
+	};
+}

packages/integrations/mdx/test/fixtures/mdx-optimize/astro.config.mjs

@@ -0,0 +1,9 @@
+import mdx from '@astrojs/mdx';
+
+export default {
+	integrations: [mdx({
+		optimize: {
+			customComponentNames: ['strong']
+		}
+	})]
+}

packages/integrations/mdx/test/fixtures/mdx-optimize/package.json

@@ -0,0 +1,8 @@
+{
+  "name": "@test/mdx-optimize",
+  "private": true,
+  "dependencies": {
+    "@astrojs/mdx": "workspace:*",
+    "astro": "workspace:*"
+  }
+}

packages/integrations/mdx/test/fixtures/mdx-optimize/src/components/Blockquote.astro

@@ -0,0 +1,3 @@
+<blockquote {...Astro.props} class="custom-blockquote">
+  <slot />
+</blockquote>

packages/integrations/mdx/test/fixtures/mdx-optimize/src/components/Strong.astro

@@ -0,0 +1,3 @@
+<strong {...Astro.props} class="custom-strong">
+  <slot />
+</strong>

packages/integrations/mdx/test/fixtures/mdx-optimize/src/pages/_imported.mdx

@@ -0,0 +1,3 @@
+I once heard a very **inspirational** quote:
+
+> I like pancakes

packages/integrations/mdx/test/fixtures/mdx-optimize/src/pages/import.astro

@@ -0,0 +1,15 @@
+---
+import { Content, components } from './index.mdx'
+import Strong from '../components/Strong.astro'
+---
+
+<!DOCTYPE html>
+<html>
+  <head>
+    <title>Import MDX component</title>
+  </head>
+  <body>
+    <h1>Astro page</h1>
+    <Content components={{ ...components, strong: Strong }} />
+  </body>
+</html>

packages/integrations/mdx/test/fixtures/mdx-optimize/src/pages/index.mdx

@@ -0,0 +1,15 @@
+import Blockquote from '../components/Blockquote.astro'
+
+export const components = {
+  blockquote: Blockquote
+}
+
+# MDX page
+
+I once heard a very inspirational quote:
+
+> I like pancakes
+
+```js
+const pancakes = 'yummy'
+```

packages/integrations/mdx/test/mdx-optimize.test.js

@@ -0,0 +1,47 @@
+import { expect } from 'chai';
+import { parseHTML } from 'linkedom';
+import { loadFixture } from '../../../astro/test/test-utils.js';
+
+const FIXTURE_ROOT = new URL('./fixtures/mdx-optimize/', import.meta.url);
+
+describe('MDX optimize', () => {
+	let fixture;
+	before(async () => {
+		fixture = await loadFixture({
+			root: FIXTURE_ROOT,
+		});
+		await fixture.build();
+	});
+
+	it('renders an MDX page fine', async () => {
+		const html = await fixture.readFile('/index.html');
+		const { document } = parseHTML(html);
+
+		expect(document.querySelector('h1').textContent).include('MDX page');
+		expect(document.querySelector('p').textContent).include(
+			'I once heard a very inspirational quote:'
+		);
+
+		const blockquote = document.querySelector('blockquote.custom-blockquote');
+		expect(blockquote).to.not.be.null;
+		expect(blockquote.textContent).to.include('I like pancakes');
+
+		const code = document.querySelector('pre.astro-code');
+		expect(code).to.not.be.null;
+		expect(code.textContent).to.include(`const pancakes = 'yummy'`);
+	});
+
+	it('renders an Astro page that imports MDX fine', async () => {
+		const html = await fixture.readFile('/import/index.html');
+		const { document } = parseHTML(html);
+
+		expect(document.querySelector('h1').textContent).include('Astro page');
+		expect(document.querySelector('p').textContent).include(
+			'I once heard a very inspirational quote:'
+		);
+
+		const blockquote = document.querySelector('blockquote.custom-blockquote');
+		expect(blockquote).to.not.be.null;
+		expect(blockquote.textContent).to.include('I like pancakes');
+	});
+});