1 files changed, 50 insertions, 0 deletions
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)} }`)
+		);
+	};
+}