Skip to content

Commit a07ab77

Browse files
ovflowdovflowd
committed
perf(rehype-shiki): cache tokens for all themes given by name
Shiki tokenizes code once for each theme it's highlighted in: every theme's pass runs the grammar over every line again, though matching a grammar doesn't depend on the theme, only the styles the matched scopes get do. Highlighting in a light and a dark theme, as doc-kit and the site do, tokenizes everything twice. `withTokenizationCache`, in the new `tokenization-cache` module, wraps a highlighter so that with two or more themes given by name each line's tokens are cached for all of them: the line is tokenized once, and each theme's token metadata is computed from the scopes found, as vscode-textmate computes it, and cached by scope path. Everything after that is Shiki's own code, so the output doesn't change: it's byte-identical on the 3,357 code blocks of Node.js's API docs with both regex engines, and in full doc-kit builds. Highlighting them takes 32% less time on Oniguruma and 43% less on the JavaScript engine. Theme objects, a single theme and `includeExplanation` take Shiki's usual path, as does a grammar with token types of its own. Assisted-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 53d777b commit a07ab77

4 files changed

Lines changed: 479 additions & 1 deletion

File tree

‎packages/rehype-shiki/src/__tests__/highlighter.test.mjs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@ mock.module('@shikijs/core', {
1717
},
1818
});
1919

20+
mock.module('../tokenization-cache.mjs', {
21+
namedExports: { withTokenizationCache: shiki => shiki },
22+
});
23+
2024
mock.module('@shikijs/engine-javascript', {
2125
namedExports: { createJavaScriptRegexEngine: () => ({}) },
2226
});

