MD/MDX collect headings refactor (#5654)

10 files changed, 212 insertions, 108 deletions

diff --git a/.changeset/fast-baboons-prove.md b/.changeset/fast-baboons-prove.md
new file mode 100644
index 000000000..42cd710b2
--- /dev/null
+++ b/.changeset/fast-baboons-prove.md
@@ -0,0 +1,27 @@
+---
+'@astrojs/mdx': minor
+---
+
+Run heading ID injection after user plugins
+
+⚠️ BREAKING CHANGE ⚠️
+
+If you are using a rehype plugin that depends on heading IDs injected by Astro, the IDs will no longer be available when your plugin runs by default.
+
+To inject IDs before your plugins run, import and add the `rehypeHeadingIds` plugin to your `rehypePlugins` config:
+
+```diff
+// astro.config.mjs
++ import { rehypeHeadingIds } from '@astrojs/markdown-remark';
+import mdx from '@astrojs/mdx';
+
+export default {
+  integrations: [mdx()],
+  markdown: {
+    rehypePlugins: [
++     rehypeHeadingIds,
+      otherPluginThatReliesOnHeadingIDs,
+    ],
+  },
+}
+```
diff --git a/.changeset/violet-mice-push.md b/.changeset/violet-mice-push.md
new file mode 100644
index 000000000..8b111f6c7
--- /dev/null
+++ b/.changeset/violet-mice-push.md
@@ -0,0 +1,7 @@
+---
+'@astrojs/markdown-remark': minor
+---
+
+Refactor and export `rehypeHeadingIds` plugin
+
+The `rehypeHeadingIds` plugin injects IDs for all headings in a Markdown document and can now also handle MDX inputs if needed. You can import and use this plugin if you need heading IDs to be injected _before_ other rehype plugins run.
diff --git a/packages/integrations/mdx/package.json b/packages/integrations/mdx/package.json
index 01bbeb610..08a681097 100644
--- a/packages/integrations/mdx/package.json
+++ b/packages/integrations/mdx/package.json
@@ -30,6 +30,7 @@
     "test:match": "mocha --timeout 20000 -g"
   },
   "dependencies": {
+    "@astrojs/markdown-remark": "^1.1.3",
     "@astrojs/prism": "^1.0.2",
     "@mdx-js/mdx": "^2.1.2",
     "@mdx-js/rollup": "^2.1.1",
diff --git a/packages/integrations/mdx/src/plugins.ts b/packages/integrations/mdx/src/plugins.ts
index 0ae148746..d46ab6cd4 100644
--- a/packages/integrations/mdx/src/plugins.ts
+++ b/packages/integrations/mdx/src/plugins.ts
@@ -1,3 +1,4 @@
+import { rehypeHeadingIds } from '@astrojs/markdown-remark';
 import { nodeTypes } from '@mdx-js/mdx';
 import type { PluggableList } from '@mdx-js/mdx/lib/core.js';
 import type { Options as MdxRollupPluginOptions } from '@mdx-js/rollup';
@@ -10,7 +11,7 @@ import remarkGfm from 'remark-gfm';
 import remarkSmartypants from 'remark-smartypants';
 import type { Data, VFile } from 'vfile';
 import { MdxOptions } from './index.js';
-import rehypeCollectHeadings from './rehype-collect-headings.js';
+import { rehypeInjectHeadingsExport } from './rehype-collect-headings.js';
 import rehypeMetaString from './rehype-meta-string.js';
 import remarkPrism from './remark-prism.js';
 import remarkShiki from './remark-shiki.js';
@@ -153,8 +154,6 @@ export function getRehypePlugins(
 	config: AstroConfig
 ): MdxRollupPluginOptions['rehypePlugins'] {
 	let rehypePlugins: PluggableList = [
-		// getHeadings() is guaranteed by TS, so we can't allow user to override
-		rehypeCollectHeadings,
 		// ensure `data.meta` is preserved in `properties.metastring` for rehype syntax highlighters
 		rehypeMetaString,
 		// rehypeRaw allows custom syntax highlighters to work without added config
@@ -175,7 +174,14 @@ export function getRehypePlugins(
 			break;
 	}
 
-	rehypePlugins = [...rehypePlugins, ...(mdxOptions.rehypePlugins ?? [])];
+	rehypePlugins = [
+		...rehypePlugins,
+		...(mdxOptions.rehypePlugins ?? []),
+		// getHeadings() is guaranteed by TS, so this must be included.
+		// We run `rehypeHeadingIds` _last_ to respect any custom IDs set by user plugins.
+		rehypeHeadingIds,
+		rehypeInjectHeadingsExport,
+	];
 	return rehypePlugins;
 }
 
diff --git a/packages/integrations/mdx/src/rehype-collect-headings.ts b/packages/integrations/mdx/src/rehype-collect-headings.ts
index 64bd7182b..e6cd20a8d 100644
--- a/packages/integrations/mdx/src/rehype-collect-headings.ts
+++ b/packages/integrations/mdx/src/rehype-collect-headings.ts
@@ -1,48 +1,9 @@
-import Slugger from 'github-slugger';
-import { visit } from 'unist-util-visit';
+import { MarkdownVFile, MarkdownHeading } from '@astrojs/markdown-remark';
 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 });
