perf(plugin-shiki): lazy load languages (#347)

Author: Mister-Hope

docs/plugins/markdown/shiki.md

@@ -41,13 +41,11 @@ export default {
 
 - Details:
 
-  Languages of code blocks to be parsed by Shiki.
+  Additional languages to be parsed by Shiki.
 
-  This option will be forwarded to `createHighlighter()` method of Shiki.
+  ::: tip
 
-  ::: warning
-
-  We recommend you to provide the languages list you are using explicitly, otherwise Shiki will load all languages and can affect performance.
+  The plugin now automatically loads the languages used in your markdown files, so you don't need to specify them manually.
 
   :::
 

docs/zh/plugins/markdown/shiki.md

@@ -41,13 +41,11 @@ export default {
 
 - 详情：
 
-  Shiki 要解析的代码块的语言。
+  被 Shiki 解析的额外语言。
 
-  该配置项会被传递到 Shiki 的 `createHighlighter()` 方法中。
+  ::: tip
 
-  ::: warning
-
-  我们建议明确传入所有你使用的语言列表，否则 Shiki 会加载所有语言，并可能影响性能。
+  插件现在会自动加载你的 markdown 文件中使用的语言，所以你不需要手动指定它们。
 
   :::
 

plugins/markdown/plugin-shiki/package.json

@@ -23,6 +23,7 @@
   "type": "module",
   "exports": {
     ".": "./lib/node/index.js",
+    "./resolveLang": "./lib/node/resolveLang.js",
     "./styles/*": "./lib/client/styles/*",
     "./package.json": "./package.json"
   },
@@ -42,7 +43,8 @@
     "@vuepress/helper": "workspace:*",
     "@vuepress/highlighter-helper": "workspace:*",
     "nanoid": "^5.0.9",
-    "shiki": "^2.1.0"
+    "shiki": "^2.1.0",
+    "synckit": "^0.9.2"
   },
   "peerDependencies": {
     "vuepress": "2.0.0-rc.19"

plugins/markdown/plugin-shiki/rollup.config.ts

@@ -1,5 +1,8 @@
 import { rollupBundle } from '../../../scripts/rollup.js'
 
-export default rollupBundle('node/index', {
-  external: ['@shikijs/transformers', 'nanoid', 'shiki'],
-})
+export default rollupBundle(
+  { base: 'node', files: ['index', 'resolveLang'] },
+  {
+    external: ['@shikijs/transformers', 'nanoid', 'shiki', 'synckit'],
+  },
+)

plugins/markdown/plugin-shiki/src/node/index.ts

@@ -1,4 +1,3 @@
 export type * from './options.js'
-export * from './shiki.js'
 export * from './shikiPlugin.js'
 export type * from './types.js'

plugins/markdown/plugin-shiki/src/node/markdown/highlighter/createShikiHighlighter.ts

@@ -1,27 +1,66 @@
+import { createRequire } from 'node:module'
 import type { BundledLanguage, BundledTheme, HighlighterGeneric } from 'shiki'
-import { createHighlighter } from 'shiki'
-import { bundledLanguageNames } from '../../shiki.js'
+import { createHighlighter, isSpecialLang } from 'shiki'
+import { createSyncFn } from 'synckit'
+import type { ShikiResolveLang } from '../../resolveLang.js'
 import type { ShikiHighlightOptions } from '../../types.js'
+import { resolveLanguage } from '../../utils.js'
+
+const require = createRequire(import.meta.url)
+
+const resolveLangSync = createSyncFn<ShikiResolveLang>(
+  require.resolve('@vuepress/plugin-shiki/resolveLang'),
+)
+
+export type ShikiLoadLang = (lang: string) => boolean
 
 export const createShikiHighlighter = async ({
-  langs = bundledLanguageNames,
+  langs = [],
   langAlias = {},
   defaultLang,
   shikiSetup,
   ...options
-}: ShikiHighlightOptions = {}): Promise<
-  HighlighterGeneric<BundledLanguage, BundledTheme>
-> => {
-  const shikiHighlighter = await createHighlighter({
-    langs,
+}: ShikiHighlightOptions = {}): Promise<{
+  highlighter: HighlighterGeneric<BundledLanguage, BundledTheme>
+  loadLang: ShikiLoadLang
+}> => {
+  const highlighter = await createHighlighter({
+    langs: [...langs, ...Object.values(langAlias)],
     langAlias,
     themes:
       'themes' in options
         ? Object.values(options.themes)
         : [options.theme ?? 'nord'],
   })
 
-  await shikiSetup?.(shikiHighlighter)
+  const loadLang = (lang: string): boolean => {
+    if (isSpecialLang(lang)) return true
+
+    const loadedLangs = highlighter.getLoadedLanguages()
+
+    if (!loadedLangs.includes(lang)) {
+      const resolvedLang = resolveLangSync(lang)
+
+      if (!resolvedLang.length) return false
+
+      highlighter.loadLanguageSync(resolvedLang)
+    }
+
+    return true
+  }
+
+  // patch for twoslash - https://github.com/vuejs/vitepress/issues/4334
+  const rawGetLanguage = highlighter.getLanguage
+
+  highlighter.getLanguage = (name) => {
+    const lang = typeof name === 'string' ? name : name.name
+
+    loadLang(resolveLanguage(lang))
+
+    return rawGetLanguage.call(highlighter, name)
+  }
+
+  await shikiSetup?.(highlighter)
 
-  return shikiHighlighter
+  return { highlighter, loadLang }
 }

plugins/markdown/plugin-shiki/src/node/markdown/highlighter/getHighLightFunction.ts

@@ -7,6 +7,7 @@ import {
 import type { ShikiHighlightOptions } from '../../types.js'
 import { attrsToLines } from '../../utils.js'
 import type { MarkdownFilePathGetter } from './createMarkdownFilePathGetter.js'
+import type { ShikiLoadLang } from './createShikiHighlighter.js'
 import { getLanguage } from './getLanguage.js'
 import { handleMustache } from './handleMustache.js'
 
@@ -19,20 +20,15 @@ type MarkdownItHighlight = (
 export const getHighLightFunction = (
   highlighter: HighlighterGeneric<BundledLanguage, BundledTheme>,
   options: ShikiHighlightOptions,
+  loadLang: ShikiLoadLang,
   markdownFilePathGetter: MarkdownFilePathGetter,
 ): MarkdownItHighlight => {
   const transformers = getTransformers(options)
-  const loadedLanguages = highlighter.getLoadedLanguages()
 
   return (content, language, attrs) =>
     handleMustache(content, (str) =>
       highlighter.codeToHtml(str, {
-        lang: getLanguage(
-          language,
-          loadedLanguages,
-          options,
-          markdownFilePathGetter,
-        ),
+        lang: getLanguage(language, options, loadLang, markdownFilePathGetter),
         meta: {
           /**
            * Custom `transformers` passed by users may require `attrs`.

plugins/markdown/plugin-shiki/src/node/markdown/highlighter/getLanguage.ts

@@ -1,20 +1,20 @@
-import { isSpecialLang } from 'shiki'
 import { colors } from 'vuepress/utils'
 import type { ShikiHighlightOptions } from '../../types.js'
 import { logger, resolveLanguage } from '../../utils.js'
 import type { MarkdownFilePathGetter } from './createMarkdownFilePathGetter.js'
+import type { ShikiLoadLang } from './createShikiHighlighter.js'
 
 const WARNED_LANGS = new Set<string>()
 
 export const getLanguage = (
   lang: string,
-  loadedLanguages: string[],
   { defaultLang, logLevel }: ShikiHighlightOptions,
+  loadLang: ShikiLoadLang,
   markdownFilePathGetter: MarkdownFilePathGetter,
 ): string => {
   let result = resolveLanguage(lang)
 
-  if (result && !loadedLanguages.includes(result) && !isSpecialLang(result)) {
+  if (result && !loadLang(result)) {
     // warn for unknown languages only once
     if (logLevel !== 'silent' && !WARNED_LANGS.has(result)) {
       logger.warn(

plugins/markdown/plugin-shiki/src/node/resolveLang.ts

@@ -0,0 +1,23 @@
+import type {
+  DynamicImportLanguageRegistration,
+  LanguageRegistration,
+} from 'shiki'
+import { bundledLanguages } from 'shiki'
+import { runAsWorker } from 'synckit'
+
+async function resolveLang(lang: string): Promise<LanguageRegistration[]> {
+  return (
+    (
+      bundledLanguages as Record<
+        string,
+        DynamicImportLanguageRegistration | undefined
+      >
+    )
+      [lang]?.()
+      .then((m) => m.default) ?? []
+  )
+}
+
+runAsWorker(resolveLang)
+
+export type ShikiResolveLang = typeof resolveLang

plugins/markdown/plugin-shiki/src/node/shiki.ts

@@ -41,11 +41,12 @@ export const shikiPlugin = (_options: ShikiPluginOptions = {}): Plugin => {
         const { preWrapper, lineNumbers, collapsedLines } = options
 
         const markdownFilePathGetter = createMarkdownFilePathGetter(md)
-        const shikiHighlighter = await createShikiHighlighter(options)
+        const { highlighter, loadLang } = await createShikiHighlighter(options)
 
         md.options.highlight = getHighLightFunction(
-          shikiHighlighter,
+          highlighter,
           options,
+          loadLang,
           markdownFilePathGetter,
         )
 

plugins/markdown/plugin-shiki/src/node/shikiPlugin.ts

@@ -18,21 +18,22 @@ import {
 } from '../src/node/markdown/index.js'
 import type { ShikiPluginOptions } from '../src/node/options.js'
 
-const shikiHighlighter = await createShikiHighlighter()
+const { highlighter, loadLang } = await createShikiHighlighter()
 
 const createMarkdown = ({
   preWrapper = true,
   lineNumbers = true,
   collapsedLines = false,
   ...options
 }: ShikiPluginOptions = {}): MarkdownIt => {
-  const md = MarkdownIt()
+  const md = new MarkdownIt()
 
   const markdownFilePathGetter = createMarkdownFilePathGetter(md)
 
   md.options.highlight = getHighLightFunction(
-    shikiHighlighter,
+    highlighter,
     options,
+    loadLang,
     markdownFilePathGetter,
   )