perf(core): load Shiki grammars on demand

The Shiki plugin registered every bundled language (~250 grammars) in
each thread that highlighted code, which cost ~2s and ~100MB per thread
and made every highlight several times slower, as each code block was
matched against grammars it never uses. Importing it also imported all
of them, through `@node-core/rehype-shiki`'s `LANGS` and its plugin,
on every thread loading `jsx-ast`, the main thread included.

The highlighter now registers a bundled language the first time code in
it is highlighted, along with the bundled languages a configured one
embeds, and lists the bundled ones from their metadata alone, without
importing `LANGS`. The themes are given to Shiki by name, which it keeps
parsed instead of parsing them for every highlight, named after the
configured themes, as Shiki loads a theme given for both schemes once.

A bundled language is registered with the bundled languages it embeds,
lazily too (a fenced code block in Markdown): Shiki only highlights
those once they are registered, so registering them later would make
code highlight differently depending on what was highlighted before it.
For the same reason, a bundled grammar injecting into others (`ts-tags`,
into JavaScript's tagged templates) only injects into the language it
was imported for.

`@node-core/rehype-shiki` is updated to 1.5.1, whose plugin only imports
every grammar when it isn't given a highlighter (nodejs/nodejs.org#9212).

The grammars now come from doc-kit's own `shiki` dependency (4.4.3),
which 1.5.1 shares, rather than the copy 1.5.0 pinned (4.3.1). Its C++
grammar highlights types and template arguments differently, which
shows on the Node.js docs' C++ examples.

Assisted-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: ovflowd

docs/configuration.md

@@ -218,6 +218,11 @@ bundles. To configure it, list it with options in
 dependencies, so your configuration file can resolve it):
 
 - `langs` {Array} More languages: grammars, or modules default-exporting them.
+  The languages Shiki bundles are registered once code in them is highlighted,
+  so one injecting into others (such as `ts-tags`, into JavaScript's tagged
+  templates) only does so in its own code blocks; list it here (e.g.
+  `import.meta.resolve('shiki/langs/ts-tags.mjs')`) to register it up front,
+  as a transformer highlighting code in another language needs too.
 - `langAlias` {Object} Aliases of languages, such as `{ conf: 'ini' }`.
 - `themes` {Object} The `light` and `dark` themes: names of themes Shiki
   bundles, modules default-exporting themes, or themes.

docs/creating-generators.md

@@ -524,7 +524,9 @@ A plugin needing asynchronous setup can export an async `load(options)`
 returning the plugin, instead of a default export. The Shiki plugin does, and
 `getHighlighter(name)` (from `@doc-kit/core/plugins/shiki/highlighter.mjs`)
 returns the highlighter of a generator's Shiki plugin, to highlight code as its
-pipeline does.
+pipeline does. It registers the languages Shiki bundles as code in them is
+highlighted: pass a language through its `resolveLanguage` before using its
+`shiki` instance directly.
 
 ## File Output
 

packages/core/package.json

@@ -53,7 +53,7 @@
   ],
   "dependencies": {
     "@actions/core": "^3.0.0",
-    "@node-core/rehype-shiki": "^1.5.0",
+    "@node-core/rehype-shiki": "^1.5.1",
     "@swc/html-wasm": "^1.16.2",
     "@swc/wasm": "^1.15.46",
     "acorn": "^8.17.0",

packages/core/src/plugins/shiki/__tests__/highlighter.test.mjs

@@ -86,6 +86,57 @@ describe('createHighlighter', () => {
     );
   });
 
