Use shikiji (#8502)

Co-authored-by: Sarah Rainsberger <sarah@rainsberger.ca>

2 files changed, 52 insertions, 38 deletions

diff --git a/packages/markdown/remark/src/remark-shiki.ts b/packages/markdown/remark/src/remark-shiki.ts
index 58ed16369..bf3dd0b78 100644
--- a/packages/markdown/remark/src/remark-shiki.ts
+++ b/packages/markdown/remark/src/remark-shiki.ts
@@ -1,56 +1,52 @@
-import type * as shiki from 'shiki';
-import { getHighlighter } from 'shiki';
+import { bundledLanguages, getHighlighter, type Highlighter } from 'shikiji';
 import { visit } from 'unist-util-visit';
 import type { RemarkPlugin, ShikiConfig } from './types.js';
 
+const ASTRO_COLOR_REPLACEMENTS: Record<string, string> = {
+	'#000001': 'var(--astro-code-color-text)',
+	'#000002': 'var(--astro-code-color-background)',
+	'#000004': 'var(--astro-code-token-constant)',
+	'#000005': 'var(--astro-code-token-string)',
+	'#000006': 'var(--astro-code-token-comment)',
+	'#000007': 'var(--astro-code-token-keyword)',
+	'#000008': 'var(--astro-code-token-parameter)',
+	'#000009': 'var(--astro-code-token-function)',
+	'#000010': 'var(--astro-code-token-string-expression)',
+	'#000011': 'var(--astro-code-token-punctuation)',
+	'#000012': 'var(--astro-code-token-link)',
+};
+const COLOR_REPLACEMENT_REGEX = new RegExp(
+	`(${Object.keys(ASTRO_COLOR_REPLACEMENTS).join('|')})`,
+	'g'
+);
+
 /**
  * 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 highlighterCacheAsync = new Map<string, Promise<Highlighter>>();
 
 export function remarkShiki({
 	langs = [],
 	theme = 'github-dark',
 	wrap = false,
 }: ShikiConfig = {}): ReturnType<RemarkPlugin> {
-	const cacheID: string = typeof theme === 'string' ? theme : theme.name;
-	let highlighterAsync = highlighterCacheAsync.get(cacheID);
+	const cacheId =
+		(typeof theme === 'string' ? theme : theme.name ?? '') +
+		langs.map((l) => l.name ?? (l as any).id).join(',');
+
+	let highlighterAsync = highlighterCacheAsync.get(cacheId);
 	if (!highlighterAsync) {
-		highlighterAsync = getHighlighter({ theme }).then((hl) => {
-			hl.setColorReplacements({
-				'#000001': 'var(--astro-code-color-text)',
-				'#000002': 'var(--astro-code-color-background)',
-				'#000004': 'var(--astro-code-token-constant)',
-				'#000005': 'var(--astro-code-token-string)',
-				'#000006': 'var(--astro-code-token-comment)',
-				'#000007': 'var(--astro-code-token-keyword)',
-				'#000008': 'var(--astro-code-token-parameter)',
-				'#000009': 'var(--astro-code-token-function)',
-				'#000010': 'var(--astro-code-token-string-expression)',
-				'#000011': 'var(--astro-code-token-punctuation)',
-				'#000012': 'var(--astro-code-token-link)',
-			});
-			return hl;
+		highlighterAsync = getHighlighter({
+			langs: langs.length ? langs : Object.keys(bundledLanguages),
+			themes: [theme],
 		});
-		highlighterCacheAsync.set(cacheID, highlighterAsync);
+		highlighterCacheAsync.set(cacheId, highlighterAsync);
 	}
 
-	let highlighter: shiki.Highlighter;
-
 	return async (tree: any) => {
-		// Lazily assign the highlighter as async can only happen within this function,
-		// and not on `remarkShiki` directly.
-		if (!highlighter) {
-			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);
-			}
-		}
+		const highlighter = await highlighterAsync!;
 
 		visit(tree, 'code', (node) => {
 			let lang: string;
@@ -68,7 +64,7 @@ export function remarkShiki({
 				lang = 'plaintext';
 			}
 
-			let html = highlighter.codeToHtml(node.value, { lang });
+			let html = highlighter.codeToHtml(node.value, { lang, theme });
 
 			// Q: Couldn't these regexes match on a user's inputted code blocks?
 			// A: Nope! All rendered HTML is properly escaped.
@@ -96,9 +92,22 @@ export function remarkShiki({
 				);
 			}
 
+			// theme.id for shiki -> shikiji compat
+			const themeName = typeof theme === 'string' ? theme : theme.name;
+			if (themeName === 'css-variables') {
+				html = html.replace(/style="(.*?)"/g, (m) => replaceCssVariables(m));
+			}
+
 			node.type = 'html';
 			node.value = html;
 			node.children = [];
 		});
 	};
 }
+
+/**
+ * shiki -> shikiji compat as we need to manually replace it
+ */
+function replaceCssVariables(str: string) {
+	return str.replace(COLOR_REPLACEMENT_REGEX, (match) => ASTRO_COLOR_REPLACEMENTS[match] || match);
+}
diff --git a/packages/markdown/remark/src/types.ts b/packages/markdown/remark/src/types.ts
index bcab97041..4abcf578d 100644
--- a/packages/markdown/remark/src/types.ts
+++ b/packages/markdown/remark/src/types.ts
@@ -5,7 +5,12 @@ import type {
 	all as Handlers,
 	Options as RemarkRehypeOptions,
 } from 'remark-rehype';
-import type { ILanguageRegistration, IThemeRegistration, Theme } from 'shiki';
+import type {
+	BuiltinTheme,
+	LanguageRegistration,
+	ThemeRegistration,
+	ThemeRegistrationRaw,
+} from 'shikiji';
 import type * as unified from 'unified';
 import type { VFile } from 'vfile';
 
@@ -35,8 +40,8 @@ export type RemarkRehype = Omit<RemarkRehypeOptions, 'handlers' | 'unknownHandle
 };
 
 export interface ShikiConfig {
-	langs?: ILanguageRegistration[];
-	theme?: Theme | IThemeRegistration;
+	langs?: LanguageRegistration[];
+	theme?: BuiltinTheme | ThemeRegistration | ThemeRegistrationRaw;
 	wrap?: boolean | null;
 }