diff --git a/apps/mobile/src/features/review/shikiReviewHighlighter.test.ts b/apps/mobile/src/features/review/shikiReviewHighlighter.test.ts index cfb28051cb12..dc8de530001c 100644 --- a/apps/mobile/src/features/review/shikiReviewHighlighter.test.ts +++ b/apps/mobile/src/features/review/shikiReviewHighlighter.test.ts @@ -27,6 +27,24 @@ describe("highlightSourceFile", () => { ); }); + it("keeps colors for a block comment that spans highlight batches", async () => { + // Long enough that the comment body crosses at least one batch boundary. + const body = Array.from({ length: 400 }, (_, index) => `const insideComment${index} = 1;`); + const highlighted = await highlightSourceFile({ + path: "example.ts", + contents: ["/*", ...body, "*/", "const after = 2;"].join("\n"), + theme: "dark", + }); + + const commentColors = new Set( + highlighted.slice(1, -2).flatMap((line) => line.map((token) => token.color)), + ); + expect(commentColors.size).toBe(1); + expect(highlighted.at(-1)?.map((token) => token.color)).not.toEqual( + highlighted[1]?.map((token) => token.color), + ); + }); + it("falls back to plain tokens for very long lines", async () => { const longLine = `const value = "${"a".repeat(1_100)}";`; diff --git a/apps/mobile/src/features/review/shikiReviewHighlighter.ts b/apps/mobile/src/features/review/shikiReviewHighlighter.ts index 9b69c8d552ea..b0ab139d0d27 100644 --- a/apps/mobile/src/features/review/shikiReviewHighlighter.ts +++ b/apps/mobile/src/features/review/shikiReviewHighlighter.ts @@ -1,4 +1,4 @@ -import { createHighlighterCore, type HighlighterCore } from "@shikijs/core"; +import { createHighlighterCore, type GrammarState, type HighlighterCore } from "@shikijs/core"; import { createJavaScriptRegexEngine } from "@shikijs/engine-javascript"; import bashLanguage from "@shikijs/langs/bash"; import javascriptLanguage from "@shikijs/langs/javascript"; @@ -50,7 +50,12 @@ const REVIEW_HIGHLIGHTER_ENGINE_PREFERENCE = resolveReviewHighlighterEnginePrefe REVIEW_HIGHLIGHTER_ENGINE_ENV_VALUE, ); const REVIEW_HIGHLIGHT_CHUNK_LINE_THRESHOLD = 8; -const REVIEW_HIGHLIGHT_CHUNK_SIZE = 200; +// Bounds each tokenizing call by characters, not lines, so long-line files +// still yield to touches and renders between short batches. +const REVIEW_HIGHLIGHT_CHUNK_CHARACTERS = 2_000; +// A yield waits for the next frame, so batches run back to back until about one +// frame of work has passed instead of yielding after every batch. +const REVIEW_HIGHLIGHT_YIELD_AFTER_MS = 16; const REVIEW_TOKENIZE_MAX_LINE_LENGTH = 1_000; const REVIEW_INITIAL_LANGUAGE_MODULES = [ bashLanguage, @@ -513,46 +518,53 @@ async function highlightLines( const highlighter = await getHighlighter(); const sourceLines = code.split("\n"); const highlightedLines: Array> = []; - const shortLineBatch: string[] = []; - - const flushShortLineBatch = async (): Promise => { - if (shortLineBatch.length === 0) { - return; - } - - const tokenLines = highlighter.codeToTokensBase(shortLineBatch.join("\n"), { - lang: language, - theme, - }); - highlightedLines.push(...normalizeHighlightedLines(tokenLines)); - shortLineBatch.length = 0; - }; - - for (let lineIndex = 0; lineIndex < sourceLines.length; lineIndex += 1) { - const line = sourceLines[lineIndex] ?? ""; - - if (line.length > REVIEW_TOKENIZE_MAX_LINE_LENGTH) { - await flushShortLineBatch(); - highlightedLines.push([{ content: line, color: null, fontStyle: null }]); + // Batches resume from the previous batch's grammar state, so a comment or + // template string that spans a batch boundary keeps its colors. + let grammarState: GrammarState | undefined; + let start = 0; + let sliceStartedAt = performance.now(); + + while (start < sourceLines.length) { + // A skipped line leaves its ending state unknown; resume from a fresh state. + if (sourceLines[start]!.length > REVIEW_TOKENIZE_MAX_LINE_LENGTH) { + highlightedLines.push([{ content: sourceLines[start]!, color: null, fontStyle: null }]); + grammarState = undefined; + start += 1; } else { - shortLineBatch.push(line); - } + let end = start; + let characters = 0; + while (end < sourceLines.length) { + const length = sourceLines[end]!.length; + if ( + length > REVIEW_TOKENIZE_MAX_LINE_LENGTH || + (end > start && characters + length + 1 > REVIEW_HIGHLIGHT_CHUNK_CHARACTERS) + ) { + break; + } + characters += length + (end > start ? 1 : 0); + end += 1; + } - if (shortLineBatch.length >= REVIEW_HIGHLIGHT_CHUNK_SIZE) { - await flushShortLineBatch(); + const tokenLines = highlighter.codeToTokensBase(sourceLines.slice(start, end).join("\n"), { + lang: language, + theme, + grammarState, + }); + grammarState = highlighter.getLastGrammarState(tokenLines); + highlightedLines.push(...normalizeHighlightedLines(tokenLines)); + start = end; } if ( sourceLines.length > REVIEW_HIGHLIGHT_CHUNK_LINE_THRESHOLD && - lineIndex + 1 < sourceLines.length && - (shortLineBatch.length === 0 || line.length > REVIEW_TOKENIZE_MAX_LINE_LENGTH) + start < sourceLines.length && + performance.now() - sliceStartedAt >= REVIEW_HIGHLIGHT_YIELD_AFTER_MS ) { await waitForNextFrame(); + sliceStartedAt = performance.now(); } } - await flushShortLineBatch(); - return highlightedLines; }