fix(markdown-remark): reuse shiki highlighter instance

.changeset/slow-coins-write.md

@@ -0,0 +1,5 @@
+---
+'@astrojs/markdown-remark': patch
+---
+
+Reuses cached Shiki highlighter instances across languages.

packages/markdown/remark/src/shiki.ts

@@ -1,15 +1,16 @@
 import type { Properties, Root } from 'hast';
 import {
+	type BuiltinLanguage,
 	type BundledLanguage,
-	type BundledTheme,
 	createCssVariablesTheme,
 	createHighlighter,
 	type HighlighterCoreOptions,
-	type HighlighterGeneric,
 	isSpecialLang,
+	type LanguageInput,
 	type LanguageRegistration,
 	type RegexEngine,
 	type ShikiTransformer,
+	type SpecialLanguage,
 	type ThemeRegistration,
 	type ThemeRegistrationRaw,
 } from 'shiki';
@@ -29,6 +30,13 @@ export interface ShikiHighlighter {
 	): Promise<string>;
 }
 
+type ShikiLanguage = LanguageInput | BuiltinLanguage | SpecialLanguage;
+
+interface ShikiHighlighterInternal extends ShikiHighlighter {
+	loadLanguage(...langs: ShikiLanguage[]): Promise<void>;
+	getLoadedLanguages(): string[];
+}
+
 export interface CreateShikiHighlighterOptions {
 	langs?: LanguageRegistration[];
 	theme?: ThemePresets | ThemeRegistration | ThemeRegistrationRaw;
@@ -73,41 +81,100 @@ const cssVariablesTheme = () =>
 		variablePrefix: '--astro-code-',
 	}));
 
-// Caches Promise<ShikiHighlighter> for reuse when the same theme and langs are provided
-const cachedHighlighters = new Map();
+// Caches Promise<ShikiHighlighterInternal> for reuse when the same `themes` and `langAlias`.
+const cachedHighlighters = new Map<string, Promise<ShikiHighlighterInternal>>();
+
+/**
+ * Only used for testing.
+ *
+ * @internal
+ */
+export function clearShikiHighlighterCache(): void {
+	cachedHighlighters.clear();
+}
+
+export function createShikiHighlighter(
+	options?: CreateShikiHighlighterOptions,
+): Promise<ShikiHighlighter> {
+	// Although this function returns a promise, its body runs synchronously so
+	// that the cache lookup happens immediately. Without this, calling
+	// `Promise.all([createShikiHighlighter(), createShikiHighlighter()])` would
+	// bypass the cache and create duplicate highlighters.
+
+	const key: string = getCacheKey(options);
+	let highlighterPromise = cachedHighlighters.get(key);
+	if (!highlighterPromise) {
+		highlighterPromise = createShikiHighlighterInternal(options);
+		cachedHighlighters.set(key, highlighterPromise);
+	}
+	return ensureLanguagesLoaded(highlighterPromise, options?.langs);
+}
+
+/**
+ * Gets the cache key for the highlighter.
+ *
+ * Notice that we don't use `langs` in the cache key because we can dynamically
+ * load languages. This allows us to reuse the same highlighter instance for
+ * different languages.
+ */
+function getCacheKey(options?: CreateShikiHighlighterOptions): string {
+	const keyCache: unknown[] = [];
+	const { theme, themes, langAlias } = options ?? {};
+	if (theme) {
+		keyCache.push(theme);
+	}
+	if (themes) {
+		keyCache.push(Object.entries(themes).sort());
+	}
+	if (langAlias) {
+		keyCache.push(Object.entries(langAlias).sort());
+	}
+	return keyCache.length > 0 ? JSON.stringify(keyCache) : '';
+}
+
+/**
+ * Ensures that the languages are loaded into the highlighter. This is
+ * especially important when the languages are objects representing custom
+ * user-defined languages.
+ */
+async function ensureLanguagesLoaded(
+	promise: Promise<ShikiHighlighterInternal>,
+	langs?: ShikiLanguage[],
+): Promise<ShikiHighlighterInternal> {
+	const highlighter = await promise;
+	if (!langs) {
+		return highlighter;
+	}
+	const loadedLanguages = highlighter.getLoadedLanguages();
+	for (const lang of langs) {
+		if (typeof lang === 'string' && (isSpecialLang(lang) || loadedLanguages.includes(lang))) {
+			continue;
+		}
+		await highlighter.loadLanguage(lang);
+	}
+	return highlighter;
+}
 
 let shikiEngine: RegexEngine | undefined = undefined;
 