+  it('registers a bundled language once code in it is highlighted', async () => {
+    const highlighter = await createHighlighter({
+      langAlias: { py: 'python' },
+    });
+    const loaded = () => highlighter.shiki.getLoadedLanguages();
+
+    assert.ok(!loaded().includes('python'));
+    assert.ok(!loaded().includes('javascript'));
+
+    // By its name, an alias of its own, or one of the options
+    assert.equal(highlighter.resolveLanguage('py'), 'py');
+    assert.equal(highlighter.resolveLanguage('mjs'), 'mjs');
+
+    assert.ok(loaded().includes('python'));
+    assert.ok(loaded().includes('javascript'));
+    assert.ok(loaded().includes('cjs'));
+
+    assert.equal(highlighter.resolveLanguage(undefined), 'text');
+    assert.equal(highlighter.resolveLanguage('plaintext'), 'plaintext');
+  });
+
+  it('lists the bundled languages without registering them', async () => {
+    const highlighter = await createHighlighter({ langs: [grammar] });
+
+    assert.deepStrictEqual(
+      highlighter.langs.find(({ name }) => name === 'rust'),
+      { name: 'rust', displayName: 'Rust', aliases: ['rs'] }
+    );
+    assert.equal(highlighter.langs.at(-1), grammar);
+    assert.ok(!highlighter.shiki.getLoadedLanguages().includes('rust'));
+  });
+
+  it('registers the bundled languages a language embeds', async () => {
+    const highlighter = await createHighlighter({
+      langs: [
+        {
+          ...grammar,
+          name: 'oxcscript',
+          scopeName: 'source.oxcscript',
+          embeddedLangs: ['javascript'],
+          patterns: [{ include: 'source.js' }],
+        },
+      ],
+    });
+
+    assert.match(
+      highlighter.highlightToHtml('const on = 1', 'oxcscript'),
+      /--shiki-dark/
+    );
+  });
+
   it('gives the same highlighter for the same options', async () => {
     const highlighter = await createHighlighter({ langs: [grammar] });
 

packages/core/src/plugins/shiki/highlighter.mjs

@@ -1,14 +1,24 @@
 'use strict';
 
+import { createRequire } from 'node:module';
 import { endianness } from 'node:os';
 
-import { LANGS } from '@node-core/rehype-shiki';
 import createSyntaxHighlighter from '@node-core/rehype-shiki/highlighter';
+import { isSpecialLang } from 'shiki/core';
+import { bundledLanguagesInfo } from 'shiki/langs';
 import { bundledThemes } from 'shiki/themes';
 
 import { importFromURL } from '#utils/loaders.mjs';
 import { getMarkdownPlugins } from '#utils/markdown/plugins.mjs';
 
+const require = createRequire(import.meta.url);
+
+/**
+ * A language a highlighter highlights.
+ *
+ * @typedef {Pick<import('shiki').LanguageRegistration, 'name' | 'displayName' | 'aliases'>} Language
+ */
+
 /**
  * A syntax highlighter, creating its Shiki instance on first use.
  *
@@ -17,7 +27,7 @@ import { getMarkdownPlugins } from '#utils/markdown/plugins.mjs';
  * @property {(languageId?: string) => string} resolveLanguage - Resolves a language, falling back to plain text for unknown ones
  * @property {(code: string, lang: string, meta?: Record<string, unknown>) => string} highlightToHtml - Highlights code, returning the inner HTML of its `<code>` element
  * @property {(code: string, lang: string, meta?: Record<string, unknown>) => ReturnType<import('shiki').HighlighterCore['codeToHast']>} highlightToHast - Highlights code, returning a HAST tree
- * @property {Array<import('shiki').LanguageRegistration>} langs - The languages it highlights
+ * @property {Array<Language>} langs - The languages it highlights
  */
 
 /**
@@ -53,6 +63,87 @@ let engine;
 // The highlighters of the options given, by their JSON
 const highlighters = new Map();
 
+// What code in an unknown language is highlighted as
+const FALLBACK_LANGUAGE = 'text';
+
+// The languages Shiki bundles. Registering a grammar takes time and memory,
+// and slows down every highlight after it, so each one is only imported and
+// registered once code in it is highlighted (see `resolveLanguage`).
+const BUNDLED_LANGUAGES = bundledLanguagesInfo.map(({ id, name, aliases }) => ({
+  name: id,
+  displayName: name,
+  aliases,
+}));
+
+// The bundled language each of their names and aliases stands for, as code
+// names its language by either
+const BUNDLED_NAMES = new Map(
+  BUNDLED_LANGUAGES.flatMap(({ name, aliases = [] }) =>
+    [name, ...aliases].map(alias => [alias, name])
+  )
+);
+
+/**
+ * Imports a bundled language: its grammar, and those of the languages it
+ * embeds, lazily too. Shiki only highlights a language embedded lazily (a
+ * fenced code block in Markdown) once it is registered, so registering it
+ * later would make code highlight differently depending on what was
+ * highlighted before it.
+ *
+ * For the same reason, a grammar injecting into others' (`ts-tags`', into
+ * JavaScript's tagged templates) only injects into the language it was
+ * imported for: Shiki only applies an injection to the grammars it creates
+ * after it. A language given in `langs` keeps its injections.
+ *
+ * @param {string} name - A name or alias of the language
+ * @returns {Array<import('shiki').LanguageRegistration>}
+ */
+const importBundledLanguage = name => {
+  const id = BUNDLED_NAMES.get(name);
+  const langs = new Map();
+  const pending = [id];
+
+  while (pending.length > 0) {
+    const grammars = require(`shiki/langs/${pending.pop()}.mjs`).default;
+
+    for (const lang of grammars) {
+      if (!langs.has(lang.name)) {
+        langs.set(lang.name, lang);
+
+        pending.push(
+          ...[...(lang.embeddedLangs ?? []), ...(lang.embeddedLangsLazy ?? [])]
+            .filter(embedded => BUNDLED_NAMES.has(embedded))
+            .map(embedded => BUNDLED_NAMES.get(embedded))
+        );
+      }
+    }
+  }
+
+  const { scopeName } = langs.get(id);
+
+  return [...langs.values()].map(lang =>
+    lang.injectTo ? { ...lang, injectTo: [scopeName] } : lang
+  );
+};
+
+/**
+ * Adds the bundled languages that languages embed, which Shiki registers
+ * along with them.
+ *
+ * @param {Array<import('shiki').LanguageRegistration>} langs - The languages
+ * @returns {Array<import('shiki').LanguageRegistration>}
+ */
+const withEmbeddedLanguages = langs => {
+  const names = new Set(langs.map(({ name }) => name));
+
+  const embedded = langs
+    .flatMap(({ embeddedLangs = [] }) => embeddedLangs)
+    .filter(name => !names.has(name) && BUNDLED_NAMES.has(name))
+    .flatMap(importBundledLanguage);
+
+  return [...embedded, ...langs];
+};
+
 /**
  * Imports a list of options, each given as it is, or as a module (a path
  * relative to the working directory, or a URL) default-exporting it, or a list
@@ -119,13 +210,12 @@ const importHighlighter = async ({
 
   const coreOptions = {
     engine: regexEngine,
-    langs: [...LANGS, ...importedLangs],
-    // A copy, as Shiki adds the aliases of the languages it bundles to it
+    // The bundled languages are registered as code uses them
+    langs: withEmbeddedLanguages(importedLangs),
+    // A copy, as Shiki adds the aliases of the languages it registers to it
     langAlias: { ...langAlias },
   };
 
-  const highlighterOptions = { transformers: importedTransformers };
-
   // Without themes of its own, the highlighter has a default light and dark one
   if (themes) {
     const [light, dark] = await Promise.all([
@@ -134,58 +224,117 @@ const importHighlighter = async ({
     ]);
 
     coreOptions.themes = [light, dark];
-    highlighterOptions.themes = { light: light.name, dark: dark.name };
-    highlighterOptions.defaultColor = 'light';
   }
 
-  let highlighter;
+  let shiki;
+  let highlightOptions;
+
+  /**
+   * Gives the Shiki instance, creating it on first use.
+   *
+   * @returns {import('shiki').HighlighterCore}
+   */
+  const current = () => {
+    if (!shiki) {
+      ({ shiki } = createSyntaxHighlighter({ coreOptions }));
+
+      // The themes are given by name: Shiki keeps the themes it loaded parsed,
+      // but parses a theme object it's given again for every highlight
+      const [light, dark] =
+        coreOptions.themes?.map(theme => theme.name) ?? shiki.getLoadedThemes();
+
+      highlightOptions = {
+        themes: { light, dark },
+        defaultColor: 'light',
+        transformers: importedTransformers,
+      };
+    }
+
+    return shiki;
+  };
 
   /**
-   * Gives the Shiki highlighter, creating it on first use.
+   * Resolves a language, falling back to plain text for unknown ones. A
+   * bundled language is registered the first time it's resolved.
+   *
+   * @param {string} [languageId]
+   * @returns {string}
    */
-  const current = () =>
-    (highlighter ??= createSyntaxHighlighter({
-      coreOptions,
-      highlighterOptions,
-    }));
+  const resolveLanguage = languageId => {
+    if (!languageId) {
+      return FALLBACK_LANGUAGE;
+    }
+
+    const instance = current();
+    const name = instance.resolveLangAlias(languageId.toLowerCase());
+
+    if (isSpecialLang(name) || instance.getLoadedLanguages().includes(name)) {
+      return languageId;
+    }
+
+    if (!BUNDLED_NAMES.has(name)) {
+      return FALLBACK_LANGUAGE;
+    }
+
+    instance.loadLanguageSync(importBundledLanguage(name));
+
+    return languageId;
+  };
+
+  /**
+   * The options Shiki highlights code with.
+   *
+   * @param {string} lang - The language of the code
+   * @param {Record<string, unknown>} meta - Its metadata
+   */
+  const optionsFor = (lang, meta) => ({
+    lang: resolveLanguage(lang),
+    ...highlightOptions,
+    meta,
+  });
 
   return {
-    langs: coreOptions.langs,
+    langs: [...BUNDLED_LANGUAGES, ...importedLangs],
 
     /**
      * The Shiki instance.
      */
     get shiki() {
-      return current().shiki;
+      return current();
     },
 
-    /**
-     * Resolves a language, falling back to plain text for unknown ones.
-     *
-     * @param {string} [languageId]
-     */
-    resolveLanguage: languageId => current().resolveLanguage(languageId),
+    resolveLanguage,
 
     /**
      * Highlights code, returning the inner HTML of its `<code>` element.
      *
-     * @param {...any} args - The code, its language, and its metadata
+     * @param {string} code - The code
+     * @param {string} lang - Its language
+     * @param {Record<string, unknown>} [meta] - Its metadata
      */
-    highlightToHtml: (...args) => current().highlightToHtml(...args),
+    highlightToHtml: (code, lang, meta = {}) =>
+      current()
+        .codeToHtml(code, optionsFor(lang, meta))
+        // Shiki wraps the highlighted code in a <pre> and a <code>
+        .match(/<code>(.+?)<\/code>/s)[1],
 
     /**
      * Highlights code, returning a HAST tree.
      *
-     * @param {...any} args - The code, its language, and its metadata
+     * @param {string} code - The code
+     * @param {string} lang - Its language
+     * @param {Record<string, unknown>} [meta] - Its metadata
      */
-    highlightToHast: (...args) => current().highlightToHast(...args),
+    highlightToHast: (code, lang, meta = {}) =>
+      current().codeToHast(code, optionsFor(lang, meta)),
   };
 };
 
 /**
  * Creates a highlighter of every language Shiki bundles, with the given
  * options (see `./rehype.mjs`). Its Shiki instance is created on first use,
- * and the same options give the same highlighter.
+ * each bundled language is registered once code in it is highlighted, and the
+ * same options give the same highlighter.
  *
  * @param {import('./rehype.mjs').ShikiOptions} [options] - The options
  * @returns {Promise<SyntaxHighlighter>}

packages/core/src/plugins/type-annotations/highlighter.mjs

@@ -23,11 +23,11 @@ export const createTypeAnnotationHandler = getHighlighter => (state, node) => {
     return typeAnnotationToHast(state, node);
   }
 
-  const { shiki } = getHighlighter();
+  const { shiki, resolveLanguage } = getHighlighter();
   const [lightTheme, darkTheme] = shiki.getLoadedThemes();
 
   const root = shiki.codeToHast(node.value, {
-    lang: node.data?.typescript ? 'typescript' : 'text',
+    lang: resolveLanguage(node.data?.typescript ? 'typescript' : 'text'),
     themes: { light: lightTheme, dark: darkTheme },
     decorations: links.map(({ start, end, href }) => ({
       start,