-		});
+export function rehypeInjectHeadingsExport() {
+	return function (tree: any, file: MarkdownVFile) {
+		const headings: MarkdownHeading[] = file.data.__astroHeadings || [];
 		tree.children.unshift(
 			jsToTreeNode(`export function getHeadings() { return ${JSON.stringify(headings)} }`)
 		);
diff --git a/packages/integrations/mdx/test/mdx-get-headings.test.js b/packages/integrations/mdx/test/mdx-get-headings.test.js
index 1ac7283dd..03290abc5 100644
--- a/packages/integrations/mdx/test/mdx-get-headings.test.js
+++ b/packages/integrations/mdx/test/mdx-get-headings.test.js
@@ -1,4 +1,6 @@
+import { rehypeHeadingIds } from '@astrojs/markdown-remark';
 import mdx from '@astrojs/mdx';
+import { visit } from 'unist-util-visit';
 
 import { expect } from 'chai';
 import { parseHTML } from 'linkedom';
@@ -58,3 +60,92 @@ describe('MDX getHeadings', () => {
 		);
 	});
 });
+
+describe('MDX heading IDs can be customized by user plugins', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({
+			root: new URL('./fixtures/mdx-get-headings/', import.meta.url),
+			integrations: [mdx()],
+			markdown: {
+				rehypePlugins: [
+					() => (tree) => {
+						let count = 0;
+						visit(tree, 'element', (node, index, parent) => {
+							if (!/^h\d$/.test(node.tagName)) return;
+							if (!node.properties?.id) {
+								node.properties = { ...node.properties, id: String(count++) };
+							}
+						});
+					},
+				],
+			},
+		});
+
+		await fixture.build();
+	});
+
+	it('adds user-specified IDs to HTML output', async () => {
+		const html = await fixture.readFile('/test/index.html');
+		const { document } = parseHTML(html);
+
+		const h1 = document.querySelector('h1');
+		expect(h1?.textContent).to.equal('Heading test');
+		expect(h1?.getAttribute('id')).to.equal('0');
+
+		const headingIDs = document.querySelectorAll('h1,h2,h3').map((el) => el.id);
+		expect(JSON.stringify(headingIDs)).to.equal(
+			JSON.stringify(Array.from({ length: headingIDs.length }, (_, idx) => String(idx)))
+		);
+	});
+
+	it('generates correct getHeadings() export', async () => {
+		const { headingsByPage } = JSON.parse(await fixture.readFile('/pages.json'));
+		expect(JSON.stringify(headingsByPage['./test.mdx'])).to.equal(
+			JSON.stringify([
+				{ depth: 1, slug: '0', text: 'Heading test' },
+				{ depth: 2, slug: '1', text: 'Section 1' },
+				{ depth: 3, slug: '2', text: 'Subsection 1' },
+				{ depth: 3, slug: '3', text: 'Subsection 2' },
+				{ depth: 2, slug: '4', text: 'Section 2' },
+			])
+		);
+	});
+});
+
+describe('MDX heading IDs can be injected before user plugins', () => {
+	let fixture;
+
+	before(async () => {
+		fixture = await loadFixture({
+			root: new URL('./fixtures/mdx-get-headings/', import.meta.url),
+			integrations: [
+				mdx({
+					rehypePlugins: [
+						rehypeHeadingIds,
+						() => (tree) => {
+							visit(tree, 'element', (node, index, parent) => {
+								if (!/^h\d$/.test(node.tagName)) return;
+								if (node.properties?.id) {
+									node.children.push({ type: 'text', value: ' ' + node.properties.id });
+								}
+							});
+						},
+					],
+				}),
+			],
+		});
+
+		await fixture.build();
+	});
+
+	it('adds user-specified IDs to HTML output', async () => {
+		const html = await fixture.readFile('/test/index.html');
+		const { document } = parseHTML(html);
+
+		const h1 = document.querySelector('h1');
+		expect(h1?.textContent).to.equal('Heading test heading-test');
+		expect(h1?.id).to.equal('heading-test');
+	});
+});
diff --git a/packages/markdown/remark/src/index.ts b/packages/markdown/remark/src/index.ts
index c64bdac0e..6d69bcd20 100644
--- a/packages/markdown/remark/src/index.ts
+++ b/packages/markdown/remark/src/index.ts
@@ -1,7 +1,7 @@
-import type { MarkdownRenderingOptions, MarkdownRenderingResult } from './types';
+import type { MarkdownRenderingOptions, MarkdownRenderingResult, MarkdownVFile } from './types';
 
 import { loadPlugins } from './load-plugins.js';
