fix(diff): highlight expanded context rows

Author: adit-chandra

src/ui/components/ui-components.test.tsx

@@ -2219,6 +2219,61 @@ describe("UI components", () => {
     }
   });
 
+  test("PierreDiffView highlights expanded unchanged source rows", async () => {
+    const beforeLines = Array.from({ length: 30 }, (_, index) =>
+      index === 0
+        ? "export const expandedMarker = 1;"
+        : `export const line${index + 1} = ${index + 1};`,
+    );
+    const afterLines = [...beforeLines];
+    afterLines[4] = "export const line5 = 999;";
+    const after = lines(...afterLines);
+    const file = buildTestDiffFile({
+      after,
+      before: lines(...beforeLines),
+      context: 3,
+      id: "expanded-highlight",
+      path: "expanded-highlight.ts",
+    });
+    const theme = resolveTheme("midnight", null);
+    const setup = await testRender(
+      <PierreDiffView
+        file={file}
+        layout="split"
+        theme={theme}
+        width={140}
+        selectedHunkIndex={0}
+        expandedGapKeys={new Set(["before:0"])}
+        sourceStatus={{ kind: "loaded", text: after }}
+        scrollable={false}
+      />,
+      { width: 144, height: 20 },
+    );
+
+    try {
+      let highlighted = false;
+      for (let iteration = 0; iteration < 400; iteration += 1) {
+        await act(async () => {
+          await setup.renderOnce();
+          await Bun.sleep(0);
+          await setup.renderOnce();
+          await Bun.sleep(0);
+        });
+
+        if (frameHasHighlightedMarker(setup.captureSpans(), "expandedMarker")) {
+          highlighted = true;
+          break;
+        }
+      }
+
+      expect(highlighted).toBe(true);
+    } finally {
+      await act(async () => {
+        setup.renderer.destroy();
+      });
+    }
+  });
+
   test("PierreDiffView renders word-diff spans with a visibly different background in split view", async () => {
     const file = createTestDiffFile(
       "word-diff",

src/ui/diff/PierreDiffView.tsx

@@ -1,4 +1,4 @@
-import { useMemo } from "react";
+import { useCallback, useMemo } from "react";
 import type { DiffFile, LayoutMode } from "../../core/types";
 import { AgentInlineNote, AgentInlineNoteGuideCap } from "../components/panes/AgentInlineNote";
 import type { VisibleAgentNote } from "../lib/agentAnnotations";
@@ -7,12 +7,13 @@ import { reviewRowId } from "../lib/ids";
 import type { AppTheme } from "../themes";
 import { findMaxLineNumber, findMaxLineNumberInRows } from "./codeColumns";
 import { expandCollapsedRows, type FileSourceStatus } from "./expandCollapsedRows";
-import { buildSplitRows, buildStackRows } from "./pierre";
+import { buildSplitRows, buildStackRows, spansForHighlightedSourceLine } from "./pierre";
 import { plannedReviewRowVisible } from "./plannedReviewRows";
 import { buildReviewRenderPlan } from "./reviewRenderPlan";
 import { resolveVisiblePlannedRowWindow, type VisibleBodyBounds } from "./rowWindowing";
 import { diffMessage, DiffRowView, fitText } from "./renderRows";
 import { useHighlightedDiff } from "./useHighlightedDiff";
+import { useHighlightedSource } from "./useHighlightedSource";
 
 const EMPTY_ANNOTATED_HUNK_INDICES = new Set<number>();
 const EMPTY_VISIBLE_AGENT_NOTES: VisibleAgentNote[] = [];
@@ -65,6 +66,23 @@ export function PierreDiffView({
     appearance: theme.appearance,
     shouldLoadHighlight,
   });
+  const sourceTextForHighlight =
+    sourceStatus?.kind === "loaded" && expandedGapKeys.size > 0 ? sourceStatus.text : undefined;
+  const resolvedHighlightedSource = useHighlightedSource({
+    file,
+    text: sourceTextForHighlight,
+    appearance: theme.appearance,
+    shouldLoadHighlight: shouldLoadHighlight && expandedGapKeys.size > 0,
+  });
+  const sourceLineSpans = useCallback(
+    (line: string | undefined, sourceLineNumber: number) =>
+      spansForHighlightedSourceLine(
+        line,
+        resolvedHighlightedSource?.lines[sourceLineNumber],
+        theme,
+      ),
+    [resolvedHighlightedSource, theme],
+  );
 
   const rows = useMemo(
     () =>
@@ -86,10 +104,11 @@ export function PierreDiffView({
       expandCollapsedRows(rows, {
         layout,
         expandedKeys: expandedGapKeys,
+        sourceLineSpans,
         sourceStatus,
         side: expansionSide,
       }),
-    [rows, layout, expandedGapKeys, sourceStatus, expansionSide],
+    [rows, layout, expandedGapKeys, sourceLineSpans, sourceStatus, expansionSide],
   );
   const plannedRows = useMemo(
     () =>

src/ui/diff/expandCollapsedRows.test.ts

@@ -252,6 +252,33 @@ describe("expandCollapsedRows", () => {
     }
     expect(inserted.cell.spans[0]?.text.includes("\t")).toBe(false);
   });
+
+  test("uses caller-provided spans for expanded source lines", () => {
+    const rows: DiffRow[] = [makeCollapsedRow("before", 0, [2, 3], [2, 3]), makeHunkHeader(0)];
+    const calls: Array<{ line: string | undefined; sourceLineNumber: number }> = [];
+
+    const result = expandCollapsedRows(rows, {
+      layout: "stack",
+      expandedKeys: new Set([gapKey("before", 0)]),
+      sourceStatus: { kind: "loaded", text: SOURCE },
+      sourceLineSpans: (line, sourceLineNumber) => {
+        calls.push({ line, sourceLineNumber });
+        return [{ text: `highlighted:${line ?? ""}`, fg: "#abcdef" }];
+      },
+      side: "new",
+    });
+
+    expect(calls).toEqual([
+      { line: "beta", sourceLineNumber: 1 },
+      { line: "gamma", sourceLineNumber: 2 },
+    ]);
+
+    const inserted = result[1];
+    if (!inserted || inserted.type !== "stack-line") {
+      throw new Error("expected stack-line context row");
+    }
+    expect(inserted.cell.spans).toEqual([{ text: "highlighted:beta", fg: "#abcdef" }]);
+  });
 });
 
 describe("selectGapForKeyboardToggle", () => {

src/ui/diff/expandCollapsedRows.ts

@@ -19,6 +19,8 @@ export interface ExpandCollapsedRowsOptions {
   layout: ExpansionLayout;
   expandedKeys: ReadonlySet<string>;
   sourceStatus: FileSourceStatus | undefined;
+  /** Optional syntax-aware span resolver for a zero-based source line. */
+  sourceLineSpans?: (line: string | undefined, sourceLineNumber: number) => RenderSpan[];
   // Whose side's line indices in the source text. Defaults to "new".
   // For deleted files (no new side) callers should pass "old" instead.
   side?: "old" | "new";
@@ -89,9 +91,8 @@ function buildSplitContextRow(
   index: number,
   oldLineNumber: number,
   newLineNumber: number,
-  line: string | undefined,
+  spans: RenderSpan[],
 ): Extract<DiffRow, { type: "split-line" }> {
-  const spans = spansFor(line);
   const cell = (lineNumber: number): SplitLineCell => ({
     kind: "context",
     sign: " ",
@@ -117,14 +118,14 @@ function buildStackContextRow(
   index: number,
   oldLineNumber: number,
   newLineNumber: number,
-  line: string | undefined,
+  spans: RenderSpan[],
 ): Extract<DiffRow, { type: "stack-line" }> {
   const cell: StackLineCell = {
     kind: "context",
     sign: " ",
     oldLineNumber,
     newLineNumber,
-    spans: spansFor(line),
+    spans,
   };
 
   return {
@@ -148,7 +149,7 @@ export function expandCollapsedRows(
   rows: DiffRow[],
   options: ExpandCollapsedRowsOptions,
 ): DiffRow[] {
-  const { layout, expandedKeys, sourceStatus, side = "new" } = options;
+  const { layout, expandedKeys, sourceLineSpans, sourceStatus, side = "new" } = options;
 
   if (expandedKeys.size === 0) {
     return rows;
@@ -199,6 +200,7 @@ export function expandCollapsedRows(
       const newLineNumber = row.newRange[0] + offset;
       const sourceLineNumber = (side === "old" ? oldLineNumber : newLineNumber) - 1;
       const text = sourceLines[sourceLineNumber];
+      const spans = sourceLineSpans?.(text, sourceLineNumber) ?? spansFor(text);
 
       result.push(
         layout === "split"
@@ -209,7 +211,7 @@ export function expandCollapsedRows(
               offset,
               oldLineNumber,
               newLineNumber,
-              text,
+              spans,
             )
           : buildStackContextRow(
               row.fileId,
@@ -218,7 +220,7 @@ export function expandCollapsedRows(
               offset,
               oldLineNumber,
               newLineNumber,
-              text,
+              spans,
             ),
       );
     }

src/ui/diff/pierre.test.ts

@@ -1,7 +1,14 @@
 import { describe, expect, test } from "bun:test";
 import { parseDiffFromFile } from "@pierre/diffs";
 import type { DiffFile } from "../../core/types";
-import { buildSplitRows, buildStackRows, loadHighlightedDiff, type DiffRow } from "./pierre";
+import {
+  buildSplitRows,
+  buildStackRows,
+  loadHighlightedDiff,
+  loadHighlightedSourceLines,
+  spansForHighlightedSourceLine,
+  type DiffRow,
+} from "./pierre";
 import { resolveTheme } from "../themes";
 
 function createDiffFile(): DiffFile {
@@ -181,6 +188,27 @@ describe("Pierre diff rows", () => {
     }
   });
 
+  test("builds syntax spans for highlighted full-source lines", async () => {
+    const file = createDiffFile();
+    const theme = resolveTheme("midnight", null);
+    const text = "export const hiddenMarker = true;\n";
+    const highlighted = await loadHighlightedSourceLines({
+      file,
+      text,
+      appearance: theme.appearance,
+    });
+    const spans = spansForHighlightedSourceLine(
+      "export const hiddenMarker = true;",
+      highlighted.lines[0],
+      theme,
+    );
+
+    expect(spans.map((span) => span.text).join("")).toBe("export const hiddenMarker = true;");
+    expect(spans.some((span) => span.text.includes("export") && typeof span.fg === "string")).toBe(
+      true,
+    );
+  });
+
   test("remaps Pierre markdown reds and greens away from diff-semantic hues", async () => {
     const file = createMarkdownDiffFile();