astro/packages/markdown/remark/src/rehype-collect-headings.ts

import Slugger from 'github-slugger';
import { toHtml } from 'hast-util-to-html';
import { visit } from 'unist-util-visit';

import type { MarkdownHeading, MarkdownVFile, RehypePlugin } from './types.js';

const rawNodeTypes = new Set(['text', 'raw', 'mdxTextExpression']);
const codeTagNames = new Set(['code', 'pre']);

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) {
					return;
				}
				if (child.type === 'raw') {
					if (child.value.match(/^\n?<.*>\n?$/)) {
						return;
					}
				}
				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);

					if (slug.endsWith('-')) slug = slug.slice(0, -1);

					node.properties.id = slug;
				}
			}

			headings.push({ depth, slug: node.properties.id, text });
		});

		file.data.__astroHeadings = headings;
	};
}

function isMDXFile(file: MarkdownVFile) {
	return Boolean(file.history[0]?.endsWith('.mdx'));
}