Optimize static mdx (withastro#3184)

astro.config.ts

@@ -13,6 +13,7 @@ import { sitemap } from './integrations/sitemap';
 import { autolinkConfig } from './plugins/rehype-autolink-config';
 import { rehypei18nAutolinkHeadings } from './plugins/rehype-i18n-autolink-headings';
 import { rehypeTasklistEnhancer } from './plugins/rehype-tasklist-enhancer';
+import { rehypeOptimizeStatic } from './plugins/rehype-optimize-static';
 import { remarkFallbackLang } from './plugins/remark-fallback-lang';
 import { theme } from './syntax-highlighting-theme';
 
@@ -47,6 +48,8 @@ export default defineConfig({
 			rehypeTasklistEnhancer(),
 			// Translates the autolink headings anchors
 			rehypei18nAutolinkHeadings(),
+			// Collapse static parts of the hast to html
+			rehypeOptimizeStatic,
 		],
 	},
 });

package.json

@@ -60,6 +60,7 @@
     "eslint-plugin-react": "^7.32.1",
     "fast-glob": "^3.2.11",
     "hast-util-from-html": "^1.0.0",
+    "hast-util-to-html": "^8.0.4",
     "hast-util-to-string": "^2.0.0",
     "hastscript": "^7.0.2",
     "html-escaper": "^3.0.3",
@@ -88,6 +89,7 @@
     "unified": "^10.1.2",
     "unist-util-remove": "^3.1.0",
     "unist-util-visit": "^4.1.0",
+    "unist-util-walker": "^1.0.0",
     "vfile": "^5.3.6",
     "vitest": "^0.28.5"
   },

plugins/rehype-optimize-static.ts

@@ -0,0 +1,79 @@
+import type { Root } from 'hast';
+import type { Transformer } from 'unified';
+import { walk } from 'unist-util-walker';
+import { toHtml } from 'hast-util-to-html';
+
+// accessing untyped hast and mdx types
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type Node = any;
+
+const headingRe = /h([0-6])/;
+
+/**
+ * For MDX only, collapse static subtrees of the hast into `set:html`. Subtrees
+ * do not include any MDX elements or headings (for `rehypeHeadingIds` to work).
+ * 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(): Transformer<Root, Root> {
+	return (tree) => {
+		// 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[] = [];
+
+		walk(tree, {
+			enter(node) {
+				// @ts-expect-error test tagName naively
+				const isHeading = node.tagName && headingRe.test(node.tagName);
+				// For nodes that can't be optimized, eliminate all elements in the
+				// `elementStack` from the `allPossibleElements` set.
+				if (node.type.startsWith('mdx') || isHeading) {
+					for (const el of elementStack) {
+						allPossibleElements.delete(el);
+					}
+				}
+				// If is heading node, skip it and its children. This prevents the content
+				// from being optimized, as the content is used to generate the heading text.
+				if (isHeading) {
+					this.skip();
+					return;
+				}
+				// For possible subtree root nodes, record them
+				if (node.type === 'element' || node.type === 'mdxJsxFlowElement') {
+					elementStack.push(node);
+					allPossibleElements.add(node);
+				}
+			},
+			leave(node, parent) {
+				// 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.
+					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 = [];
+		}
+	};
+}