From 33afd5a5aa361890c9ce22479e8f52f6b18248aa Mon Sep 17 00:00:00 2001 From: Julius Marminge Date: Wed, 7 Oct 2026 22:48:58 -0700 Subject: [PATCH] perf(mobile): highlight source files in small batches that keep grammar state (#16729) Co-authored-by: Claude Opus 5.5 (1M context) --- .../review/shikiReviewHighlighter.test.ts | 18 +++++ .../features/review/shikiReviewHighlighter.ts | 74 +++++++++++-------- 2 files changed, 61 insertions(+), 31 deletions(-) diff --git a/apps/mobile/src/features/review/shikiReviewHighlighter.test.ts b/apps/mobile/src/features/review/shikiReviewHighlighter.test.ts index cfb28051cb1225767744cc89008186df481c8b28..dc8de530001cc8b699213243e6183bd664290305 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 9b69c8d552ea2dad0345586df7aec616f27d9195..b0ab139d0d279f1dd1069a4233cff1a53b6c615e 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; } -- 2.54.0