更快地初始化 Shiki 代码高亮

Shiki 作为热门的代码高亮库，通过 rehype 插件使用时由于默认情况下需要加载所有主题和语言定义，导致初始化的性能不太能被接受。本文简单分享一些在本网站中使用的按需加载 Shiki 资源的思路。

问题

我是以 rehype 插件的形式使用 Shiki 的，这种方法默认会在内部创建一个包含了所有主题和语言定义的 highlighter 实例。而由于主题和语言定义都是通过动态 import() 加载的，导致第一次创建时的性能非常差：在本地机器（M2 Pro）上，需要大约 2500ms，而在 Vercel 的构建机器上需要 8000ms 以上。

(feed.xml 是第一个需要渲染代码高亮的页面，通过去除代码高亮功能来进行对比，可以发现它主要的开销就是 Shiki 的初始化)

而如果是实时服务端渲染的场景，如果是 Serverless 冷启动的情况，那每个请求都会需要初始化一次 Shiki，加载速度完全没法接受。本站之前使用 Next.js + RSC 渲染博客文章时非常依赖 Route Cache，否则每次访问都需要等待5秒以上。

按需加载

对于这种初始化慢的问题，Shiki 文档中给出的解决方法是使用 Fine-grained Bundle，只加载需要的主题和语言定义。

最容易优化的是主题的加载：一个站点使用的主题肯定是提前选择好的，因此可以直接指定好需要的主题。

const highlighter = await createHighlighterCore({
  themes: [import('@shikijs/themes/vitesse-light')], // 只加载一个主题
  langs: [], // 稍后再说语言的加载
  engine: createOnigurumaEngine(() => import('shiki/wasm'))
})

然而语言定义的加载就相对复杂：Shiki 支持的语言中，我有一大部分这辈子可能都不会用到；但另一方面，我既不知道以前曾经在博客文章中在代码块中使用过哪些语言，更不知道未来可能会使用哪些语言，因此无法提前在代码中指定需要加载哪些语言定义。

因此，语言的加载只能是实时、按需的，我的思路是，在 shiki 的 rehype 插件之前再添加一个 rehype 插件，它会扫描出代码块中使用过的语言，将它们添加进 highlighter 的语言定义中，之后的 shiki rehype 插件就可以直接使用这个 highlighter 了。

完整代码：

import rehypeShikiFromHighlighter from "@shikijs/rehype/core";
import type { Element } from "hast";
import { createHighlighterCore } from "shiki/core";
import { createOnigurumaEngine } from "shiki/engine/oniguruma";
import { bundledLanguages, type BundledLanguage } from "shiki/langs";
import { bundledThemes } from "shiki/themes";
import type { Pluggable, Plugin, Preset } from "unified";
import { visit, CONTINUE, SKIP } from "unist-util-visit";

const highlighter = await createHighlighterCore({
  themes: [bundledThemes["catppuccin-macchiato"]],
  // initially empty, the `rehypeCodeLangDetecter` will load the languages
  langs: [],
  engine: createOnigurumaEngine(() => import("shiki/wasm")),
});

const loadedLangs = new Set<string>();

const shikiPluggable: Pluggable = [
  rehypeShikiFromHighlighter,
  highlighter,
  { theme: "catppuccin-macchiato" },
];

export const rehypeShikiPreset: Preset = {
  plugins: [rehypeCodeLangDetecter, shikiPluggable],
};

function rehypeCodeLangDetecter(): ReturnType<Plugin> {
  return async (tree, file) => {
    const langs: string[] = [];
    visit(tree, "element", (node: Element) => {
      if (node.tagName !== "code") return CONTINUE;
      if (!node.properties.className) return CONTINUE;
      const className = node.properties.className;
      if (!Array.isArray(className)) return CONTINUE;
      const lang = className.find(
        (cls): cls is string =>
          typeof cls === "string" && cls.startsWith("language-"),
      );
      if (!lang) return CONTINUE;
      langs.push(lang.slice("language-".length));
      return SKIP;
    });
    file.data.langs = langs;
    for (const lang of langs) {
      if (loadedLangs.has(lang)) continue;
      if (!(lang in bundledLanguages)) continue;
      const langRegistration = bundledLanguages[lang as BundledLanguage];
      await highlighter.loadLanguage(langRegistration);
      loadedLangs.add(lang);
    }
  };
}

1. highlighter 是一个模块级别单例，在当前 JS Context 中，所有的代码高亮都由这个实例完成。
2. loadedLangs Set 用于记录已经加载过的语言，避免重复加载带来的开销。
3. rehypeCodeLangDetecter 插件会扫描出代码块中使用过的语言，并将它们加载到 highlighter 中。
   1. 由于 remark-rehype 插件会将非 inline 的 Code 节点转换成 <pre><code class="language-xxx">...</code></pre> 的形式，而 inline 或者没有指定语言的代码块不会含有 language-xxx 的 class，因此只需要扫描 className 即可。
   2. 同时会将语言信息存储在 file.data.langs 中，方便后续其他插件使用。
4. rehypeShikiFromHighlighter 插件使用这个 highlighter 来进行代码高亮。
5. rehypeShikiPreset 是一个 rehype preset，可以直接在 unified 的处理链中使用。

使用示例：

const vfile = await unified()
      .use(remarkParse)
      .use(remarkRehype)
      .use(shikiPreset) // 直接使用上面定义的 preset
      .use(rehypeStringify)
      .process("your markdown content");
const html = String(vfile);

效果

一个简单的 benchmark 可以发现加载3个语言比加载所有语言快 3000 倍。（作为一个 micro benchmark，这个结果仅供图一乐）

而在实际渲染时，第一个使用 highlighter 的路由渲染时间从 10s+ 降低到了不到 2s，性能差距非常明显；而后续的请求虽然不需要重新初始化 highlighter，但可能是因为使用的语言少，普遍也有约 20% 的性能提升（从这也可以看出来，代码高亮是主要的性能瓶颈）。

参考

• Shiki 文档
• Shiki rehype Fine-grained Bundle
• Shiki Fine-grained Bundle overview
• unified 文档