-export async function createShikiHighlighter({
+async function createShikiHighlighterInternal({
 	langs = [],
 	theme = 'github-dark',
 	themes = {},
 	langAlias = {},
-}: CreateShikiHighlighterOptions = {}): Promise<ShikiHighlighter> {
+}: CreateShikiHighlighterOptions = {}): Promise<ShikiHighlighterInternal> {
 	theme = theme === 'css-variables' ? cssVariablesTheme() : theme;
 
 	if (shikiEngine === undefined) {
 		shikiEngine = await loadShikiEngine();
 	}
 
-	const highlighterOptions = {
+	const highlighter = await createHighlighter({
 		langs: ['plaintext', ...langs],
 		langAlias,
 		themes: Object.values(themes).length ? Object.values(themes) : [theme],
 		engine: shikiEngine,
-	};
-
-	const key = JSON.stringify(highlighterOptions, Object.keys(highlighterOptions).sort());
-
-	let highlighter: HighlighterGeneric<BundledLanguage, BundledTheme>;
-
-	// Highlighter has already been requested, reuse the same instance
-	if (cachedHighlighters.has(key)) {
-		highlighter = cachedHighlighters.get(key);
-	} else {
-		highlighter = await createHighlighter(highlighterOptions);
-		cachedHighlighters.set(key, highlighter);
-	}
+	});
 
 	async function highlight(
 		code: string,
@@ -220,6 +287,12 @@ export async function createShikiHighlighter({
 		codeToHtml(code, lang, options = {}) {
 			return highlight(code, lang, options, 'html') as Promise<string>;
 		},
+		loadLanguage(...newLangs) {
+			return highlighter.loadLanguage(...newLangs);
+		},
+		getLoadedLanguages() {
+			return highlighter.getLoadedLanguages();
+		},
 	};
 }
 

packages/markdown/remark/test/shiki.test.js

@@ -1,6 +1,7 @@
 import assert from 'node:assert/strict';
 import { describe, it } from 'node:test';
 import { createMarkdownProcessor, createShikiHighlighter } from '../dist/index.js';
+import { clearShikiHighlighterCache } from '../dist/shiki.js';
 
 describe('shiki syntax highlighting', () => {
 	it('does not add is:raw to the output', async () => {
@@ -48,6 +49,51 @@ describe('shiki syntax highlighting', () => {
 		assert.match(hast.children[0].properties.style, /background-color:#24292e;color:#e1e4e8;/);
 	});
 
+	it('createShikiHighlighter can reuse the same instance for different languages', async () => {
+		const langs = [
+			'abap',
+			'ada',
+			'adoc',
+			'angular-html',
+			'angular-ts',
+			'apache',
+			'apex',
+			'apl',
+			'applescript',
+			'ara',
+			'asciidoc',
+			'asm',
+			'astro',
+			'awk',
+			'ballerina',
+			'bash',
+			'bat',
+			'batch',
+			'be',
+			'beancount',
+			'berry',
+			'bibtex',
+			'bicep',
+			'blade',
+			'bsl',
+		];
+
+		const highlighters = new Set();
+		for (const lang of langs) {
+			highlighters.add(await createShikiHighlighter({ langs: [lang] }));
+		}
+
+		// Ensure that we only have one highlighter instance.
+		assert.strictEqual(highlighters.size, 1);
+
+		// Ensure that this highlighter instance can highlight different languages.
+		const highlighter = Array.from(highlighters)[0];
+		const html1 = await highlighter.codeToHtml('const foo = "bar";', 'js');
+		const html2 = await highlighter.codeToHtml('const foo = "bar";', 'ts');
+		assert.match(html1, /color:#F97583/);
+		assert.match(html2, /color:#F97583/);
+	});
+
 	it('diff +/- text has user-select: none', async () => {
 		const highlighter = await createShikiHighlighter();
 
@@ -137,4 +183,43 @@ describe('shiki syntax highlighting', () => {
 
 		assert.match(code, /data-language="cjs"/);
 	});
+
+	it("the cached highlighter won't load the same language twice", async () => {
+		clearShikiHighlighterCache();
+
+		const theme = 'github-light';
+		const highlighter = await createShikiHighlighter({ theme });
+
+		// loadLanguage is an internal method
+		const loadLanguageArgs = [];
+		const originalLoadLanguage = highlighter['loadLanguage'];
+		highlighter['loadLanguage'] = async (...args) => {
+			loadLanguageArgs.push(...args);
+			return await originalLoadLanguage(...args);
+		};
+
+		// No languages loaded yet
+		assert.equal(loadLanguageArgs.length, 0);
+
+		// Load a new language
+		const h1 = await createShikiHighlighter({ theme, langs: ['js'] });
+		assert.equal(loadLanguageArgs.length, 1);
+
+		// Load the same language again
+		const h2 = await createShikiHighlighter({ theme, langs: ['js'] });
+		assert.equal(loadLanguageArgs.length, 1);
+
+		// Load another language
+		const h3 = await createShikiHighlighter({ theme, langs: ['ts'] });
+		assert.equal(loadLanguageArgs.length, 2);
+
+		// Load the same language again
+		const h4 = await createShikiHighlighter({ theme, langs: ['ts'] });
+		assert.equal(loadLanguageArgs.length, 2);
+
+		// All highlighters should be the same instance
+		assert.equal(new Set([highlighter, h1, h2, h3, h4]).size, 1);
+
+		clearShikiHighlighterCache();
+	});
 });