Add MDX optimize option (#7151)

Co-authored-by: Sarah Rainsberger <[email protected]>

Author: bluwy

.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

@@ -83,6 +83,7 @@ You can configure how your MDX is rendered with the following options:
 - [Options inherited from Markdown config](#options-inherited-from-markdown-config)
 - [`extendMarkdownConfig`](#extendmarkdownconfig)
 - [`recmaPlugins`](#recmaplugins)
+- [`optimize`](#optimize)
 
 ### Options inherited from Markdown config
 
@@ -183,6 +184,71 @@ 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`
+
+- **Type:** `boolean | { customComponentNames?: string[] }`
+
+This is an optional configuration setting to optimize the MDX output for faster builds and rendering via an internal rehype plugin. This may be useful if you have many MDX files and notice slow builds. However, this option may generate some unescaped HTML, so make sure your site's interactive parts still work correctly after enabling it.
+
+This is disabled by default. To enable MDX optimization, add the following to your MDX integration configuration:
+
+__`astro.config.mjs`__
+
+```js
+import { defineConfig } from 'astro/config';
+import mdx from '@astrojs/mdx';
+
+export default defineConfig({
+  integrations: [
+    mdx({
+      optimize: true,
+    })
+  ]
+});
+```
+
+#### `customComponentNames`
+
+- **Type:** `string[]`
+
+An optional property of `optimize` to prevent the MDX optimizer from handling any [custom components passed to imported MDX content via the components prop](https://docs.astro.build/en/guides/markdown-content/#custom-components-with-imported-mdx).
+
+You will need to exclude these components from optimization as the optimizer eagerly converts content into a static string, which will break custom components that needs to be dynamically rendered. 
+
+For example, the intended MDX output of the following is `<Heading>...</Heading>` in place of every `"<h1>...</h1>"`:
+
+```astro
+---
+import { Content, components } from '../content.mdx';
+import Heading from '../Heading.astro';
+---
+
+<Content components={{...components, h1: Heading }} />
+```
+
+To configure optimization for this using the `customComponentNames` property, specify an array of HTML element names that should be treated as custom components:
+
+__`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
+        // These will be treated as custom components
+        customComponentNames: ['h1'],
+      },
+    })
+  ]
+});
+```
+
+Note that if your MDX file [configures custom components using `export const components = { ... }`](https://docs.astro.build/en/guides/markdown-content/#assigning-custom-components-to-html-elements), then you do not need to manually configure this option. The optimizer will automatically detect them.
+
 ## 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

@@ -42,6 +42,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

@@ -11,6 +11,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'> & {
@@ -21,6 +22,7 @@ export type MdxOptions = Omit<typeof markdownConfigDefaults, 'remarkPlugins' | '
 	remarkPlugins: PluggableList;
 	rehypePlugins: PluggableList;
 	remarkRehype: RemarkRehypeOptions;
+	optimize: boolean | OptimizeOptions;
 };
 
 type SetupHookParams = HookParameters<'astro:config:setup'> & {
@@ -194,6 +196,7 @@ function markdownConfigToMdxOptions(markdownConfig: typeof markdownConfigDefault
 		remarkPlugins: ignoreStringPlugins(markdownConfig.remarkPlugins),
 		rehypePlugins: ignoreStringPlugins(markdownConfig.rehypePlugins),
 		remarkRehype: (markdownConfig.remarkRehype as any) ?? {},
+		optimize: false,
 	};
 }
 
@@ -214,6 +217,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

@@ -15,6 +15,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';
@@ -144,6 +145,13 @@ export function getRehypePlugins(mdxOptions: MdxOptions): PluggableList {
 		// computed from `astro.data.frontmatter` in VFile data
 		rehypeApplyFrontmatterExport,
 	];
+
+	if (mdxOptions.optimize) {
+		// Convert user `optimize` option to compatible `rehypeOptimizeStatic` option
+		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,105 @@
+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) => {
+		// A set of non-static components to avoid collapsing when walking the tree
+		// as they need to be preserved as JSX to be rendered dynamically.
+		const customComponentNames = new Set<string>(options?.customComponentNames);
+
+		// Find `export const components = { ... }` and get it's object's keys to be
+		// populated into `customComponentNames`. This configuration is used to render
+		// some HTML elements as custom components, and we also want to avoid collapsing them.
+		for (const child of tree.children) {
+			if (child.type === 'mdxjsEsm' && exportConstComponentsRe.test(child.value)) {
+				// Try to loosely get the object property nodes
+				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 in `elementStack` and
+				// `allPossibleElements` to be used in the "leave" hook below.
+				if (node.type === 'element' || node.type === 'mdxJsxFlowElement') {
+					elementStack.push(node);
+					allPossibleElements.add(node);
+				}
+			},
+			leave(node, _, __, parents) {
+				// Do the reverse of the if condition above, popping the `elementStack`,
+				// and consolidating `allPossibleElements` as a subtree root.
+				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>