Skip to content

Commit 95b8cdf

Browse files
committed
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. Importing `@node-core/rehype-shiki`'s plugin still imports every grammar until nodejs/nodejs.org#9212 is released. The grammars now come from doc-kit's own `shiki` dependency (4.4.3) rather than the copy `@node-core/rehype-shiki` pins (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>
1 parent 7238bb1 commit 95b8cdf

5 files changed

Lines changed: 217 additions & 39 deletions

File tree

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,57 @@ describe('createHighlighter', () => {
8686
);
8787
});
8888

89+
it('registers a bundled language once code in it is highlighted', async () => {
90+
const highlighter = await createHighlighter({
91+
langAlias: { py: 'python' },
92+
});
93+
const loaded = () => highlighter.shiki.getLoadedLanguages();
94+
95+
assert.ok(!loaded().includes('python'));
96+
assert.ok(!loaded().includes('javascript'));
97+
98+
// By its name, an alias of its own, or one of the options
99+
assert.equal(highlighter.resolveLanguage('py'), 'py');
100+
assert.equal(highlighter.resolveLanguage('mjs'), 'mjs');
101+
102+
assert.ok(loaded().includes('python'));
103+
assert.ok(loaded().includes('javascript'));
104+
assert.ok(loaded().includes('cjs'));
105+
106+
assert.equal(highlighter.resolveLanguage(undefined), 'text');
107+
assert.equal(highlighter.resolveLanguage('plaintext'), 'plaintext');
108+
});
109+
110+
it('lists the bundled languages without registering them', async () => {
111+
const highlighter = await createHighlighter({ langs: [grammar] });
112+
113+
assert.deepStrictEqual(
114+
highlighter.langs.find(({ name }) => name === 'rust'),
115+
{ name: 'rust', displayName: 'Rust', aliases: ['rs'] }
116+
);
117+
assert.equal(highlighter.langs.at(-1), grammar);
118+
assert.ok(!highlighter.shiki.getLoadedLanguages().includes('rust'));
119+
});
120+
121+
it('registers the bundled languages a language embeds', async () => {
122+
const highlighter = await createHighlighter({
123+
langs: [
124+
{
125+
...grammar,
126+
name: 'oxcscript',
127+
scopeName: 'source.oxcscript',
128+
embeddedLangs: ['javascript'],
129+
patterns: [{ include: 'source.js' }],
130+
},
131+
],
132+
});
133+
134+
assert.match(
135+
highlighter.highlightToHtml('const on = 1', 'oxcscript'),
136+
/--shiki-dark/
137+
);
138+
});
139+
89140
it('gives the same highlighter for the same options', async () => {
90141
const highlighter = await createHighlighter({ langs: [grammar] });
91142

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

Lines changed: 143 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,24 @@
11
'use strict';
22

3+
import { createRequire } from 'node:module';
34
import { endianness } from 'node:os';
45

5-
import { LANGS } from '@node-core/rehype-shiki';
66
import createSyntaxHighlighter from '@node-core/rehype-shiki/highlighter';
7+
import { isSpecialLang } from 'shiki/core';
8+
import { bundledLanguagesInfo } from 'shiki/langs';
79
import { bundledThemes } from 'shiki/themes';
810

911
import { importFromURL } from '#utils/loaders.mjs';
1012
import { getMarkdownPlugins } from '#utils/markdown/plugins.mjs';
1113

14+
const require = createRequire(import.meta.url);
15+
16+
/**
17+
* A language a highlighter highlights.
18+
*
19+
* @typedef {Pick<import('shiki').LanguageRegistration, 'name' | 'displayName' | 'aliases'>} Language
20+
*/
21+
1222
/**
1323
* A syntax highlighter, creating its Shiki instance on first use.
1424
*
@@ -17,7 +27,7 @@ import { getMarkdownPlugins } from '#utils/markdown/plugins.mjs';
1727
* @property {(languageId?: string) => string} resolveLanguage - Resolves a language, falling back to plain text for unknown ones
1828
* @property {(code: string, lang: string, meta?: Record<string, unknown>) => string} highlightToHtml - Highlights code, returning the inner HTML of its `<code>` element
1929
* @property {(code: string, lang: string, meta?: Record<string, unknown>) => ReturnType<import('shiki').HighlighterCore['codeToHast']>} highlightToHast - Highlights code, returning a HAST tree
20-
* @property {Array<import('shiki').LanguageRegistration>} langs - The languages it highlights
30+
* @property {Array<Language>} langs - The languages it highlights
2131
*/
2232

2333
/**
@@ -53,6 +63,54 @@ let engine;
5363
// The highlighters of the options given, by their JSON
5464
const highlighters = new Map();
5565

66+
// What code in an unknown language is highlighted as
67+
const FALLBACK_LANGUAGE = 'text';
68+
69+
// The languages Shiki bundles. Registering a grammar takes time and memory,
70+
// and slows down every highlight after it, so each one is only imported and
71+
// registered once code in it is highlighted (see `resolveLanguage`).
72+
const BUNDLED_LANGUAGES = bundledLanguagesInfo.map(({ id, name, aliases }) => ({
73+
name: id,
74+
displayName: name,
75+
aliases,
76+
}));
77+
78+
// The bundled language each of their names and aliases stands for, as code
79+
// names its language by either
80+
const BUNDLED_NAMES = new Map(
81+
BUNDLED_LANGUAGES.flatMap(({ name, aliases = [] }) =>
82+
[name, ...aliases].map(alias => [alias, name])
83+
)
84+
);
85+
86+
/**
87+
* Imports a bundled language: its grammar, and those of the languages it
88+
* embeds.
89+
*
90+
* @param {string} name - A name or alias of the language
91+
* @returns {Array<import('shiki').LanguageRegistration>}
92+
*/
93+
const importBundledLanguage = name =>
94+
require(`shiki/langs/${BUNDLED_NAMES.get(name)}.mjs`).default;
95+
96+
/**
97+
* Adds the bundled languages that languages embed, which Shiki registers
98+
* along with them.
99+
*
100+
* @param {Array<import('shiki').LanguageRegistration>} langs - The languages
101+
* @returns {Array<import('shiki').LanguageRegistration>}
102+
*/
103+
const withEmbeddedLanguages = langs => {
104+
const names = new Set(langs.map(({ name }) => name));
105+
106+
const embedded = langs
107+
.flatMap(({ embeddedLangs = [] }) => embeddedLangs)
108+
.filter(name => !names.has(name) && BUNDLED_NAMES.has(name))
109+
.flatMap(importBundledLanguage);
110+
111+
return [...embedded, ...langs];
112+
};
113+
56114
/**
57115
* Imports a list of options, each given as it is, or as a module (a path
58116
* relative to the working directory, or a URL) default-exporting it, or a list
@@ -119,13 +177,12 @@ const importHighlighter = async ({
119177

120178
const coreOptions = {
121179
engine: regexEngine,
122-
langs: [...LANGS, ...importedLangs],
123-
// A copy, as Shiki adds the aliases of the languages it bundles to it
180+
// The bundled languages are registered as code uses them
181+
langs: withEmbeddedLanguages(importedLangs),
182+
// A copy, as Shiki adds the aliases of the languages it registers to it
124183
langAlias: { ...langAlias },
125184
};
126185

127-
const highlighterOptions = { transformers: importedTransformers };
128-
129186
// Without themes of its own, the highlighter has a default light and dark one
130187
if (themes) {
131188
const [light, dark] = await Promise.all([
@@ -134,58 +191,116 @@ const importHighlighter = async ({
134191
]);
135192

136193
coreOptions.themes = [light, dark];
137-
highlighterOptions.themes = { light: light.name, dark: dark.name };
138-
highlighterOptions.defaultColor = 'light';
139194
}
140195

141-
let highlighter;
196+
let shiki;
197+
let highlightOptions;
142198

143199
/**
144-
* Gives the Shiki highlighter, creating it on first use.
200+
* Gives the Shiki instance, creating it on first use.
201+
*
202+
* @returns {import('shiki').HighlighterCore}
145203
*/
146-
const current = () =>
147-
(highlighter ??= createSyntaxHighlighter({
148-
coreOptions,
149-
highlighterOptions,
150-
}));
204+
const current = () => {
205+
if (!shiki) {
206+
({ shiki } = createSyntaxHighlighter({ coreOptions }));
207+
208+
// The themes are given by name: Shiki keeps the themes it loaded parsed,
209+
// but parses a theme object it's given again for every highlight
210+
const [light, dark] = shiki.getLoadedThemes();
211+
212+
highlightOptions = {
213+
themes: { light, dark },
214+
defaultColor: 'light',
215+
transformers: importedTransformers,
216+
};
217+
}
218+
219+
return shiki;
220+
};
221+
222+
/**
223+
* Resolves a language, falling back to plain text for unknown ones. A
224+
* bundled language is registered the first time it's resolved.
225+
*
226+
* @param {string} [languageId]
227+
* @returns {string}
228+
*/
229+
const resolveLanguage = languageId => {
230+
if (!languageId) {
231+
return FALLBACK_LANGUAGE;
232+
}
233+
234+
const instance = current();
235+
const name = instance.resolveLangAlias(languageId.toLowerCase());
236+
237+
if (isSpecialLang(name) || instance.getLoadedLanguages().includes(name)) {
238+
return languageId;
239+
}
240+
241+
if (!BUNDLED_NAMES.has(name)) {
242+
return FALLBACK_LANGUAGE;
243+
}
244+
245+
instance.loadLanguageSync(importBundledLanguage(name));
246+
247+
return languageId;
248+
};
249+
250+
/**
251+
* The options Shiki highlights code with.
252+
*
253+
* @param {string} lang - The language of the code
254+
* @param {Record<string, unknown>} meta - Its metadata
255+
*/
256+
const optionsFor = (lang, meta) => ({
257+
lang: resolveLanguage(lang),
258+
...highlightOptions,
259+
meta,
260+
});
151261

152262
return {
153-
langs: coreOptions.langs,
263+
langs: [...BUNDLED_LANGUAGES, ...importedLangs],
154264

155265
/**
156266
* The Shiki instance.
157267
*/
158268
get shiki() {
159-
return current().shiki;
269+
return current();
160270
},
161271

162-
/**
163-
* Resolves a language, falling back to plain text for unknown ones.
164-
*
165-
* @param {string} [languageId]
166-
*/
167-
resolveLanguage: languageId => current().resolveLanguage(languageId),
272+
resolveLanguage,
168273

169274
/**
170275
* Highlights code, returning the inner HTML of its `<code>` element.
171276
*
172-
* @param {...any} args - The code, its language, and its metadata
277+
* @param {string} code - The code
278+
* @param {string} lang - Its language
279+
* @param {Record<string, unknown>} [meta] - Its metadata
173280
*/
174-
highlightToHtml: (...args) => current().highlightToHtml(...args),
281+
highlightToHtml: (code, lang, meta = {}) =>
282+
current()
283+
.codeToHtml(code, optionsFor(lang, meta))
284+
// Shiki wraps the highlighted code in a <pre> and a <code>
285+
.match(/<code>(.+?)<\/code>/s)[1],
175286

176287
/**
177288
* Highlights code, returning a HAST tree.
178289
*
179-
* @param {...any} args - The code, its language, and its metadata
290+
* @param {string} code - The code
291+
* @param {string} lang - Its language
292+
* @param {Record<string, unknown>} [meta] - Its metadata
180293
*/
181-
highlightToHast: (...args) => current().highlightToHast(...args),
294+
highlightToHast: (code, lang, meta = {}) =>
295+
current().codeToHast(code, optionsFor(lang, meta)),
182296
};
183297
};
184298

185299
/**
186300
* Creates a highlighter of every language Shiki bundles, with the given
187301
* options (see `./rehype.mjs`). Its Shiki instance is created on first use,
188-
* and the same options give the same highlighter.
302+
* each bundled language is registered once code in it is highlighted, and the
303+
* same options give the same highlighter.
189304
*
190305
* @param {import('./rehype.mjs').ShikiOptions} [options] - The options
191306
* @returns {Promise<SyntaxHighlighter>}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@ export const createTypeAnnotationHandler = getHighlighter => (state, node) => {
2323
return typeAnnotationToHast(state, node);
2424
}
2525

26-
const { shiki } = getHighlighter();
26+
const { shiki, resolveLanguage } = getHighlighter();
2727
const [lightTheme, darkTheme] = shiki.getLoadedThemes();
2828

2929
const root = shiki.codeToHast(node.value, {
30-
lang: node.data?.typescript ? 'typescript' : 'text',
30+
lang: resolveLanguage(node.data?.typescript ? 'typescript' : 'text'),
3131
themes: { light: lightTheme, dark: darkTheme },
3232
decorations: links.map(({ start, end, href }) => ({
3333
start,

‎packages/node-legacy/src/legacy-html/plugins/shiki.mjs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ function isCodeBlock(node) {
4141
* Creates a HAST transformer for Shiki which is used for transforming our codeboxes
4242
*/
4343
export default function rehypeShikiji() {
44+
const [lightTheme, darkTheme] = shikiConfig.themes;
45+
46+
// Loaded so they can be given by name, as Shiki parses a theme object it's
47+
// given again for every code block
48+
highlighter.shiki.loadThemeSync(lightTheme, darkTheme);
49+
4450
/**
4551
* @param {import('hast').Root} tree - The HAST tree to be transformed.
4652
*/
@@ -90,9 +96,9 @@ export default function rehypeShikiji() {
9096
const { children } = highlighter.shiki.codeToHast(
9197
preElement.children[0].value,
9298
{
93-
lang: languageId,
99+
lang: highlighter.resolveLanguage(languageId),
94100
// Allows support for dual themes (light, dark) for Shiki
95-
themes: { light: shikiConfig.themes[0], dark: shikiConfig.themes[1] },
101+
themes: { light: lightTheme.name, dark: darkTheme.name },
96102
}
97103
);
98104

‎packages/react/src/html/utils/__tests__/config.test.mjs‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,24 @@
11
import assert from 'node:assert/strict';
2+
import { createRequire } from 'node:module';
23
import { describe, it, mock } from 'node:test';
4+
import { fileURLToPath } from 'node:url';
35

46
import { setConfig } from '@doc-kit/core/utils/configuration/index.mjs';
57
import { loadMarkdownPlugins } from '@doc-kit/core/utils/markdown/plugins.mjs';
68
import { SemVer } from 'semver';
79

8-
mock.module('@node-core/rehype-shiki', {
10+
// The languages Shiki bundles, as `@doc-kit/core` imports them
11+
const core = createRequire(
12+
fileURLToPath(import.meta.resolve('@doc-kit/core/package.json'))
13+
);
14+
15+
mock.module(core.resolve('shiki/langs'), {
916
exports: {
10-
LANGS: [
11-
{ name: 'javascript', aliases: ['js'], displayName: 'JavaScript' },
12-
{ name: 'typescript', aliases: ['ts'], displayName: 'TypeScript' },
13-
{ name: 'python', displayName: 'Python' },
17+
bundledLanguagesInfo: [
18+
{ id: 'javascript', name: 'JavaScript', aliases: ['js'] },
19+
{ id: 'typescript', name: 'TypeScript', aliases: ['ts'] },
20+
{ id: 'python', name: 'Python' },
1421
],
15-
default: async () => ({}),
1622
},
1723
});
1824

@@ -44,7 +50,7 @@ const config = await setConfig({
4450
});
4551

4652
// Loading the real `html` generator would pull in the full rendering stack
47-
// (which the `rehype-shiki` mock above cannot satisfy), so its resolved
53+
// (which the `shiki/langs` mock above cannot satisfy), so its resolved
4854
// configuration is stubbed in directly.
4955
config.html = {
5056
...config.global,

0 commit comments

Comments
 (0)