perf(mobile): highlight source files in small batches that keep grammar state

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)}";`;
 

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<ReadonlyArray<ReviewHighlightedToken>> = [];
-  const shortLineBatch: string[] = [];
-
-  const flushShortLineBatch = async (): Promise<void> => {
-    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;
 }