From a07ab77e0ef5902b01f93058fb6a0bf202d1af60 Mon Sep 17 00:00:00 2001 From: Claudio Wunder Date: Sat, 31 Oct 2026 14:35:51 +0100 Subject: [PATCH 1/4] 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 --- .../src/__tests__/highlighter.test.mjs | 4 + .../src/__tests__/tokenization-cache.test.mjs | 236 ++++++++++++++++++ packages/rehype-shiki/src/highlighter.mjs | 4 +- .../rehype-shiki/src/tokenization-cache.mjs | 236 ++++++++++++++++++ 4 files changed, 479 insertions(+), 1 deletion(-) create mode 100644 packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs create mode 100644 packages/rehype-shiki/src/tokenization-cache.mjs diff --git a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs index 49479bedc367f..233710139c5b3 100644 --- a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs +++ b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs @@ -17,6 +17,10 @@ mock.module('@shikijs/core', { }, }); +mock.module('../tokenization-cache.mjs', { + namedExports: { withTokenizationCache: shiki => shiki }, +}); + mock.module('@shikijs/engine-javascript', { namedExports: { createJavaScriptRegexEngine: () => ({}) }, }); diff --git a/packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs b/packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs new file mode 100644 index 0000000000000..7e602bbcc201a --- /dev/null +++ b/packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs @@ -0,0 +1,236 @@ +import assert from 'node:assert/strict'; +import { describe, it } from 'node:test'; + +import { createHighlighterCoreSync } from '@shikijs/core'; +import { createJavaScriptRegexEngine } from '@shikijs/engine-javascript'; +import { createOnigurumaEngine } from '@shikijs/engine-oniguruma'; +import c from 'shiki/langs/c.mjs'; +import html from 'shiki/langs/html.mjs'; +import javascript from 'shiki/langs/javascript.mjs'; +import json from 'shiki/langs/json.mjs'; +import markdown from 'shiki/langs/markdown.mjs'; +import shellscript from 'shiki/langs/shellscript.mjs'; +import typescript from 'shiki/langs/typescript.mjs'; +import githubDark from 'shiki/themes/github-dark.mjs'; +import githubLight from 'shiki/themes/github-light-default.mjs'; +import nord from 'shiki/themes/nord.mjs'; + +import { withTokenizationCache } from '../tokenization-cache.mjs'; + +// Code that starts lines in states other than the first: open comments, +// template literals and embedded languages, and lines repeated in a state +const SAMPLES = { + javascript: [ + "import { readFile } from 'node:fs/promises';", + '', + '/**', + ' * Reads a file.', + ' */', + 'const read = async (path, { encoding = "utf8" } = {}) => {', + ' const text = `${path}:', + ' ${await readFile(path, encoding)}`;', + ' return /^[a-z]+$/i.test(text) ? text : null;', + '};', + '};', + '', + 'class Reader extends EventEmitter {', + ' #path = null;', + '}', + ].join('\n'), + typescript: [ + 'type Options = Partial & { signal?: AbortSignal };', + 'function f(value: T): Promise> {', + ' return Promise.resolve([value as T]);', + '}', + ].join('\n'), + markdown: [ + '# Title', + '', + 'Some **bold** text and `code`.', + '', + '```js', + 'const x = 1;', + '```', + '', + '- [link](https://nodejs.org)', + ].join('\n'), + html: [ + '', + '', + '

Hello

