Add MDX optimize option (#7151)

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

3 files changed, 117 insertions, 0 deletions

diff --git a/packages/integrations/mdx/src/index.ts b/packages/integrations/mdx/src/index.ts
index 0d1ff9d13..e11cd1ac5 100644
--- a/packages/integrations/mdx/src/index.ts
+++ b/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,
 	};
 }
 
diff --git a/packages/integrations/mdx/src/plugins.ts b/packages/integrations/mdx/src/plugins.ts
index af9950451..94c3c10ba 100644
--- a/packages/integrations/mdx/src/plugins.ts
+++ b/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;
 }
 
diff --git a/packages/integrations/mdx/src/rehype-optimize-static.ts b/packages/integrations/mdx/src/rehype-optimize-static.ts
new file mode 100644
index 000000000..c476f7c83
--- /dev/null
+++ b/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 = [];
+		}
+	};
+}