Isolate native diff highlight grammar state

apps/mobile/src/features/diffs/nativeReviewDiffHighlighter.test.ts

@@ -0,0 +1,189 @@
+import { describe, expect, it } from "vite-plus/test";
+
+import type { NativeReviewDiffRow } from "./nativeReviewDiffSurface";
+import type { NativeReviewDiffFile } from "./nativeReviewDiffTypes";
+import { highlightNativeReviewDiffVisibleRows } from "./nativeReviewDiffHighlighter";
+
+const TYPESCRIPT_FILE: NativeReviewDiffFile = {
+  id: "file-1",
+  path: "example.ts",
+  language: "typescript",
+  additions: 0,
+  deletions: 0,
+};
+
+function makeLine(
+  input: Pick<NativeReviewDiffRow, "id" | "content" | "change" | "oldLineNumber" | "newLineNumber">,
+): NativeReviewDiffRow {
+  return {
+    kind: "line",
+    fileId: TYPESCRIPT_FILE.id,
+    ...input,
+  };
+}
+
+function makeHunk(id: string): NativeReviewDiffRow {
+  return {
+    kind: "hunk",
+    id,
+    fileId: TYPESCRIPT_FILE.id,
+    text: "@@",
+  };
+}
+
+function highlight(
+  rows: ReadonlyArray<NativeReviewDiffRow>,
+  alreadyHighlightedRowIds?: ReadonlySet<string>,
+) {
+  return highlightNativeReviewDiffVisibleRows({
+    rows,
+    files: [TYPESCRIPT_FILE],
+    scheme: "dark",
+    engine: "javascript",
+    firstRowIndex: 0,
+    lastRowIndex: rows.length - 1,
+    overscanRows: 0,
+    maxRows: 100,
+    alreadyHighlightedRowIds,
+  });
+}
+
+describe("highlightNativeReviewDiffVisibleRows", () => {
+  it("does not carry grammar state across hunk boundaries", async () => {
+    const exportRow = makeLine({
+      id: "export-row",
+      content: "export async function run() {}",
+      change: "add",
+      oldLineNumber: null,
+      newLineNumber: 100,
+    });
+    const rows = [
+      makeHunk("hunk-1"),
+      makeLine({
+        id: "import-open",
+        content: "import {",
+        change: "context",
+        oldLineNumber: 1,
+        newLineNumber: 1,
+      }),
+      makeLine({
+        id: "import-entry",
+        content: "  Model,",
+        change: "context",
+        oldLineNumber: 2,
+        newLineNumber: 2,
+      }),
+      makeHunk("hunk-2"),
+      exportRow,
+    ];
+
+    const [highlighted, standalone] = await Promise.all([
+      highlight(rows),
+      highlight([makeHunk("standalone-hunk"), exportRow]),
+    ]);
+
+    expect(highlighted.tokensByRowId[exportRow.id]).toEqual(standalone.tokensByRowId[exportRow.id]);
+  });
+
+  it("keeps grammar state across inline comment rows", async () => {
+    const openingRow = makeLine({
+      id: "template-open",
+      content: "const message = `open",
+      change: "add",
+      oldLineNumber: null,
+      newLineNumber: 1,
+    });
+    const closingRow = makeLine({
+      id: "template-close",
+      content: "closed`;",
+      change: "add",
+      oldLineNumber: null,
+      newLineNumber: 2,
+    });
+    const trailingRow = makeLine({
+      id: "trailing-row",
+      content: "export const answer = 42;",
+      change: "add",
+      oldLineNumber: null,
+      newLineNumber: 3,
+    });
+    const commentRow: NativeReviewDiffRow = {
+      kind: "comment",
+      id: "comment-1",
+      fileId: TYPESCRIPT_FILE.id,
+      commentText: "Review note",
+    };
+
+    const [withComment, contiguous] = await Promise.all([
+      highlight([openingRow, commentRow, closingRow, trailingRow]),
+      highlight([openingRow, closingRow, trailingRow]),
+    ]);
+
+    expect(withComment.tokensByRowId).toEqual(contiguous.tokensByRowId);
+  });
+
+  it("does not join unhighlighted rows across cached gaps", async () => {
+    const trailingRow = makeLine({
+      id: "trailing-row",
+      content: "export const answer = 42;",
+      change: "add",
+      oldLineNumber: null,
+      newLineNumber: 3,
+    });
+    const rows = [
+      makeLine({
+        id: "template-open",
+        content: "const message = `open",
+        change: "add",
+        oldLineNumber: null,
+        newLineNumber: 1,
+      }),
+      makeLine({
+        id: "template-close",
+        content: "closed`;",
+        change: "add",
+        oldLineNumber: null,
+        newLineNumber: 2,
+      }),
+      trailingRow,
+    ];
+
+    const [highlighted, standalone] = await Promise.all([
+      highlight(rows, new Set(["template-close"])),
+      highlight([trailingRow]),
+    ]);
+
+    expect(highlighted.tokensByRowId[trailingRow.id]).toEqual(
+      standalone.tokensByRowId[trailingRow.id],
+    );
+  });
+
+  it("keeps deletion grammar state out of addition rows", async () => {
+    const additionRow = makeLine({
+      id: "addition-row",
+      content: "export const answer = 42;",
+      change: "add",
+      oldLineNumber: null,
+      newLineNumber: 1,
+    });
+    const rows = [
+      makeLine({
+        id: "deletion-row",
+        content: "const removed = `open",
+        change: "delete",
+        oldLineNumber: 1,
+        newLineNumber: null,
+      }),
+      additionRow,
+    ];
+
+    const [highlighted, standalone] = await Promise.all([
+      highlight(rows),
+      highlight([additionRow]),
+    ]);
+
+    expect(highlighted.tokensByRowId[additionRow.id]).toEqual(
+      standalone.tokensByRowId[additionRow.id],
+    );
+  });
+});

