diff --git a/.changeset/rehype-shiki-tokenization-cache.md b/.changeset/rehype-shiki-tokenization-cache.md new file mode 100644 index 0000000000000..e66ac42ea19a9 --- /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/cache` module exports `withTokenizationCache`, which does this for any Shiki highlighter. diff --git a/packages/rehype-shiki/src/__tests__/cache.test.mjs b/packages/rehype-shiki/src/__tests__/cache.test.mjs new file mode 100644 index 0000000000000..ea0de9bd58571 --- /dev/null +++ b/packages/rehype-shiki/src/__tests__/cache.test.mjs @@ -0,0 +1,265 @@ +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 '../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('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'); + + 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/__tests__/highlighter.test.mjs b/packages/rehype-shiki/src/__tests__/highlighter.test.mjs index 49479bedc367f..15bac6e28c3c1 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('../cache.mjs', { + namedExports: { withTokenizationCache: shiki => shiki }, +}); + mock.module('@shikijs/engine-javascript', { namedExports: { createJavaScriptRegexEngine: () => ({}) }, }); @@ -83,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/cache.mjs b/packages/rehype-shiki/src/cache.mjs new file mode 100644 index 0000000000000..8e6e136d5c3bd --- /dev/null +++ b/packages/rehype-shiki/src/cache.mjs @@ -0,0 +1,217 @@ +import { + codeToHast, + codeToHtml, + codeToTokens, + codeToTokensWithThemes, +} from '@shikijs/core'; +import { EncodedTokenMetadata } from 'shiki/textmate'; + +// 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 {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} + */ +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 + ); + + 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) { + /** @type {any} */ + const grammar = primitive.getLanguage(name); + + if ( + !grammar._basicScopeAttributesProvider || + grammar._tokenTypeMatchers?.length !== 0 + ) { + return grammar; + } + + return { + 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), + }; +}; diff --git a/packages/rehype-shiki/src/highlighter.mjs b/packages/rehype-shiki/src/highlighter.mjs index 52920e33cbba0..9b1b7182d3297 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/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. @@ -18,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'; /** @@ -57,12 +68,12 @@ const createHighlighter = ({ coreOptions = {}, highlighterOptions = {} }) => { themes: Object.values(DEFAULT_THEMES), ...coreOptions, }; - const shiki = createHighlighterCoreSync(options); + const shiki = withTokenizationCache(createHighlighterCoreSync(options)); const themeOptions = usesCustomTheme ? {} : coreOptions.themes ? { theme: options.themes[0] } - : { themes: DEFAULT_THEMES, defaultColor: 'light' }; + : DEFAULT_THEME_OPTIONS; const loadedLanguages = new Set( shiki.getLoadedLanguages().map(lang => lang.toLowerCase())