', + ].join('\n'), + c: '#include \n\nnapi_value Init(napi_env env, napi_value exports) {\n return exports;\n}', + shellscript: 'for f in *.md; do\n echo "$f" | tr a-z A-Z\ndone', + json: '{\n "name": "node",\n "version": [1, 2, 3],\n "private": true\n}', +}; + +const ENGINES = { + oniguruma: () => createOnigurumaEngine(import('shiki/wasm')), + javascript: () => createJavaScriptRegexEngine({ forgiving: true }), +}; + +const THEMES = { light: 'github-light-default', dark: 'nord' }; + +/** + * Creates a highlighter, and the same one caching each line's tokens. + * + * @param {import('@shikijs/core').RegexEngine} engine + * @returns {{ shiki: import('@shikijs/core').HighlighterCore, cached: import('@shikijs/core').HighlighterCore }} + */ +const createHighlighters = engine => { + const shiki = createHighlighterCoreSync({ + engine, + langs: [ + ...c, + ...html, + ...javascript, + ...json, + ...markdown, + ...shellscript, + ...typescript, + ], + themes: [githubLight, nord, githubDark], + }); + + return { shiki, cached: withTokenizationCache(shiki) }; +}; + +/** + * Creates options to highlight TypeScript in two themes by name, with a + * transformer that has it highlighted as JavaScript, as Twoslash does with + * `// @showEmit`. + * + * @param {() => void} preprocess - Called whenever the transformer runs + * @returns {import('@shikijs/core').CodeToHastOptions} + */ +const createTranspilingOptions = preprocess => ({ + lang: 'typescript', + themes: THEMES, + transformers: [ + { + preprocess() { + preprocess(); + this.options.lang = 'javascript'; + }, + }, + ], +}); + +describe('withTokenizationCache', async () => { + for (const [name, createEngine] of Object.entries(ENGINES)) { + const engine = await createEngine(); + const { shiki, cached } = createHighlighters(engine); + + it(`highlights as Shiki does, with the ${name} engine`, () => { + for (const [lang, code] of Object.entries(SAMPLES)) { + for (const options of [ + { lang, themes: THEMES, defaultColor: 'light' }, + { lang, themes: { ...THEMES, dim: 'github-dark' } }, + ]) { + for (const method of [ + 'codeToHtml', + 'codeToHast', + 'codeToTokens', + 'codeToTokensWithThemes', + ]) { + assert.deepStrictEqual( + cached[method](code, options), + shiki[method](code, options), + `${method} of ${lang}` + ); + } + } + } + }); + } + + const { shiki, cached } = createHighlighters(ENGINES.javascript()); + + const grammar = shiki.getInternalContext().getLanguage('javascript'); + + // Shiki tokenizes every line but the empty ones + const lines = SAMPLES.javascript.split('\n').filter(Boolean); + + it('tokenizes each line once for all the themes', t => { + const tokenizeLine = t.mock.method(grammar, 'tokenizeLine'); + const tokenizeLine2 = t.mock.method(grammar, 'tokenizeLine2'); + + cached.codeToHtml(SAMPLES.javascript, { + lang: 'javascript', + themes: { ...THEMES, dim: 'github-dark' }, + }); + + assert.equal(tokenizeLine.mock.callCount(), lines.length); + assert.equal(tokenizeLine2.mock.callCount(), 0); + }); + + it('highlights as usual with theme objects, or a single theme', t => { + const tokenizeLine = t.mock.method(grammar, 'tokenizeLine'); + + for (const options of [ + { lang: 'javascript', themes: { light: githubLight, dark: nord } }, + { lang: 'javascript', theme: 'nord' }, + ]) { + assert.equal( + cached.codeToHtml(SAMPLES.javascript, options), + shiki.codeToHtml(SAMPLES.javascript, options) + ); + } + + assert.equal(tokenizeLine.mock.callCount(), 0); + }); + + it('tokenizes code for each theme with a grammar that sets token types', t => { + t.mock.property(grammar, '_tokenTypeMatchers', [ + { matcher: () => false, type: 0 }, + ]); + + const tokenizeLine = t.mock.method(grammar, 'tokenizeLine'); + const tokenizeLine2 = t.mock.method(grammar, 'tokenizeLine2'); + const preprocess = t.mock.fn(); + + const html = cached.codeToHtml( + SAMPLES.javascript, + createTranspilingOptions(preprocess) + ); + + assert.equal(preprocess.mock.callCount(), 1); + assert.equal(tokenizeLine.mock.callCount(), 0); + assert.equal(tokenizeLine2.mock.callCount(), 2 * lines.length); + + assert.equal( + html, + shiki.codeToHtml(SAMPLES.javascript, createTranspilingOptions(preprocess)) + ); + }); + + it('tokenizes a line that stops at the time limit once for all the themes', t => { + const tokenize = grammar.tokenizeLine.bind(grammar); + + const tokenizeLine = t.mock.method(grammar, 'tokenizeLine', (...args) => ({ + ...tokenize(...args), + stoppedEarly: true, + })); + + const tokenizeLine2 = t.mock.method(grammar, 'tokenizeLine2'); + const preprocess = t.mock.fn(); + + const html = cached.codeToHtml( + SAMPLES.javascript, + createTranspilingOptions(preprocess) + ); + + assert.equal(preprocess.mock.callCount(), 1); + assert.equal(tokenizeLine.mock.callCount(), lines.length); + assert.equal(tokenizeLine2.mock.callCount(), 0); + + assert.equal( + html, + shiki.codeToHtml(SAMPLES.javascript, createTranspilingOptions(preprocess)) + ); + }); +}); diff --git a/packages/rehype-shiki/src/highlighter.mjs b/packages/rehype-shiki/src/highlighter.mjs index 52920e33cbba0..ff1dedfc13e91 100644 --- a/packages/rehype-shiki/src/highlighter.mjs +++ b/packages/rehype-shiki/src/highlighter.mjs @@ -2,6 +2,8 @@ import { createHighlighterCoreSync, isSpecialLang } from '@shikijs/core'; import shikiGitHubLightTheme from 'shiki/themes/github-light-default.mjs'; import shikiNordTheme from 'shiki/themes/nord.mjs'; +import { withTokenizationCache } from '#rs/tokenization-cache.mjs'; + const DEFAULT_DARK_THEME = { // We are updating this color because the background color and comment text color // in the Codebox component do not comply with accessibility standards. @@ -57,7 +59,7 @@ const createHighlighter = ({ coreOptions = {}, highlighterOptions = {} }) => { themes: Object.values(DEFAULT_THEMES), ...coreOptions, }; - const shiki = createHighlighterCoreSync(options); + const shiki = withTokenizationCache(createHighlighterCoreSync(options)); const themeOptions = usesCustomTheme ? {} : coreOptions.themes diff --git a/packages/rehype-shiki/src/tokenization-cache.mjs b/packages/rehype-shiki/src/tokenization-cache.mjs new file mode 100644 index 0000000000000..d435606748435 --- /dev/null +++ b/packages/rehype-shiki/src/tokenization-cache.mjs @@ -0,0 +1,236 @@ +import { + codeToHast, + codeToHtml, + codeToTokens, + codeToTokensWithThemes, +} from '@shikijs/core'; +import { EncodedTokenMetadata } from 'shiki/textmate'; + +/** + * A vscode-textmate grammar, with the internals its tokens' styles are + * computed from, which the interface Shiki types it with leaves out. + * + * @typedef {Object} Grammar + * @property {string} name + * @property {(line: string, stack: import('shiki/textmate').StateStack | null, timeLimit?: number) => import('shiki/textmate').ITokenizeLineResult} tokenizeLine + * @property {(scopeName: string) => { languageId: number, tokenType: number }} getMetadataForScope + * @property {{ getDefaults: () => Style, themeMatch: (scopePath: ScopePath) => Style | null }} themeProvider + * @property {{ matchesAlways: boolean, matchesNever: boolean, match: (scopes: Array) => boolean } | null} balancedBracketSelectors + * @property {{ getDefaultAttributes: () => { languageId: number, tokenType: number } }} _basicScopeAttributesProvider + * @property {Array} [_tokenTypeMatchers] + * + * @typedef {{ fontStyle: number, foregroundId: number, backgroundId: number }} Style + * @typedef {{ scopeName: string, parent: ScopePath | null }} ScopePath + */ + +// The metadata of each path of scopes, by grammar and theme +/** @type {WeakMap>>} */ +const metadataCache = new WeakMap(); + +/** + * Gets a map's value for a key, setting it first if it has none. + * + * @template K, V + * @param {Map | (K extends object ? WeakMap : never)} map + * @param {K} key + * @param {() => V} create + * @returns {V} + */ +const getOrSet = (map, key, create) => { + let value = map.get(key); + + if (value === undefined) { + value = create(); + map.set(key, value); + } + + return value; +}; + +/** + * Computes the metadata vscode-textmate gives a token with these scopes in the + * theme in effect: the theme's defaults, with each scope's attributes merged + * over them (`AttributedScopeStack`), and whether it holds balanced brackets + * (`LineTokens.produceFromScopes`). + * + * @param {Grammar} grammar + * @param {Array} scopes + * @returns {number} + */ +const computeMetadata = (grammar, scopes) => { + const { languageId, tokenType } = + grammar._basicScopeAttributesProvider.getDefaultAttributes(); + const defaults = grammar.themeProvider.getDefaults(); + + let metadata = EncodedTokenMetadata.set( + 0, + languageId, + tokenType, + null, + defaults.fontStyle, + defaults.foregroundId, + defaults.backgroundId + ); + + /** @type {ScopePath | null} */ + let path = null; + + for (const scopeName of scopes) { + path = { scopeName, parent: path }; + + const basic = grammar.getMetadataForScope(scopeName); + const style = grammar.themeProvider.themeMatch(path); + + metadata = EncodedTokenMetadata.set( + metadata, + basic.languageId, + basic.tokenType, + null, + style ? style.fontStyle : -1, + style ? style.foregroundId : 0, + style ? style.backgroundId : 0 + ); + } + + const brackets = grammar.balancedBracketSelectors; + + if ( + brackets && + (brackets.matchesAlways || + (!brackets.matchesNever && brackets.match(scopes))) + ) { + metadata = EncodedTokenMetadata.set(metadata, 0, 8, true, -1, 0, 0); + } + + return metadata; +}; + +/** + * Creates the primitive one highlight goes through, caching each line's + * tokens: Shiki tokenizes code once for each theme, each time from the same + * states, so the themes after the first only give the cached tokens their + * styles. A grammar with token types of its own is tokenized as usual. + * + * @param {import('@shikijs/core').ShikiPrimitive} primitive + * @returns {import('@shikijs/core').ShikiPrimitive} + */ +const createCachingPrimitive = primitive => { + // Each line's tokens, by grammar, then by the state the line starts in + /** @type {Map>>} */ + const lines = new Map(); + + /** @type {object} */ + let theme; + + return { + ...primitive, + + setTheme(name) { + const result = primitive.setTheme(name); + + theme = result.theme; + + return result; + }, + + getLanguage(name) { + const language = primitive.getLanguage(name); + const grammar = /** @type {Grammar} */ ( + /** @type {unknown} */ (language) + ); + + if ( + !grammar._basicScopeAttributesProvider || + grammar._tokenTypeMatchers?.length !== 0 + ) { + return language; + } + + return /** @type {any} */ ({ + name: grammar.name, + + tokenizeLine: grammar.tokenizeLine.bind(grammar), + + /** @type {import('@shikijs/core').Grammar['tokenizeLine2']} */ + tokenizeLine2(line, stack, timeLimit) { + const states = getOrSet(lines, grammar, () => new Map()); + + const { tokens, ruleStack, stoppedEarly } = getOrSet( + getOrSet(states, stack, () => new Map()), + line, + () => grammar.tokenizeLine(line, stack, timeLimit) + ); + + const metadata = getOrSet( + getOrSet(metadataCache, grammar, () => new WeakMap()), + theme, + () => new Map() + ); + + /** @type {Array} */ + const binary = []; + + for (const { startIndex, scopes } of tokens) { + const tokenMetadata = getOrSet(metadata, scopes.join(' '), () => + computeMetadata(grammar, scopes) + ); + + // Neighboring tokens with the same metadata are one + if (binary.at(-1) !== tokenMetadata) { + binary.push(startIndex, tokenMetadata); + } + } + + return { tokens: Uint32Array.from(binary), ruleStack, stoppedEarly }; + }, + }); + }, + }; +}; + +/** + * Makes a Shiki highlighter cache each line's tokens for all the themes it + * highlights code in, rather than tokenizing it again for each: matching a + * grammar doesn't depend on the theme, only the styles the matched scopes get + * do. That's when two or more themes are given by name (Shiki prepares a theme + * object again for every highlight), without `includeExplanation`; otherwise + * code is highlighted as usual. The grammar state a cached highlight ends in + * holds one theme's styles for every theme, so it only continues highlights + * made the same way. + * + * @template {import('@shikijs/core').HighlighterCore} T + * @param {T} shiki + * @returns {T} + */ +export const withTokenizationCache = shiki => { + const primitive = shiki.getInternalContext(); + + /** + * Wraps a highlighter method to highlight through the cache when it can. + * + * @template R + * @param {(primitive: import('@shikijs/core').ShikiPrimitive, code: string, options: any) => R} highlight - The `@shikijs/core` function behind the method + * @param {(code: string, options: any) => R} usual - The highlighter's own method + * @returns {(code: string, options: any) => R} + */ + const cached = (highlight, usual) => (code, options) => { + const themes = Object.values(options?.themes ?? {}).filter(Boolean); + + return themes.length > 1 && + themes.every(theme => typeof theme === 'string') && + !options.includeExplanation + ? highlight(createCachingPrimitive(primitive), code, options) + : usual(code, options); + }; + + return { + ...shiki, + codeToTokensWithThemes: cached( + codeToTokensWithThemes, + shiki.codeToTokensWithThemes + ), + codeToTokens: cached(codeToTokens, shiki.codeToTokens), + codeToHast: cached(codeToHast, shiki.codeToHast), + codeToHtml: cached(codeToHtml, shiki.codeToHtml), + }; +}; From 61e78d9689253550e595cdc30f3a15f2a9ae4aaf Mon Sep 17 00:00:00 2001 From: Claudio Wunder Date: Sat, 31 Oct 2026 14:35:51 +0100 Subject: [PATCH 2/4] perf(rehype-shiki): give the default themes by name The default themes were passed as theme objects, which Shiki normalizes, and parses into a new vscode-textmate theme, on every highlight, and whose tokens can't be cached for both. They're registered with the highlighter anyway, so they're now given by name, in a `DEFAULT_THEME_OPTIONS` constant. The output is the same. Given by name, a theme's own color replacements no longer apply to the code block's foreground and background, but the default themes' only replacement (`#616e88`) is neither. Assisted-by: Claude Opus 5.5 --- .changeset/rehype-shiki-tokenization-cache.md | 5 +++++ .../rehype-shiki/src/__tests__/highlighter.test.mjs | 5 ++++- packages/rehype-shiki/src/highlighter.mjs | 11 ++++++++++- 3 files changed, 19 insertions(+), 2 deletions(-) create mode 100644 .changeset/rehype-shiki-tokenization-cache.md diff --git a/.changeset/rehype-shiki-tokenization-cache.md b/.changeset/rehype-shiki-tokenization-cache.md new file mode 100644 index 0000000000000..d5ae2d6426acd --- /dev/null +++ b/.changeset/rehype-shiki-tokenization-cache.md @@ -0,0 +1,5 @@ +--- +'@node-core/rehype-shiki': minor +--- + +Cache each line's tokens for all the themes code is highlighted in when two or more are given by name, as the default themes now are, instead of tokenizing it again for each. The new `@node-core/rehype-shiki/tokenization-cache` module exports `withTokenizationCache`, which does this for any Shiki highlighter. diff --git a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs index 233710139c5b3..15eecf92784db 100644 --- a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs +++ b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs @@ -87,7 +87,10 @@ describe('createHighlighter', async () => { highlighter.highlightToHtml('const x = 1;', 'javascript'); const [, options] = mockShiki.codeToHtml.mock.calls.at(-1).arguments; - assert.deepStrictEqual(Object.keys(options.themes), ['light', 'dark']); + assert.deepStrictEqual(options.themes, { + light: 'github-light-default', + dark: 'nord', + }); assert.strictEqual(options.defaultColor, 'light'); }); }); diff --git a/packages/rehype-shiki/src/highlighter.mjs b/packages/rehype-shiki/src/highlighter.mjs index ff1dedfc13e91..73a9003c586c7 100644 --- a/packages/rehype-shiki/src/highlighter.mjs +++ b/packages/rehype-shiki/src/highlighter.mjs @@ -20,6 +20,15 @@ const DEFAULT_THEMES = { dark: DEFAULT_DARK_THEME, }; +// The default themes by name: Shiki prepares a theme object again for every +// highlight, and only caches tokens for all the themes given by name +const DEFAULT_THEME_OPTIONS = { + themes: Object.fromEntries( + Object.entries(DEFAULT_THEMES).map(([color, { name }]) => [color, name]) + ), + defaultColor: 'light', +}; + const FALLBACK_LANGUAGE = 'text'; /** @@ -64,7 +73,7 @@ const createHighlighter = ({ coreOptions = {}, highlighterOptions = {} }) => { ? {} : coreOptions.themes ? { theme: options.themes[0] } - : { themes: DEFAULT_THEMES, defaultColor: 'light' }; + : DEFAULT_THEME_OPTIONS; const loadedLanguages = new Set( shiki.getLoadedLanguages().map(lang => lang.toLowerCase()) From f10f4063a63d0a5260eb17637171e61183ba65cc Mon Sep 17 00:00:00 2001 From: Claudio Wunder Date: Sat, 31 Oct 2026 23:16:08 +0100 Subject: [PATCH 3/4] refactor(rehype-shiki): rename the tokenization cache module `tokenization-cache.mjs` is now `tokenizationCache.mjs`, after review. Its grammar typedefs go too: Shiki doesn't type the grammar internals it reads, which vscode-textmate keeps private, and they leaked into the published types. A new test continues a cached highlight from inside a comment and a template literal, and compares it with Shiki continuing its own. Assisted-by: Claude Opus 5.5 --- .changeset/rehype-shiki-tokenization-cache.md | 2 +- .../src/__tests__/highlighter.test.mjs | 2 +- ...he.test.mjs => tokenizationCache.test.mjs} | 31 +++++++++++++++- packages/rehype-shiki/src/highlighter.mjs | 2 +- ...zation-cache.mjs => tokenizationCache.mjs} | 37 +++++-------------- 5 files changed, 42 insertions(+), 32 deletions(-) rename packages/rehype-shiki/src/__tests__/{tokenization-cache.test.mjs => tokenizationCache.test.mjs} (88%) rename packages/rehype-shiki/src/{tokenization-cache.mjs => tokenizationCache.mjs} (80%) diff --git a/.changeset/rehype-shiki-tokenization-cache.md b/.changeset/rehype-shiki-tokenization-cache.md index d5ae2d6426acd..5b2b531ffa191 100644 --- a/.changeset/rehype-shiki-tokenization-cache.md +++ b/.changeset/rehype-shiki-tokenization-cache.md @@ -2,4 +2,4 @@ '@node-core/rehype-shiki': minor --- -Cache each line's tokens for all the themes code is highlighted in when two or more are given by name, as the default themes now are, instead of tokenizing it again for each. The new `@node-core/rehype-shiki/tokenization-cache` module exports `withTokenizationCache`, which does this for any Shiki highlighter. +Cache each line's tokens for all the themes code is highlighted in when two or more are given by name, as the default themes now are, instead of tokenizing it again for each. The new `@node-core/rehype-shiki/tokenizationCache` module exports `withTokenizationCache`, which does this for any Shiki highlighter. diff --git a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs index 15eecf92784db..6d9f1930b019c 100644 --- a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs +++ b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs @@ -17,7 +17,7 @@ mock.module('@shikijs/core', { }, }); -mock.module('../tokenization-cache.mjs', { +mock.module('../tokenizationCache.mjs', { namedExports: { withTokenizationCache: shiki => shiki }, }); diff --git a/packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs b/packages/rehype-shiki/src/__tests__/tokenizationCache.test.mjs similarity index 88% rename from packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs rename to packages/rehype-shiki/src/__tests__/tokenizationCache.test.mjs index 7e602bbcc201a..010f9fc21c756 100644 --- a/packages/rehype-shiki/src/__tests__/tokenization-cache.test.mjs +++ b/packages/rehype-shiki/src/__tests__/tokenizationCache.test.mjs @@ -15,7 +15,7 @@ import githubDark from 'shiki/themes/github-dark.mjs'; import githubLight from 'shiki/themes/github-light-default.mjs'; import nord from 'shiki/themes/nord.mjs'; -import { withTokenizationCache } from '../tokenization-cache.mjs'; +import { withTokenizationCache } from '../tokenizationCache.mjs'; // Code that starts lines in states other than the first: open comments, // template literals and embedded languages, and lines repeated in a state @@ -168,6 +168,35 @@ describe('withTokenizationCache', async () => { assert.equal(tokenizeLine2.mock.callCount(), 0); }); + it('continues a highlight as Shiki does, from inside a comment or a template literal', () => { + const options = { + lang: 'javascript', + themes: { ...THEMES, dim: 'github-dark' }, + }; + + const sample = SAMPLES.javascript.split('\n'); + + for (const at of [3, 7]) { + const start = sample.slice(0, at).join('\n'); + const end = sample.slice(at).join('\n'); + + const { grammarState } = cached.codeToTokens(start, options); + const { grammarState: shikiState } = shiki.codeToTokens(start, options); + + for (const method of [ + 'codeToHtml', + 'codeToHast', + 'codeToTokensWithThemes', + ]) { + assert.deepStrictEqual( + cached[method](end, { ...options, grammarState }), + shiki[method](end, { ...options, grammarState: shikiState }), + `${method} from line ${at}` + ); + } + } + }); + it('highlights as usual with theme objects, or a single theme', t => { const tokenizeLine = t.mock.method(grammar, 'tokenizeLine'); diff --git a/packages/rehype-shiki/src/highlighter.mjs b/packages/rehype-shiki/src/highlighter.mjs index 73a9003c586c7..49766d96c8497 100644 --- a/packages/rehype-shiki/src/highlighter.mjs +++ b/packages/rehype-shiki/src/highlighter.mjs @@ -2,7 +2,7 @@ import { createHighlighterCoreSync, isSpecialLang } from '@shikijs/core'; import shikiGitHubLightTheme from 'shiki/themes/github-light-default.mjs'; import shikiNordTheme from 'shiki/themes/nord.mjs'; -import { withTokenizationCache } from '#rs/tokenization-cache.mjs'; +import { withTokenizationCache } from '#rs/tokenizationCache.mjs'; const DEFAULT_DARK_THEME = { // We are updating this color because the background color and comment text color diff --git a/packages/rehype-shiki/src/tokenization-cache.mjs b/packages/rehype-shiki/src/tokenizationCache.mjs similarity index 80% rename from packages/rehype-shiki/src/tokenization-cache.mjs rename to packages/rehype-shiki/src/tokenizationCache.mjs index d435606748435..8e6e136d5c3bd 100644 --- a/packages/rehype-shiki/src/tokenization-cache.mjs +++ b/packages/rehype-shiki/src/tokenizationCache.mjs @@ -6,25 +6,8 @@ import { } from '@shikijs/core'; import { EncodedTokenMetadata } from 'shiki/textmate'; -/** - * A vscode-textmate grammar, with the internals its tokens' styles are - * computed from, which the interface Shiki types it with leaves out. - * - * @typedef {Object} Grammar - * @property {string} name - * @property {(line: string, stack: import('shiki/textmate').StateStack | null, timeLimit?: number) => import('shiki/textmate').ITokenizeLineResult} tokenizeLine - * @property {(scopeName: string) => { languageId: number, tokenType: number }} getMetadataForScope - * @property {{ getDefaults: () => Style, themeMatch: (scopePath: ScopePath) => Style | null }} themeProvider - * @property {{ matchesAlways: boolean, matchesNever: boolean, match: (scopes: Array) => boolean } | null} balancedBracketSelectors - * @property {{ getDefaultAttributes: () => { languageId: number, tokenType: number } }} _basicScopeAttributesProvider - * @property {Array} [_tokenTypeMatchers] - * - * @typedef {{ fontStyle: number, foregroundId: number, backgroundId: number }} Style - * @typedef {{ scopeName: string, parent: ScopePath | null }} ScopePath - */ - // The metadata of each path of scopes, by grammar and theme -/** @type {WeakMap>>} */ +/** @type {WeakMap>>} */ const metadataCache = new WeakMap(); /** @@ -53,7 +36,8 @@ const getOrSet = (map, key, create) => { * over them (`AttributedScopeStack`), and whether it holds balanced brackets * (`LineTokens.produceFromScopes`). * - * @param {Grammar} grammar + * @param {any} grammar - A vscode-textmate grammar, whose internals this reads: + * Shiki types it without them, and vscode-textmate keeps them private * @param {Array} scopes * @returns {number} */ @@ -72,7 +56,6 @@ const computeMetadata = (grammar, scopes) => { defaults.backgroundId ); - /** @type {ScopePath | null} */ let path = null; for (const scopeName of scopes) { @@ -116,7 +99,7 @@ const computeMetadata = (grammar, scopes) => { */ const createCachingPrimitive = primitive => { // Each line's tokens, by grammar, then by the state the line starts in - /** @type {Map>>} */ + /** @type {Map>>} */ const lines = new Map(); /** @type {object} */ @@ -134,19 +117,17 @@ const createCachingPrimitive = primitive => { }, getLanguage(name) { - const language = primitive.getLanguage(name); - const grammar = /** @type {Grammar} */ ( - /** @type {unknown} */ (language) - ); + /** @type {any} */ + const grammar = primitive.getLanguage(name); if ( !grammar._basicScopeAttributesProvider || grammar._tokenTypeMatchers?.length !== 0 ) { - return language; + return grammar; } - return /** @type {any} */ ({ + return { name: grammar.name, tokenizeLine: grammar.tokenizeLine.bind(grammar), @@ -183,7 +164,7 @@ const createCachingPrimitive = primitive => { return { tokens: Uint32Array.from(binary), ruleStack, stoppedEarly }; }, - }); + }; }, }; }; From 775c0a288b7c22ee76523a6ea1b8c9e5cbc8204a Mon Sep 17 00:00:00 2001 From: Claudio Wunder Date: Sat, 31 Oct 2026 23:19:19 +0100 Subject: [PATCH 4/4] refactor(rehype-shiki): name the cache module cache.mjs Assisted-by: Claude Opus 5.5 --- .changeset/rehype-shiki-tokenization-cache.md | 2 +- .../__tests__/{tokenizationCache.test.mjs => cache.test.mjs} | 2 +- packages/rehype-shiki/src/__tests__/highlighter.test.mjs | 2 +- packages/rehype-shiki/src/{tokenizationCache.mjs => cache.mjs} | 0 packages/rehype-shiki/src/highlighter.mjs | 2 +- 5 files changed, 4 insertions(+), 4 deletions(-) rename packages/rehype-shiki/src/__tests__/{tokenizationCache.test.mjs => cache.test.mjs} (99%) rename packages/rehype-shiki/src/{tokenizationCache.mjs => cache.mjs} (100%) diff --git a/.changeset/rehype-shiki-tokenization-cache.md b/.changeset/rehype-shiki-tokenization-cache.md index 5b2b531ffa191..e66ac42ea19a9 100644 --- a/.changeset/rehype-shiki-tokenization-cache.md +++ b/.changeset/rehype-shiki-tokenization-cache.md @@ -2,4 +2,4 @@ '@node-core/rehype-shiki': minor --- -Cache each line's tokens for all the themes code is highlighted in when two or more are given by name, as the default themes now are, instead of tokenizing it again for each. The new `@node-core/rehype-shiki/tokenizationCache` module exports `withTokenizationCache`, which does this for any Shiki highlighter. +Cache each line's tokens for all the themes code is highlighted in when two or more are given by name, as the default themes now are, instead of tokenizing it again for each. The new `@node-core/rehype-shiki/cache` module exports `withTokenizationCache`, which does this for any Shiki highlighter. diff --git a/packages/rehype-shiki/src/__tests__/tokenizationCache.test.mjs b/packages/rehype-shiki/src/__tests__/cache.test.mjs similarity index 99% rename from packages/rehype-shiki/src/__tests__/tokenizationCache.test.mjs rename to packages/rehype-shiki/src/__tests__/cache.test.mjs index 010f9fc21c756..ea0de9bd58571 100644 --- a/packages/rehype-shiki/src/__tests__/tokenizationCache.test.mjs +++ b/packages/rehype-shiki/src/__tests__/cache.test.mjs @@ -15,7 +15,7 @@ import githubDark from 'shiki/themes/github-dark.mjs'; import githubLight from 'shiki/themes/github-light-default.mjs'; import nord from 'shiki/themes/nord.mjs'; -import { withTokenizationCache } from '../tokenizationCache.mjs'; +import { withTokenizationCache } from '../cache.mjs'; // Code that starts lines in states other than the first: open comments, // template literals and embedded languages, and lines repeated in a state diff --git a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs index 6d9f1930b019c..15bac6e28c3c1 100644 --- a/packages/rehype-shiki/src/__tests__/highlighter.test.mjs +++ b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs @@ -17,7 +17,7 @@ mock.module('@shikijs/core', { }, }); -mock.module('../tokenizationCache.mjs', { +mock.module('../cache.mjs', { namedExports: { withTokenizationCache: shiki => shiki }, }); diff --git a/packages/rehype-shiki/src/tokenizationCache.mjs b/packages/rehype-shiki/src/cache.mjs similarity index 100% rename from packages/rehype-shiki/src/tokenizationCache.mjs rename to packages/rehype-shiki/src/cache.mjs diff --git a/packages/rehype-shiki/src/highlighter.mjs b/packages/rehype-shiki/src/highlighter.mjs index 49766d96c8497..9b1b7182d3297 100644 --- a/packages/rehype-shiki/src/highlighter.mjs +++ b/packages/rehype-shiki/src/highlighter.mjs @@ -2,7 +2,7 @@ import { createHighlighterCoreSync, isSpecialLang } from '@shikijs/core'; import shikiGitHubLightTheme from 'shiki/themes/github-light-default.mjs'; import shikiNordTheme from 'shiki/themes/nord.mjs'; -import { withTokenizationCache } from '#rs/tokenizationCache.mjs'; +import { withTokenizationCache } from '#rs/cache.mjs'; const DEFAULT_DARK_THEME = { // We are updating this color because the background color and comment text color