apps/mobile/src/features/diffs/nativeReviewDiffHighlighter.ts

@@ -56,6 +56,11 @@ interface NativeReviewDiffLineRow extends NativeReviewDiffRow {
   readonly content: string;
 }
 
+interface IndexedNativeReviewDiffLineRow {
+  readonly row: NativeReviewDiffLineRow;
+  readonly rowIndex: number;
+}
+
 export interface NativeReviewDiffTokenChunk {
   readonly chunkIndex: number;
   readonly fileId: string;
@@ -308,6 +313,58 @@ function isHighlightableLineRow(row: NativeReviewDiffRow): row is NativeReviewDi
   return row.kind === "line" && typeof row.fileId === "string" && typeof row.content === "string";
 }
 
+function hasConsecutiveLineNumbers(
+  previous: number | null | undefined,
+  next: number | null | undefined,
+): boolean {
+  return typeof previous === "number" && typeof next === "number" && next === previous + 1;
+}
+
+function hasOnlyCommentRowsBetween(
+  rows: ReadonlyArray<NativeReviewDiffRow>,
+  previousRowIndex: number,
+  nextRowIndex: number,
+): boolean {
+  for (let rowIndex = previousRowIndex + 1; rowIndex < nextRowIndex; rowIndex += 1) {
+    if (rows[rowIndex]?.kind !== "comment") {
+      return false;
+    }
+  }
+  return true;
+}
+
+function canShareGrammarContext(
+  previous: IndexedNativeReviewDiffLineRow,
+  next: IndexedNativeReviewDiffLineRow,
+  rows: ReadonlyArray<NativeReviewDiffRow>,
+): boolean {
+  if (
+    next.row.fileId !== previous.row.fileId ||
+    !hasOnlyCommentRowsBetween(rows, previous.rowIndex, next.rowIndex)
+  ) {
+    return false;
+  }
+
+  if (previous.row.change === "delete" || next.row.change === "delete") {
+    return (
+      previous.row.change !== "add" &&
+      next.row.change !== "add" &&
+      hasConsecutiveLineNumbers(previous.row.oldLineNumber, next.row.oldLineNumber)
+    );
+  }
+
+  if (previous.row.change === "add" || next.row.change === "add") {
+    return hasConsecutiveLineNumbers(previous.row.newLineNumber, next.row.newLineNumber);
+  }
+
+  return (
+    previous.row.change === "context" &&
+    next.row.change === "context" &&
+    hasConsecutiveLineNumbers(previous.row.oldLineNumber, next.row.oldLineNumber) &&
+    hasConsecutiveLineNumbers(previous.row.newLineNumber, next.row.newLineNumber)
+  );
+}
+
 function groupLineRowsByFileId(rows: ReadonlyArray<NativeReviewDiffRow>) {
   const rowsByFileId = new Map<string, NativeReviewDiffLineRow[]>();
   for (const row of rows) {
@@ -360,7 +417,7 @@ export async function highlightNativeReviewDiffVisibleRows(
   const maxRows = input.maxRows ?? NATIVE_REVIEW_DIFF_VISIBLE_MAX_ROWS;
   const startIndex = clampRowIndex(input.firstRowIndex - overscanRows, input.rows);
   const endIndex = clampRowIndex(input.lastRowIndex + overscanRows, input.rows);
-  const selectedRows: NativeReviewDiffLineRow[] = [];
+  const selectedRows: IndexedNativeReviewDiffLineRow[] = [];
 
   for (
     let rowIndex = startIndex;
@@ -374,12 +431,12 @@ export async function highlightNativeReviewDiffVisibleRows(
       !input.alreadyHighlightedRowIds?.has(row.id) &&
       fileMap.has(row.fileId)
     ) {
-      selectedRows.push(row);
+      selectedRows.push({ row, rowIndex });
     }
   }
 
   const tokensByRowId: Record<string, ReadonlyArray<NativeReviewDiffToken>> = {};
-  let segmentRows: NativeReviewDiffLineRow[] = [];
+  let segmentRows: IndexedNativeReviewDiffLineRow[] = [];
   let segmentFile: NativeReviewDiffFile | undefined;
 
   const flushSegment = () => {
@@ -389,27 +446,34 @@ export async function highlightNativeReviewDiffVisibleRows(
       return;
     }
 
-    const code = segmentRows.map((row) => row.content).join("\n");
+    const code = segmentRows.map(({ row }) => row.content).join("\n");
     const tokenLines = highlighter.tokenize(code, { lang: segmentFile.language, theme });
-    segmentRows.forEach((row, rowIndex) => {
+    segmentRows.forEach(({ row }, rowIndex) => {
       tokensByRowId[row.id] = tokenLines[rowIndex] ?? makePlainTokenFallback(row);
     });
     segmentRows = [];
     segmentFile = undefined;
   };
 
-  for (const row of selectedRows) {
+  for (const selectedRow of selectedRows) {
+    const { row } = selectedRow;
     const file = fileMap.get(row.fileId);
     if (!file) {
       continue;
     }
 
-    if (segmentFile && segmentFile.id !== file.id) {
+    const previousRow = segmentRows.at(-1);
+    if (
+      segmentFile &&
+      (segmentFile.id !== file.id ||
+        (previousRow !== undefined &&
+          !canShareGrammarContext(previousRow, selectedRow, input.rows)))
+    ) {
       flushSegment();
     }
 
     segmentFile = file;
-    segmentRows.push(row);
+    segmentRows.push(selectedRow);
   }
   flushSegment();