[MDX] Add `getHeadings` + generate anchor links (#4095)

* deps: mdx github-slugger

* feat: add getHeadings via rehype plugin

* chore: stray console.log

* test: getHeadings w/ & w/0 JSX expressions

* docs: add generated exports

* refactor: pass headings using vfile.data

* deps: vfile

* test: heading anchor IDs

* docs: add collect-headings to default rehype plugins

* chore: changeset

* deps: estree-util-value-to-estree

* refactor: inject getHeadings export the right way!

* deps: switch to acorn

* refactor: just use acorn

* docs: `getHeadings` info structuring

Co-authored-by: Chris Swithinbank <swithinbank@gmail.com>

* docs: clarify `url` example

Co-authored-by: Chris Swithinbank <swithinbank@gmail.com>

* fix: move slugger inside plugin call

* refactor: cleanup code reassignment

* chore: lint

* deps: mdast-util-mdx, test utils

* refactor: add jsToTreeNode util

* feat: expose utils for lib authors

* test: rehype plugins w/ and w/o extends

* test: fixture

* refactor: remove utils from package exports

Co-authored-by: Chris Swithinbank <swithinbank@gmail.com>

3 files changed, 109 insertions, 22 deletions

diff --git a/packages/integrations/mdx/src/index.ts b/packages/integrations/mdx/src/index.ts
index 8b7831f27..f7365b505 100644
--- a/packages/integrations/mdx/src/index.ts
+++ b/packages/integrations/mdx/src/index.ts
@@ -1,4 +1,5 @@
-import { nodeTypes } from '@mdx-js/mdx';
+import { nodeTypes, compile as mdxCompile } from '@mdx-js/mdx';
+import { VFile } from 'vfile';
 import mdxPlugin, { Options as MdxRollupPluginOptions } from '@mdx-js/rollup';
 import type { AstroIntegration } from 'astro';
 import { parse as parseESM } from 'es-module-lexer';
@@ -12,6 +13,7 @@ import remarkSmartypants from 'remark-smartypants';
 import type { Plugin as VitePlugin } from 'vite';
 import remarkPrism from './remark-prism.js';
 import { getFileInfo, getFrontmatter } from './utils.js';
+import rehypeCollectHeadings from './rehype-collect-headings.js';
 
 type WithExtends<T> = T | { extends: T };
 
@@ -27,6 +29,7 @@ type MdxOptions = {
 };
 
 const DEFAULT_REMARK_PLUGINS = [remarkGfm, remarkSmartypants];
+const DEFAULT_REHYPE_PLUGINS = [rehypeCollectHeadings];
 
 function handleExtends<T>(config: WithExtends<T[] | undefined>, defaults: T[] = []): T[] {
 	if (Array.isArray(config)) return config;
@@ -41,7 +44,7 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 			'astro:config:setup': ({ updateConfig, config, addPageExtension, command }: any) => {
 				addPageExtension('.mdx');
 				let remarkPlugins = handleExtends(mdxOptions.remarkPlugins, DEFAULT_REMARK_PLUGINS);
-				let rehypePlugins = handleExtends(mdxOptions.rehypePlugins);
+				let rehypePlugins = handleExtends(mdxOptions.rehypePlugins, DEFAULT_REHYPE_PLUGINS);
 
 				if (config.markdown.syntaxHighlight === 'shiki') {
 					remarkPlugins.push([
@@ -69,7 +72,7 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 					},
 				]);
 
-				const configuredMdxPlugin = mdxPlugin({
+				const mdxPluginOpts: MdxRollupPluginOptions = {
 					remarkPlugins,
 					rehypePlugins,
 					jsx: true,
@@ -77,38 +80,47 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 					// Note: disable `.md` support
 					format: 'mdx',
 					mdExtensions: [],
-				});
+				};
 
 				updateConfig({
 					vite: {
 						plugins: [
 							{
 								enforce: 'pre',
-								...configuredMdxPlugin,
-								// Override transform to inject layouts before MDX compilation
-								async transform(this, code, id) {
-									if (!id.endsWith('.mdx')) return;
+								...mdxPlugin(mdxPluginOpts),
+								// Override transform to alter code before MDX compilation
+								// ex. inject layouts
+								async transform(code, id) {
+									if (!id.endsWith('mdx')) return;
 
-									const mdxPluginTransform = configuredMdxPlugin.transform?.bind(this);
 									// If user overrides our default YAML parser,
 									// do not attempt to parse the `layout` via gray-matter
-									if (mdxOptions.frontmatterOptions?.parsers) {
-										return mdxPluginTransform?.(code, id);
-									}
-									const frontmatter = getFrontmatter(code, id);
-									if (frontmatter.layout) {
-										const { layout, ...content } = frontmatter;
-										code += `\nexport default async function({ children }) {\nconst Layout = (await import(${JSON.stringify(
-											frontmatter.layout
-										)})).default;\nreturn <Layout content={${JSON.stringify(
-											content
-										)}}>{children}</Layout> }`;
+									if (!mdxOptions.frontmatterOptions?.parsers) {
+										const frontmatter = getFrontmatter(code, id);
+										if (frontmatter.layout) {
+											const { layout, ...content } = frontmatter;
+											code += `\nexport default async function({ children }) {\nconst Layout = (await import(${JSON.stringify(
+												frontmatter.layout
+											)})).default;\nreturn <Layout content={${JSON.stringify(
+												content
+											)}}>{children}</Layout> }`;
+										}
 									}
-									return mdxPluginTransform?.(code, id);
+
+									const compiled = await mdxCompile(
+										new VFile({ value: code, path: id }),
+										mdxPluginOpts
+									);
+
+									return {
+										code: String(compiled.value),
+										map: compiled.map,
+									};
 								},
 							},
 							{
-								name: '@astrojs/mdx',
+								name: '@astrojs/mdx-postprocess',
+								// These transforms must happen *after* JSX runtime transformations
 								transform(code, id) {
 									if (!id.endsWith('.mdx')) return;
 									const [, moduleExports] = parseESM(code);
diff --git a/packages/integrations/mdx/src/rehype-collect-headings.ts b/packages/integrations/mdx/src/rehype-collect-headings.ts
new file mode 100644
index 000000000..64bd7182b
--- /dev/null
+++ b/packages/integrations/mdx/src/rehype-collect-headings.ts
@@ -0,0 +1,50 @@
+import Slugger from 'github-slugger';
+import { visit } from 'unist-util-visit';
+import { jsToTreeNode } from './utils.js';
+
+export interface MarkdownHeading {
+	depth: number;
+	slug: string;
+	text: string;
+}
+
+export default function rehypeCollectHeadings() {
+	const slugger = new Slugger();
+	return function (tree: any) {
+		const headings: MarkdownHeading[] = [];
+		visit(tree, (node) => {
+			if (node.type !== 'element') return;
+			const { tagName } = node;
+			if (tagName[0] !== 'h') return;
+			const [_, level] = tagName.match(/h([0-6])/) ?? [];
+			if (!level) return;
+			const depth = Number.parseInt(level);
+
+			let text = '';
+			visit(node, (child, __, parent) => {
+				if (child.type === 'element' || parent == null) {
+					return;
+				}
+				if (child.type === 'raw' && child.value.match(/^\n?<.*>\n?$/)) {
+					return;
+				}
+				if (new Set(['text', 'raw', 'mdxTextExpression']).has(child.type)) {
+					text += child.value;
+				}
+			});
+
+			node.properties = node.properties || {};
+			if (typeof node.properties.id !== 'string') {
+				let slug = slugger.slug(text);
+				if (slug.endsWith('-')) {
+					slug = slug.slice(0, -1);
+				}
+				node.properties.id = slug;
+			}
+			headings.push({ depth, slug: node.properties.id, text });
+		});
+		tree.children.unshift(
+			jsToTreeNode(`export function getHeadings() { return ${JSON.stringify(headings)} }`)
+		);
+	};
+}
diff --git a/packages/integrations/mdx/src/utils.ts b/packages/integrations/mdx/src/utils.ts
index ccce179c9..7b2c4f4ec 100644
--- a/packages/integrations/mdx/src/utils.ts
+++ b/packages/integrations/mdx/src/utils.ts
@@ -1,4 +1,8 @@
 import type { AstroConfig, SSRError } from 'astro';
+import type { Options as AcornOpts } from 'acorn';
+import type { MdxjsEsm } from 'mdast-util-mdx';
+import { parse } from 'acorn';
+
 import matter from 'gray-matter';
 
 function appendForwardSlash(path: string) {
@@ -58,3 +62,24 @@ export function getFrontmatter(code: string, id: string) {
 		}
 	}
 }
+
+export function jsToTreeNode(
+	jsString: string,
+	acornOpts: AcornOpts = {
+		ecmaVersion: 'latest',
+		sourceType: 'module',
+	}
+): MdxjsEsm {
+	return {
+		type: 'mdxjsEsm',
+		value: '',
+		data: {
+			estree: {
+				body: [],
+				...parse(jsString, acornOpts),
+				type: 'Program',
+				sourceType: 'module',
+			},
+		},
+	};
+}