[MDX] Switch from Shiki Twoslash -> Astro Markdown highlighter (#4292)

* freat: twoslash -> Astro shiki parser

* test: update shiki style check

* feat: always apply rehypeRaw

* deps: move remark-shiki-twoslash to dev

* test: add shiki-twoslash test

* docs: update readme with twoslash example

* chore: changeset

* nit: remove "describe('disabled')"

2 files changed, 95 insertions, 15 deletions

diff --git a/packages/integrations/mdx/src/index.ts b/packages/integrations/mdx/src/index.ts
index 17fe0cd74..72fbbeb6c 100644
--- a/packages/integrations/mdx/src/index.ts
+++ b/packages/integrations/mdx/src/index.ts
@@ -4,13 +4,13 @@ import type { AstroConfig, AstroIntegration } from 'astro';
 import { parse as parseESM } from 'es-module-lexer';
 import rehypeRaw from 'rehype-raw';
 import remarkGfm from 'remark-gfm';
-import remarkShikiTwoslash from 'remark-shiki-twoslash';
 import remarkSmartypants from 'remark-smartypants';
 import { VFile } from 'vfile';
 import type { Plugin as VitePlugin } from 'vite';
 import { rehypeApplyFrontmatterExport, remarkInitializeAstroData } from './astro-data-utils.js';
 import rehypeCollectHeadings from './rehype-collect-headings.js';
 import remarkPrism from './remark-prism.js';
+import remarkShiki from './remark-shiki.js';
 import { getFileInfo, parseFrontmatter } from './utils.js';
 
 type WithExtends<T> = T | { extends: T };
@@ -38,22 +38,17 @@ function handleExtends<T>(config: WithExtends<T[] | undefined>, defaults: T[] =
 	return [...defaults, ...(config?.extends ?? [])];
 }
 
-function getRemarkPlugins(
+async function getRemarkPlugins(
 	mdxOptions: MdxOptions,
 	config: AstroConfig
-): MdxRollupPluginOptions['remarkPlugins'] {
+): Promise<MdxRollupPluginOptions['remarkPlugins']> {
 	let remarkPlugins = [
 		// Initialize vfile.data.astroExports before all plugins are run
 		remarkInitializeAstroData,
 		...handleExtends(mdxOptions.remarkPlugins, DEFAULT_REMARK_PLUGINS),
 	];
 	if (config.markdown.syntaxHighlight === 'shiki') {
-		// Default export still requires ".default" chaining for some reason
-		// Workarounds tried:
-		// - "import * as remarkShikiTwoslash"
-		// - "import { default as remarkShikiTwoslash }"
-		const shikiTwoslash = (remarkShikiTwoslash as any).default ?? remarkShikiTwoslash;
-		remarkPlugins.push([shikiTwoslash, config.markdown.shikiConfig]);
+		remarkPlugins.push([await remarkShiki(config.markdown.shikiConfig)]);
 	}
 	if (config.markdown.syntaxHighlight === 'prism') {
 		remarkPlugins.push(remarkPrism);
@@ -65,11 +60,11 @@ function getRehypePlugins(
 	mdxOptions: MdxOptions,
 	config: AstroConfig
 ): MdxRollupPluginOptions['rehypePlugins'] {
-	let rehypePlugins = handleExtends(mdxOptions.rehypePlugins, DEFAULT_REHYPE_PLUGINS);
+	let rehypePlugins = [
+		[rehypeRaw, { passThrough: nodeTypes }] as any,
+		...handleExtends(mdxOptions.rehypePlugins, DEFAULT_REHYPE_PLUGINS),
+	];
 
-	if (config.markdown.syntaxHighlight === 'shiki' || config.markdown.syntaxHighlight === 'prism') {
-		rehypePlugins.unshift([rehypeRaw, { passThrough: nodeTypes }]);
-	}
 	// getHeadings() is guaranteed by TS, so we can't allow user to override
 	rehypePlugins.unshift(rehypeCollectHeadings);
 
@@ -80,11 +75,11 @@ export default function mdx(mdxOptions: MdxOptions = {}): AstroIntegration {
 	return {
 		name: '@astrojs/mdx',
 		hooks: {
-			'astro:config:setup': ({ updateConfig, config, addPageExtension, command }: any) => {
+			'astro:config:setup': async ({ updateConfig, config, addPageExtension, command }: any) => {
 				addPageExtension('.mdx');
 
 				const mdxPluginOpts: MdxRollupPluginOptions = {
-					remarkPlugins: getRemarkPlugins(mdxOptions, config),
+					remarkPlugins: await getRemarkPlugins(mdxOptions, config),
 					rehypePlugins: getRehypePlugins(mdxOptions, config),
 					jsx: true,
 					jsxImportSource: 'astro',
diff --git a/packages/integrations/mdx/src/remark-shiki.ts b/packages/integrations/mdx/src/remark-shiki.ts
new file mode 100644
index 000000000..76a1d275f
--- /dev/null
+++ b/packages/integrations/mdx/src/remark-shiki.ts
@@ -0,0 +1,85 @@
+import type { ShikiConfig } from 'astro';
+import type * as shiki from 'shiki';
+import { getHighlighter } from 'shiki';
+import { visit } from 'unist-util-visit';
+
+/**
+ * getHighlighter() is the most expensive step of Shiki. Instead of calling it on every page,
+ * cache it here as much as possible. Make sure that your highlighters can be cached, state-free.
+ * We make this async, so that multiple calls to parse markdown still share the same highlighter.
+ */
+const highlighterCacheAsync = new Map<string, Promise<shiki.Highlighter>>();
+
+const remarkShiki = async ({ langs = [], theme = 'github-dark', wrap = false }: ShikiConfig) => {
+	const cacheID: string = typeof theme === 'string' ? theme : theme.name;
+	let highlighterAsync = highlighterCacheAsync.get(cacheID);
+	if (!highlighterAsync) {
+		highlighterAsync = getHighlighter({ theme });
+		highlighterCacheAsync.set(cacheID, highlighterAsync);
+	}
+	const highlighter = await highlighterAsync;
+
+	// NOTE: There may be a performance issue here for large sites that use `lang`.
+	// Since this will be called on every page load. Unclear how to fix this.
+	for (const lang of langs) {
+		await highlighter.loadLanguage(lang);
+	}
+
+	return () => (tree: any) => {
+		visit(tree, 'code', (node) => {
+			let lang: string;
+
+			if (typeof node.lang === 'string') {
+				const langExists = highlighter.getLoadedLanguages().includes(node.lang);
+				if (langExists) {
+					lang = node.lang;
+				} else {
+					// eslint-disable-next-line no-console
+					console.warn(`The language "${node.lang}" doesn't exist, falling back to plaintext.`);
+					lang = 'plaintext';
+				}
+			} else {
+				lang = 'plaintext';
+			}
+
+			let html = highlighter!.codeToHtml(node.value, { lang });
+
+			// Q: Couldn't these regexes match on a user's inputted code blocks?
+			// A: Nope! All rendered HTML is properly escaped.
+			// Ex. If a user typed `<span class="line"` into a code block,
+			// It would become this before hitting our regexes:
+			// &lt;span class=&quot;line&quot;
+
+			// Replace "shiki" class naming with "astro".
+			html = html.replace('<pre class="shiki"', `<pre class="astro-code"`);
+			// Replace "shiki" css variable naming with "astro".
+			html = html.replace(
+				/style="(background-)?color: var\(--shiki-/g,
+				'style="$1color: var(--astro-code-'
+			);
+			// Add "user-select: none;" for "+"/"-" diff symbols
+			if (node.lang === 'diff') {
+				html = html.replace(
+					/<span class="line"><span style="(.*?)">([\+|\-])/g,
+					'<span class="line"><span style="$1"><span style="user-select: none;">$2</span>'
+				);
+			}
+			// Handle code wrapping
+			// if wrap=null, do nothing.
+			if (wrap === false) {
+				html = html.replace(/style="(.*?)"/, 'style="$1; overflow-x: auto;"');
+			} else if (wrap === true) {
+				html = html.replace(
+					/style="(.*?)"/,
+					'style="$1; overflow-x: auto; white-space: pre-wrap; word-wrap: break-word;"'
+				);
+			}
+
+			node.type = 'html';
+			node.value = html;
+			node.children = [];
+		});
+	};
+};
+
+export default remarkShiki;