Core: improve Shiki performance

Author: fuma-nama

packages/core/src/highlight/config.ts

@@ -15,15 +15,11 @@ export function defineShikiConfig(config: ShikiConfig): ResolvedShikiConfig {
     createHighlighter() {
       if (created) return created;
 
-      const res = config.createHighlighter();
-      if ('then' in res) {
-        return res.then((v) => {
-          created = v;
-          return v;
-        });
-      } else {
-        return (created = res);
+      created = config.createHighlighter();
+      if ('then' in created) {
+        created = created.then((v) => (created = v));
       }
+      return created;
     },
   };
 }

packages/core/src/highlight/core/client.tsx

@@ -7,8 +7,12 @@ import type { MakeOptional } from '@/types';
 
 const ShikiConfigContext = createContext<ResolvedShikiConfig | null>(null);
 
-export function useShikiConfig(defaultConfig?: ResolvedShikiConfig) {
-  if (defaultConfig) return defaultConfig;
+export function useShikiConfigOptional() {
+  return use(ShikiConfigContext);
+}
+
+export function useShikiConfig(forced?: ResolvedShikiConfig) {
+  if (forced) return forced;
   const ctx = use(ShikiConfigContext);
   if (!ctx) throw new Error(`missing <ShikiConfigProvider />`);
   return ctx;
@@ -26,16 +30,14 @@ export function ShikiConfigProvider({
 
 const promises: Record<string, Promise<ReactNode>> = {};
 
+export type UseShikiOptions = MakeOptional<CoreHighlightOptions, 'config'>;
+
 /**
  * get highlighted results, should be used with React Suspense API.
  *
  * note: results are cached with (lang, code) as keys, if this is not the desired behaviour, pass a `deps` instead.
  */
-export function useShiki(
-  code: string,
-  options: MakeOptional<CoreHighlightOptions, 'config'>,
-  deps?: DependencyList,
-): ReactNode {
+export function useShiki(code: string, options: UseShikiOptions, deps?: DependencyList): ReactNode {
   const config = useShikiConfig(options.config);
   const key = useMemo(() => {
     return deps ? JSON.stringify(deps) : `${options.lang}:${code}`;

packages/core/src/highlight/core/index.ts

@@ -1,10 +1,12 @@
 import type {
-  BundledHighlighterOptions,
   BundledLanguage,
   CodeOptionsMeta,
   CodeOptionsThemes,
   CodeToHastOptionsCommon,
   BundledTheme,
+  HighlighterCore,
+  ThemeRegistrationAny,
+  LanguageRegistration,
 } from 'shiki';
 import {
   type Components,
@@ -30,7 +32,7 @@ export async function highlightHast(
 ): Promise<Root> {
   const { lang: initialLang, fallbackLanguage, config, ...rest } = options;
   let lang = initialLang;
-  let themesToLoad: unknown[];
+  let themesToLoad: (ThemeRegistrationAny | string)[];
 
   if (!('theme' in rest) && !('themes' in rest)) {
     Object.assign(rest, config.defaultThemes);
@@ -45,14 +47,13 @@ export async function highlightHast(
   }
 
   const highlighter = await config.createHighlighter();
-  await highlighter.loadTheme(...(themesToLoad as never[]));
-
-  try {
-    await highlighter.loadLanguage(lang as never);
-  } catch {
-    lang = fallbackLanguage ?? 'text';
-    await highlighter.loadLanguage(lang as never);
-  }
+  await Promise.all([
+    loadMissingTheme(highlighter, ...themesToLoad),
+    loadMissingLanguage(highlighter, lang).catch(() => {
+      lang = fallbackLanguage ?? 'text';
+      if (fallbackLanguage) return loadMissingLanguage(highlighter, fallbackLanguage);
+    }),
+  ]);
 
   return highlighter.codeToHast(code, {
     lang,
@@ -61,6 +62,44 @@ export async function highlightHast(
   });
 }
 
+async function loadMissingTheme(
+  highlighter: HighlighterCore,
+  ...themes: (ThemeRegistrationAny | string)[]
+) {
+  const { isSpecialTheme } = await import('shiki/core');
+
+  const missingThemes = themes.filter((theme) => {
+    if (isSpecialTheme(theme)) return false;
+    try {
+      highlighter.getTheme(theme);
+      return false;
+    } catch {
+      return true;
+    }
+  });
+
+  if (missingThemes.length > 0) await highlighter.loadTheme(...(missingThemes as never[]));
+}
+
+async function loadMissingLanguage(
+  highlighter: HighlighterCore,
+  ...langs: (LanguageRegistration | string)[]
+) {
+  const { isSpecialLang } = await import('shiki/core');
+
+  const missingLangs = langs.filter((lang) => {
+    if (isSpecialLang(lang)) return false;
+    try {
+      highlighter.getLanguage(lang);
+      return false;
+    } catch {
+      return true;
+    }
+  });
+
+  if (missingLangs.length > 0) await highlighter.loadLanguage(...(missingLangs as never[]));
+}
+
 export function hastToJsx(hast: Root, options?: Partial<ToJsxOptions>) {
   return toJsxRuntime(hast, {
     jsx,
@@ -79,13 +118,16 @@ export function hastToJsx(hast: Root, options?: Partial<ToJsxOptions>) {
  */
 export async function getHighlighter(
   config: ResolvedShikiConfig,
-  options: Omit<BundledHighlighterOptions<BundledLanguage, BundledTheme>, 'engine' | 'langAlias'>,
+  options?: {
+    langs?: (BundledLanguage | LanguageRegistration)[];
+    themes?: (BundledTheme | ThemeRegistrationAny)[];
+  },
 ) {
   const highlighter = await config.createHighlighter();
 
   await Promise.all([
-    highlighter.loadLanguage(...(options.langs as never[])),
-    highlighter.loadTheme(...(options.themes as never[])),
+    options?.langs && loadMissingLanguage(highlighter, ...options.langs),
+    options?.themes && loadMissingTheme(highlighter, ...options.themes),
   ]);
 
   return highlighter;

packages/core/src/highlight/full/config.ts

@@ -1,14 +1,16 @@
-import { createHighlighter, createOnigurumaEngine } from 'shiki';
 import { defineShikiConfig } from '../config';
 
-export const withJSEngine = defineShikiConfig({
-  defaultThemes: {
-    themes: {
-      light: 'github-light',
-      dark: 'github-dark',
-    },
+const defaultThemes = {
+  themes: {
+    light: 'github-light',
+    dark: 'github-dark',
   },
+};
+
+export const withJSEngine = defineShikiConfig({
+  defaultThemes,
   async createHighlighter() {
+    const { createHighlighter } = await import('shiki');
     const { createJavaScriptRegexEngine } = await import('shiki/engine/javascript');
 
     return createHighlighter({
@@ -20,13 +22,9 @@ export const withJSEngine = defineShikiConfig({
 });
 
 export const withWASMEngine = defineShikiConfig({
-  defaultThemes: {
-    themes: {
-      light: 'github-light',
-      dark: 'github-dark',
-    },
-  },
-  createHighlighter() {
+  defaultThemes,
+  async createHighlighter() {
+    const { createHighlighter, createOnigurumaEngine } = await import('shiki');
     return createHighlighter({
       langs: [],
       themes: [],