MD/MDX collect headings refactor (#5654)

2 files changed, 14 insertions, 47 deletions

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)} }`)
 		);