‎packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
import assert from 'node:assert/strict';
2+
import { describe, it } from 'node:test';
3+
4+
import { createHighlighterCoreSync } from '@shikijs/core';
5+
import { createJavaScriptRegexEngine } from '@shikijs/engine-javascript';
6+
import { createOnigurumaEngine } from '@shikijs/engine-oniguruma';
7+
import c from 'shiki/langs/c.mjs';
8+
import html from 'shiki/langs/html.mjs';
9+
import javascript from 'shiki/langs/javascript.mjs';
10+
import json from 'shiki/langs/json.mjs';
11+
import markdown from 'shiki/langs/markdown.mjs';
12+
import shellscript from 'shiki/langs/shellscript.mjs';
13+
import typescript from 'shiki/langs/typescript.mjs';
14+
import githubDark from 'shiki/themes/github-dark.mjs';
15+
import githubLight from 'shiki/themes/github-light-default.mjs';
16+
import nord from 'shiki/themes/nord.mjs';
17+
18+
import { withTokenizationCache } from '../tokenization-cache.mjs';
19+
20+
// Code that starts lines in states other than the first: open comments,
21+
// template literals and embedded languages, and lines repeated in a state
22+
const SAMPLES = {
23+
javascript: [
24+
"import { readFile } from 'node:fs/promises';",
25+
'',
26+
'/**',
27+
' * Reads a file.',
28+
' */',
29+
'const read = async (path, { encoding = "utf8" } = {}) => {',
30+
' const text = `${path}:',
31+
' ${await readFile(path, encoding)}`;',
32+
' return /^[a-z]+$/i.test(text) ? text : null;',
33+
'};',
34+
'};',
35+
'',
36+
'class Reader extends EventEmitter {',
37+
' #path = null;',
38+
'}',
39+
].join('\n'),
40+
typescript: [
41+
'type Options<T extends object = {}> = Partial<T> & { signal?: AbortSignal };',
42+
'function f<T>(value: T): Promise<Array<T>> {',
43+
' return Promise.resolve([value as T]);',
44+
'}',
45+
].join('\n'),
46+
markdown: [
47+
'# Title',
48+
'',
49+
'Some **bold** text and `code`.',
50+
'',
51+
'```js',
52+
'const x = 1;',
53+
'```',
54+
'',
55+
'- [link](https://nodejs.org)',
56+
].join('\n'),
57+
html: [
58+
'<style>',
59+
' body { color: #333; }',
60+
'</style>',
61+
'<script>',
62+
' document.body.append("hi");',
63+
'</script>',
64+
'<p class="note">Hello</p>',
65+
].join('\n'),
66+
c: '#include <node_api.h>\n\nnapi_value Init(napi_env env, napi_value exports) {\n return exports;\n}',
67+
shellscript: 'for f in *.md; do\n echo "$f" | tr a-z A-Z\ndone',
68+
json: '{\n "name": "node",\n "version": [1, 2, 3],\n "private": true\n}',
69+
};
70+
71+
const ENGINES = {
72+
oniguruma: () => createOnigurumaEngine(import('shiki/wasm')),
73+
javascript: () => createJavaScriptRegexEngine({ forgiving: true }),
74+
};
75+
76+
const THEMES = { light: 'github-light-default', dark: 'nord' };
77+
78+
/**
79+
* Creates a highlighter, and the same one caching each line's tokens.
80+
*
81+
* @param {import('@shikijs/core').RegexEngine} engine
82+
* @returns {{ shiki: import('@shikijs/core').HighlighterCore, cached: import('@shikijs/core').HighlighterCore }}
83+
*/
84+
const createHighlighters = engine => {
85+
const shiki = createHighlighterCoreSync({
86+
engine,
87+
langs: [
88+
...c,
89+
...html,
90+
...javascript,
91+
...json,
92+
...markdown,
93+
...shellscript,
94+
...typescript,
95+
],
96+
themes: [githubLight, nord, githubDark],
97+
});
98+
99+
return { shiki, cached: withTokenizationCache(shiki) };
100+
};
101+
102+
/**
103+
* Creates options to highlight TypeScript in two themes by name, with a
104+
* transformer that has it highlighted as JavaScript, as Twoslash does with
105+
* `// @showEmit`.
106+
*
107+
* @param {() => void} preprocess - Called whenever the transformer runs
108+
* @returns {import('@shikijs/core').CodeToHastOptions}
109+
*/
110+
const createTranspilingOptions = preprocess => ({
111+
lang: 'typescript',
112+
themes: THEMES,
113+
transformers: [
114+
{
115+
preprocess() {
116+
preprocess();
117+
this.options.lang = 'javascript';
118+
},
119+
},
120+
],
121+
});
122+
123+
describe('withTokenizationCache', async () => {
124+
for (const [name, createEngine] of Object.entries(ENGINES)) {
125+
const engine = await createEngine();
126+
const { shiki, cached } = createHighlighters(engine);
127+
128+
it(`highlights as Shiki does, with the ${name} engine`, () => {
129+
for (const [lang, code] of Object.entries(SAMPLES)) {
130+
for (const options of [
131+
{ lang, themes: THEMES, defaultColor: 'light' },
132+
{ lang, themes: { ...THEMES, dim: 'github-dark' } },
133+
]) {
134+
for (const method of [
135+
'codeToHtml',
136+
'codeToHast',
137+
'codeToTokens',
138+
'codeToTokensWithThemes',
139+
]) {
140+
assert.deepStrictEqual(
141+
cached[method](code, options),
142+
shiki[method](code, options),
143+
`${method} of ${lang}`
144+
);
145+
}
146+
}
147+
}
148+
});
149+
}
150+
151+
const { shiki, cached } = createHighlighters(ENGINES.javascript());
152+
153+
const grammar = shiki.getInternalContext().getLanguage('javascript');
154+
155+
// Shiki tokenizes every line but the empty ones
156+
const lines = SAMPLES.javascript.split('\n').filter(Boolean);
157+
158+
it('tokenizes each line once for all the themes', t => {
159+
const tokenizeLine = t.mock.method(grammar, 'tokenizeLine');
160+
const tokenizeLine2 = t.mock.method(grammar, 'tokenizeLine2');
161+
162+
cached.codeToHtml(SAMPLES.javascript, {
163+
lang: 'javascript',
164+
themes: { ...THEMES, dim: 'github-dark' },
165+
});
166+
167+
assert.equal(tokenizeLine.mock.callCount(), lines.length);
168+
assert.equal(tokenizeLine2.mock.callCount(), 0);
169+
});
170+
171+
it('highlights as usual with theme objects, or a single theme', t => {
172+
const tokenizeLine = t.mock.method(grammar, 'tokenizeLine');
173+
174+
for (const options of [
175+
{ lang: 'javascript', themes: { light: githubLight, dark: nord } },
176+
{ lang: 'javascript', theme: 'nord' },
177+
]) {
178+
assert.equal(
179+
cached.codeToHtml(SAMPLES.javascript, options),
180+
shiki.codeToHtml(SAMPLES.javascript, options)
181+
);
182+
}
183+
184+
assert.equal(tokenizeLine.mock.callCount(), 0);
185+
});
186+
187+
it('tokenizes code for each theme with a grammar that sets token types', t => {
188+
t.mock.property(grammar, '_tokenTypeMatchers', [
189+
{ matcher: () => false, type: 0 },
190+
]);
191+
192+
const tokenizeLine = t.mock.method(grammar, 'tokenizeLine');
193+
const tokenizeLine2 = t.mock.method(grammar, 'tokenizeLine2');
194+
const preprocess = t.mock.fn();
195+
196+
const html = cached.codeToHtml(
197+
SAMPLES.javascript,
198+
createTranspilingOptions(preprocess)
199+
);
200+
201+
assert.equal(preprocess.mock.callCount(), 1);
202+
assert.equal(tokenizeLine.mock.callCount(), 0);
203+
assert.equal(tokenizeLine2.mock.callCount(), 2 * lines.length);
204+
205+
assert.equal(
206+
html,
207+
shiki.codeToHtml(SAMPLES.javascript, createTranspilingOptions(preprocess))
208+
);
209+
});
210+
211+
it('tokenizes a line that stops at the time limit once for all the themes', t => {
212+
const tokenize = grammar.tokenizeLine.bind(grammar);
213+
214+
const tokenizeLine = t.mock.method(grammar, 'tokenizeLine', (...args) => ({
215+
...tokenize(...args),
216+
stoppedEarly: true,
217+
}));
218+
219+
const tokenizeLine2 = t.mock.method(grammar, 'tokenizeLine2');
220+
const preprocess = t.mock.fn();
221+
222+
const html = cached.codeToHtml(
223+
SAMPLES.javascript,
224+
createTranspilingOptions(preprocess)
225+
);
226+
227+
assert.equal(preprocess.mock.callCount(), 1);
228+
assert.equal(tokenizeLine.mock.callCount(), lines.length);
229+
assert.equal(tokenizeLine2.mock.callCount(), 0);
230+
231+
assert.equal(
232+
html,
233+
shiki.codeToHtml(SAMPLES.javascript, createTranspilingOptions(preprocess))
234+
);
235+
});
236+
});

‎packages/rehype-shiki/src/highlighter.mjs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ import { createHighlighterCoreSync, isSpecialLang } from '@shikijs/core';
22
import shikiGitHubLightTheme from 'shiki/themes/github-light-default.mjs';
33
import shikiNordTheme from 'shiki/themes/nord.mjs';
44

5+
import { withTokenizationCache } from '#rs/tokenization-cache.mjs';
6+
57
const DEFAULT_DARK_THEME = {
68
// We are updating this color because the background color and comment text color
79
// in the Codebox component do not comply with accessibility standards.
@@ -57,7 +59,7 @@ const createHighlighter = ({ coreOptions = {}, highlighterOptions = {} }) => {
5759
themes: Object.values(DEFAULT_THEMES),
5860
...coreOptions,
5961
};
60-
const shiki = createHighlighterCoreSync(options);
62+
const shiki = withTokenizationCache(createHighlighterCoreSync(options));
6163
const themeOptions = usesCustomTheme
6264
? {}
6365
: coreOptions.themes

0 commit comments

Comments
 (0)