-import createCollectHeadings from './rehype-collect-headings.js';
+import { rehypeHeadingIds } from './rehype-collect-headings.js';
 import rehypeEscape from './rehype-escape.js';
 import rehypeExpressions from './rehype-expressions.js';
 import rehypeIslands from './rehype-islands.js';
@@ -22,6 +22,7 @@ import markdownToHtml from 'remark-rehype';
 import { unified } from 'unified';
 import { VFile } from 'vfile';
 
+export { rehypeHeadingIds } from './rehype-collect-headings.js';
 export * from './types.js';
 
 export const DEFAULT_REMARK_PLUGINS = ['remark-gfm', 'remark-smartypants'];
@@ -44,7 +45,6 @@ export async function renderMarkdown(
 	} = opts;
 	const input = new VFile({ value: content, path: fileURL });
 	const scopedClassName = opts.$?.scopedClassName;
-	const { headings, rehypeCollectHeadings } = createCollectHeadings();
 
 	let parser = unified()
 		.use(markdown)
@@ -99,12 +99,12 @@ export async function renderMarkdown(
 	parser
 		.use(
 			isAstroFlavoredMd
-				? [rehypeJsx, rehypeExpressions, rehypeEscape, rehypeIslands, rehypeCollectHeadings]
-				: [rehypeCollectHeadings, rehypeRaw]
+				? [rehypeJsx, rehypeExpressions, rehypeEscape, rehypeIslands, rehypeHeadingIds]
+				: [rehypeHeadingIds, rehypeRaw]
 		)
 		.use(rehypeStringify, { allowDangerousHtml: true });
 
-	let vfile: VFile;
+	let vfile: MarkdownVFile;
 	try {
 		vfile = await parser.process(input);
 	} catch (err) {
@@ -116,6 +116,7 @@ export async function renderMarkdown(
 		throw err;
 	}
 
+	const headings = vfile?.data.__astroHeadings || [];
 	return {
 		metadata: { headings, source: content, html: String(vfile.value) },
 		code: String(vfile.value),
diff --git a/packages/markdown/remark/src/rehype-collect-headings.ts b/packages/markdown/remark/src/rehype-collect-headings.ts
index b42ed9030..03a3c6a23 100644
--- a/packages/markdown/remark/src/rehype-collect-headings.ts
+++ b/packages/markdown/remark/src/rehype-collect-headings.ts
@@ -2,72 +2,74 @@ import Slugger from 'github-slugger';
 import { toHtml } from 'hast-util-to-html';
 import { visit } from 'unist-util-visit';
 
-import type { MarkdownHeading, RehypePlugin } from './types.js';
+import type { MarkdownHeading, MarkdownVFile, RehypePlugin } from './types.js';
 
-export default function createCollectHeadings() {
-	const headings: MarkdownHeading[] = [];
-	const slugger = new Slugger();
+const rawNodeTypes = new Set(['text', 'raw', 'mdxTextExpression']);
+const codeTagNames = new Set(['code', 'pre']);
 
-	function rehypeCollectHeadings(): ReturnType<RehypePlugin> {
-		return function (tree) {
-			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);
+export function rehypeHeadingIds(): ReturnType<RehypePlugin> {
+	return function (tree, file: MarkdownVFile) {
+		const headings: MarkdownHeading[] = [];
+		const slugger = new Slugger();
+		const isMDX = isMDXFile(file);
+		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 = '';
-				let isJSX = false;
-				visit(node, (child, __, parent) => {
-					if (child.type === 'element' || parent == null) {
+			let text = '';
+			let isJSX = false;
+			visit(node, (child, __, parent) => {
+				if (child.type === 'element' || parent == null) {
+					return;
+				}
+				if (child.type === 'raw') {
+					if (child.value.match(/^\n?<.*>\n?$/)) {
 						return;
 					}
-					if (child.type === 'raw') {
-						if (child.value.match(/^\n?<.*>\n?$/)) {
-							return;
-						}
-					}
-					if (child.type === 'text' || child.type === 'raw') {
-						if (new Set(['code', 'pre']).has(parent.tagName)) {
-							text += child.value;
-						} else {
-							text += child.value.replace(/\{/g, '${');
-							isJSX = isJSX || child.value.includes('{');
-						}
+				}
+				if (rawNodeTypes.has(child.type)) {
+					if (isMDX || codeTagNames.has(parent.tagName)) {
+						text += child.value;
+					} else {
+						text += child.value.replace(/\{/g, '${');
+						isJSX = isJSX || child.value.includes('{');
 					}
-				});
+				}
+			});
 
-				node.properties = node.properties || {};
-				if (typeof node.properties.id !== 'string') {
-					if (isJSX) {
-						// HACK: serialized JSX from internal plugins, ignore these for slug
-						const raw = toHtml(node.children, { allowDangerousHtml: true })
-							.replace(/\n(<)/g, '<')
-							.replace(/(>)\n/g, '>');
-						// HACK: for ids that have JSX content, use $$slug helper to generate slug at runtime
-						node.properties.id = `$$slug(\`${text}\`)`;
-						(node as any).type = 'raw';
-						(
-							node as any
-						).value = `<${node.tagName} id={${node.properties.id}}>${raw}</${node.tagName}>`;
-					} else {
-						let slug = slugger.slug(text);
+			node.properties = node.properties || {};
+			if (typeof node.properties.id !== 'string') {
+				if (isJSX) {
+					// HACK: serialized JSX from internal plugins, ignore these for slug
+					const raw = toHtml(node.children, { allowDangerousHtml: true })
+						.replace(/\n(<)/g, '<')
+						.replace(/(>)\n/g, '>');
+					// HACK: for ids that have JSX content, use $$slug helper to generate slug at runtime
+					node.properties.id = `$$slug(\`${text}\`)`;
+					(node as any).type = 'raw';
+					(
+						node as any
+					).value = `<${node.tagName} id={${node.properties.id}}>${raw}</${node.tagName}>`;
+				} else {
+					let slug = slugger.slug(text);
 
-						if (slug.endsWith('-')) slug = slug.slice(0, -1);
+					if (slug.endsWith('-')) slug = slug.slice(0, -1);
 
-						node.properties.id = slug;
-					}
+					node.properties.id = slug;
 				}
+			}
 
-				headings.push({ depth, slug: node.properties.id, text });
-			});
-		};
-	}
+			headings.push({ depth, slug: node.properties.id, text });
+		});
 
-	return {
-		headings,
-		rehypeCollectHeadings,
+		file.data.__astroHeadings = headings;
 	};
 }
+
+function isMDXFile(file: MarkdownVFile) {
+	return Boolean(file.history[0]?.endsWith('.mdx'));
+}
diff --git a/packages/markdown/remark/src/types.ts b/packages/markdown/remark/src/types.ts
index 5b40c9f9d..76dfe9b73 100644
--- a/packages/markdown/remark/src/types.ts
+++ b/packages/markdown/remark/src/types.ts
@@ -68,6 +68,12 @@ export interface MarkdownMetadata {
 	html: string;
 }
 
+export interface MarkdownVFile extends VFile {
+	data: {
+		__astroHeadings?: MarkdownHeading[];
+	};
+}
+
 export interface MarkdownRenderingResult {
 	metadata: MarkdownMetadata;
 	vfile: VFile;
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index d391b5cf6..cf302cdf7 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -2880,6 +2880,7 @@ importers:
 
   packages/integrations/mdx:
     specifiers:
+      '@astrojs/markdown-remark': ^1.1.3
       '@astrojs/prism': ^1.0.2
       '@mdx-js/mdx': ^2.1.2
       '@mdx-js/rollup': ^2.1.1
@@ -2916,6 +2917,7 @@ importers:
       vfile: ^5.3.2
       vite: ^3.0.0
     dependencies:
+      '@astrojs/markdown-remark': link:../../markdown/remark
       '@astrojs/prism': link:../../astro-prism
       '@mdx-js/mdx': 2.1.5
       '@mdx-js/rollup': 2.1.5