fix(tui): virtualize large added-file diffs (#51122)

Co-authored-by: vimtor <36263538+vimtor@users.noreply.github.com>
Co-authored-by: vimtor <vn4varro@gmail.com>

packages/tui/src/component/patch-diff.tsx

@@ -1,31 +1,61 @@
 /** @jsxImportSource @opentui/solid */
-import { DiffRenderable, LineNumberRenderable, type ColorInput } from "@opentui/core"
+import {
+  BoxRenderable,
+  CodeRenderable,
+  DiffRenderable,
+  getTreeSitterClient,
+  LineNumberRenderable,
+  type ColorInput,
+  type OnHighlightCallback,
+  type Renderable,
+  type ScrollBoxRenderable,
+  type SimpleHighlight,
+} from "@opentui/core"
 import type { JSX } from "@opentui/solid"
-import { createMemo, For, Show, splitProps } from "solid-js"
-import { splitPatchHunks } from "../util/diff"
+import { useRenderer } from "@opentui/solid"
+import { batch, createMemo, createSignal, For, onCleanup, Show, splitProps } from "solid-js"
+import { splitAddedPatch, splitPatchHunks, type AddedPatchChunk } from "../util/diff"
 import { stringWidth } from "../util/string-width"
 
 export interface PatchDiffRef {
-  readonly hunks: () => readonly DiffRenderable[]
+  readonly hunks: () => readonly (DiffRenderable | BoxRenderable)[]
 }
 
+// Smaller patches render fine as a single DiffRenderable; only split files large enough to stall the TUI.
+const VIRTUAL_MIN_LINES = 3000
+const VIRTUAL_CHUNK_LINES = 128
+
 type Props = Omit<JSX.IntrinsicElements["diff"], "diff" | "lineNumberBg" | "ref"> & {
   diff: string
   hunkFg: ColorInput
   lineNumberBg: ColorInput
   ref?: (value: PatchDiffRef) => void
+  scroll?: () => ScrollBoxRenderable | undefined
 }
 
 export function PatchDiff(props: Props) {
-  const [local, diffProps] = splitProps(props, ["diff", "hunkFg", "lineNumberBg", "ref"])
+  const [local, diffProps] = splitProps(props, ["diff", "hunkFg", "lineNumberBg", "ref", "scroll"])
   const hunks = createMemo(() => splitPatchHunks(local.diff))
+  const chunks = createMemo(() => {
+    if (!local.scroll) return
+    const result = splitAddedPatch(local.diff, VIRTUAL_CHUNK_LINES)
+    return result && lineCount(result) > VIRTUAL_MIN_LINES ? result : undefined
+  })
+  // Virtual chunks mount independently, so size the gutter for the whole file rather than the mounted chunks.
+  const minDigits = createMemo(() => {
+    const items = chunks()
+    return items ? String(lineCount(items)).length : 0
+  })
   const nodes = new Map<number, DiffRenderable>()
+  let virtualRoot: BoxRenderable | undefined
   local.ref?.({
-    hunks: () =>
-      [...nodes.entries()]
+    hunks: () => {
+      if (chunks()) return virtualRoot && !virtualRoot.isDestroyed ? [virtualRoot] : []
+      return [...nodes.entries()]
         .sort(([left], [right]) => left - right)
         .map(([, node]) => node)
-        .filter((node) => !node.isDestroyed),
+        .filter((node) => !node.isDestroyed)
+    },
   })
   const syncGutters = (attempt = 0) => {
     requestAnimationFrame(() => {
@@ -44,10 +74,11 @@ export function PatchDiff(props: Props) {
       const maxAfter = Math.max(...after)
       if (!maxDigits && attempt < 2) return syncGutters(attempt + 1)
       if (!maxDigits) return
+      const width = Math.max(maxDigits, minDigits())
       sides.forEach((side) => {
         const index = sides.indexOf(side)
         const signs = new Map([...side.getLineSigns()].filter(([line]) => line >= 0))
-        signs.set(-1, { after: " ".repeat(maxAfter + maxDigits - digits[index]) })
+        signs.set(-1, { after: " ".repeat(maxAfter + width - digits[index]) })
         side.setLineNumbers(lineNumbers[index])
         side.setLineSigns(signs)
       })
@@ -55,29 +86,143 @@ export function PatchDiff(props: Props) {
   }
   const register = (index: number, node: DiffRenderable) => {
     nodes.set(index, node)
+    onCleanup(() => nodes.delete(index))
     syncGutters()
   }
 
   return (
-    <For each={hunks()}>
-      {(hunk, index) => (
-        <>
-          <Show when={index() > 0}>
-            <box width="100%" height={1} backgroundColor={local.lineNumberBg}>
-              <text fg={local.hunkFg} bg={local.lineNumberBg}>
-                {` ${hunk.header ?? ""}`}
-              </text>
-            </box>
-          </Show>
-          <diff
-            {...diffProps}
-            ref={(node: DiffRenderable) => register(index(), node)}
-            diff={hunk.patch}
-            minHeight={hunk.rows}
-            lineNumberBg={local.lineNumberBg}
-          />
-        </>
+    <Show
+      when={chunks()}
+      fallback={
+        <For each={hunks()}>
+          {(hunk, index) => (
+            <>
+              <Show when={index() > 0}>
+                <box width="100%" height={1} backgroundColor={local.lineNumberBg}>
+                  <text fg={local.hunkFg} bg={local.lineNumberBg}>
+                    {` ${hunk.header ?? ""}`}
+                  </text>
+                </box>
+              </Show>
+              <diff
+                {...diffProps}
+                ref={(node: DiffRenderable) => register(index(), node)}
+                diff={hunk.patch}
+                minHeight={hunk.rows}
+                lineNumberBg={local.lineNumberBg}
+              />
+            </>
+          )}
+        </For>
+      }
+    >
+      {(items) => (
+        <VirtualAddedPatch
+          chunks={items()}
+          scroll={local.scroll!}
+          diffProps={diffProps}
+          lineNumberBg={local.lineNumberBg}
+          register={register}
+          registerRoot={(root) => (virtualRoot = root)}
+        />
       )}
-    </For>
+    </Show>
   )
 }
+
+// Chunks render without wrapping so each one is exactly `rows` tall. Offscreen chunks become fixed-height
+// placeholders, and the chunks overlapping the viewport (plus one on each side) follow from the scroll offset.
+function VirtualAddedPatch(props: {
+  chunks: readonly AddedPatchChunk[]
+  scroll: () => ScrollBoxRenderable | undefined
+  diffProps: Omit<JSX.IntrinsicElements["diff"], "diff" | "lineNumberBg" | "ref">
+  lineNumberBg: ColorInput
+  register: (index: number, node: DiffRenderable) => void
+  registerRoot: (root: BoxRenderable) => void
+}) {
+  const renderer = useRenderer()
+  const [first, setFirst] = createSignal(0)
+  const [last, setLast] = createSignal(0)
+  // A chunk is not valid source on its own (a slice of a JSON object parses as an error), so highlight
+  // the whole file once and give each chunk its slice of the result.
+  const contents = createMemo(() => props.chunks.map((chunk) => chunk.lines.map((line) => line.slice(1)).join("\n")))
+  const offsets = createMemo(() =>
+    contents().map((_, index, all) => all.slice(0, index).reduce((sum, content) => sum + content.length + 1, 0)),
+  )
+  const fileHighlights = createMemo(() => {
+    const filetype = props.diffProps.filetype
+    if (!filetype) return
+    return (
+      getTreeSitterClient()
+        .highlightOnce(contents().join("\n"), filetype)
+        .then((result) => result.highlights)
+        // Rejects when the renderer tears down the client mid-parse; chunks then keep their own highlights.
+        .catch(() => undefined)
+    )
+  })
+  const chunkHighlights =
+    (index: number): OnHighlightCallback =>
+    async () => {
+      const all = await fileHighlights()
+      if (!all) return
+      const start = offsets()[index]
+      const end = start + contents()[index].length
+      return all.flatMap((highlight): SimpleHighlight[] =>
+        highlight[0] < end && highlight[1] > start
+          ? [[Math.max(highlight[0], start) - start, Math.min(highlight[1], end) - start, highlight[2], highlight[3]]]
+          : [],
+      )
+    }
+
+  return (
+    <box
+      width="100%"
+      ref={(root: BoxRenderable) => {
+        props.registerRoot(root)
+        root.onLifecyclePass = () => {
+          const scroll = props.scroll()
+          if (!scroll) return
+          // ScrollBox's scroll position is not a Solid signal; observe it during the render pass.
+          const top = scroll.scrollTop - (root.y - scroll.content.y)
+          batch(() => {
+            setFirst(Math.floor(top / VIRTUAL_CHUNK_LINES))
+            setLast(Math.floor((top + scroll.viewport.height) / VIRTUAL_CHUNK_LINES))
+          })
+        }
+        renderer.registerLifecyclePass(root)
+        onCleanup(() => renderer.unregisterLifecyclePass(root))
+      }}
+    >
+      <For each={props.chunks}>
+        {(chunk, index) => (
+          <Show when={index() >= first() - 1 && index() <= last() + 1} fallback={<box height={chunk.rows} />}>
+            <diff
+              {...props.diffProps}
+              ref={(node: DiffRenderable) => {
+                props.register(index(), node)
+                // DiffRenderable creates its CodeRenderable after ref runs; setting onHighlight re-highlights.
+                queueMicrotask(() => {
+                  const code = findCode(node)
+                  if (code) code.onHighlight = chunkHighlights(index())
+                })
+              }}
+              diff={chunk.patch}
+              wrapMode="none"
+              height={chunk.rows}
+              lineNumberBg={props.lineNumberBg}
+            />
+          </Show>
+        )}
+      </For>
+    </box>
+  )
+}
+
+function findCode(node: Renderable): CodeRenderable | undefined {
+  if (node instanceof CodeRenderable) return node
+  return node.getChildren().reduce<CodeRenderable | undefined>((found, child) => found ?? findCode(child), undefined)
+}
+
+function lineCount(chunks: readonly AddedPatchChunk[]) {
+  return chunks.reduce((count, chunk) => count + chunk.rows, 0)
+}

packages/tui/src/feature-plugins/system/diff-viewer.tsx

@@ -1024,6 +1024,7 @@ export function DiffViewerContent(props: {
                                         onCleanup(() => patchDiffByFileIndex.delete(entry.fileIndex))
                                       }}
                                       diff={patch()}
+                                      scroll={() => scroll}
                                       hunkFg={theme.diff.text.hunkHeader}
                                       view={entry.file.status === "modified" ? view() : "unified"}
                                       filetype={filetype(entry.file.file)}

packages/tui/src/util/diff.ts

@@ -4,6 +4,39 @@ export interface PatchHunk {
   readonly rows?: number
 }
 
+export interface AddedPatchChunk {
+  readonly patch: string
+  readonly lines: readonly string[]
+  readonly rows: number
+}
+
+/**
+ * Splits a new-file patch into chunks of `size` lines, each a valid patch with its own `@@ -0,0 +start,count @@`
+ * header. Returns undefined for anything else: patches with context or removed lines would need old and new line
+ * numbers recomputed at every cut, so they are not split.
+ */
+export function splitAddedPatch(patch: string, size: number): AddedPatchChunk[] | undefined {
+  const header = /^@@ -0,0 \+1,(\d+) @@[^\n]*\n/m.exec(patch)
+  if (!header) return
+  const count = Number(header[1])
+  const lines = patch
+    .slice(header.index + header[0].length)
+    .replace(/\n$/, "")
+    .split("\n")
+  const marker = lines.at(-1)?.startsWith("\\ No newline at end of file") ? lines.pop() : undefined
+  if (lines.length !== count || lines.some((line) => !line.startsWith("+"))) return
+  const prefix = patch.slice(0, header.index)
+  return Array.from({ length: Math.ceil(count / size) }, (_, index) => {
+    const start = index * size
+    const slice = lines.slice(start, start + size)
+    return {
+      patch: `${prefix}@@ -0,0 +${start + 1},${slice.length} @@\n${slice.join("\n")}${marker && start + size >= count ? `\n${marker}` : ""}`,
+      lines: slice,
+      rows: slice.length,
+    }
+  })
+}
+
 export function splitPatchHunks(patch: string): PatchHunk[] {
   const starts = [...patch.matchAll(/^@@ -\d+(?:,\d+)? \+\d+(?:,\d+)? @@.*$/gm)].map((match) => match.index)
   if (starts.length <= 1) return [{ patch }]

packages/tui/test/cli/tui/diff-viewer.test.tsx

@@ -2058,6 +2058,197 @@ const manyDiffs = Array.from({ length: 40 }, (_, index) => ({
   file: `file${String(index).padStart(2, "0")}.txt`,
 }))
 
+test.each([80, 160])("virtualizes a large added file at %i columns without losing its end", async (width) => {
+  const lines = [
+    "+{",
+    ...Array.from(
+      { length: 7500 },
+      (_, index) =>
+        `+  "row-${String(index).padStart(4, "0")}": "${"value".repeat(index === 777 ? 2000 : index % 7 === 0 ? 24 : 1)}"${index === 7499 ? "" : ","}`,
+    ),
+    "+}",
+  ]
+  const viewer = await renderDiffViewer(
+    [
+      {
+        file: "snapshot.json",
+        status: "added",
+        additions: lines.length,
+        deletions: 0,
+        patch: `diff --git a/snapshot.json b/snapshot.json\nnew file mode 100644\n--- /dev/null\n+++ b/snapshot.json\n@@ -0,0 +1,${lines.length} @@\n${lines.join("\n")}`,
+      },
+    ],
+    { width, height: 24 },
+  )
+  try {
+    expect(viewer.app.captureCharFrame()).toContain("row-0000")
+    expect(
+      findDiffs(viewer.app.renderer.root).reduce((total, node) => total + node.diff.split("\n").length, 0),
+    ).toBeLessThan(2000)
+    viewer.commands.get("diff.last")!.run()
+    await viewer.app.flush()
+    if (!viewer.app.captureCharFrame().includes("row-7499")) {
+      await viewer.app.waitForFrame((frame) => frame.includes("row-7499"))
+    }
+    expect(viewer.app.captureCharFrame()).toContain("row-7499")
+    expect(
+      findDiffs(viewer.app.renderer.root).reduce((total, node) => total + node.diff.split("\n").length, 0),
+    ).toBeLessThan(2000)
+    viewer.commands.get("diff.first")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("row-0000")
+    viewer.app.resize(width === 80 ? 160 : 80, 20)
+    await viewer.app.flush()
+    viewer.commands.get("diff.last")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("row-7499")
+  } finally {
+    viewer.app.renderer.destroy()
+  }
+})
+
+test("keeps the line-number gutter the same width across virtual chunks", async () => {
+  const additions = Array.from({ length: 10500 }, (_, index) => `+line-${String(index + 1).padStart(5, "0")}`)
+  const viewer = await renderDiffViewer(
+    [
+      {
+        file: "wide.txt",
+        status: "added",
+        additions: additions.length,
+        deletions: 0,
+        patch: `--- /dev/null\n+++ b/wide.txt\n@@ -0,0 +1,${additions.length} @@\n${additions.join("\n")}`,
+      },
+    ],
+    { width: 120, height: 24 },
+  )
+  const column = (text: string) =>
+    viewer.app
+      .captureCharFrame()
+      .split("\n")
+      .find((line) => line.includes(text))
+      ?.indexOf(text)
+  try {
+    await viewer.app.flush()
+    const top = column("line-00001")
+    viewer.commands.get("diff.last")!.run()
+    await viewer.app.flush()
+    if (!viewer.app.captureCharFrame().includes("line-10500")) {
+      await viewer.app.waitForFrame((frame) => frame.includes("line-10500"))
+    }
+    expect(top).toBeDefined()
+    expect(column("line-10500")).toBe(top)
+  } finally {
+    viewer.app.renderer.destroy()
+  }
+})
+
+test("highlights virtual chunks with whole-file syntax context", async () => {
+  const additions = Array.from({ length: 3200 }, (_, index) => {
+    if (index === 370) return "+/*"
+    if (index === 399) return "+*/"
+    if (index > 370 && index < 399) return `+  comment ${index}`
+    return `+const value${index} = ${index}`
+  })
+  const viewer = await renderDiffViewer(
+    [
+      {
+        file: "big.ts",
+        status: "added",
+        additions: additions.length,
+        deletions: 0,
+        patch: `--- /dev/null\n+++ b/big.ts\n@@ -0,0 +1,${additions.length} @@\n${additions.join("\n")}`,
+      },
+    ],
+    { width: 120, height: 40 },
+  )
+  const color = (text: string) =>
+    viewer.app
+      .captureSpans()
+      .lines.flatMap((line) => line.spans)
+      .find((span) => span.text.includes(text))?.fg
+  try {
+    findScrollBox(viewer.app.renderer.root)!.scrollTo(360)
+    await viewer.app.flush()
+    // The comment starts in the first chunk and ends in the second; wait until it is highlighted.
+    for (let attempt = 0; attempt < 100 && `${color("comment 375")}` === `${color("value365")}`; attempt++) {
+      await Bun.sleep(20)
+      await viewer.app.flush()
+    }
+    expect(`${color("comment 375")}`).not.toBe(`${color("value365")}`)
+    for (let attempt = 0; attempt < 100 && `${color("comment 390")}` !== `${color("comment 375")}`; attempt++) {
+      await Bun.sleep(20)
+      await viewer.app.flush()
+    }
+    expect(`${color("comment 390")}`).toBe(`${color("comment 375")}`)
+  } finally {
+    viewer.app.renderer.destroy()
+  }
+})
+
+test("does not virtualize added files at or below the size threshold", async () => {
+  const additions = Array.from({ length: 3000 }, (_, index) => `+small line ${index}`)
+  const viewer = await renderDiffViewer(
+    [
+      {
+        file: "small.txt",
+        status: "added",
+        additions: additions.length,
+        deletions: 0,
+        patch: `--- /dev/null\n+++ b/small.txt\n@@ -0,0 +1,${additions.length} @@\n${additions.join("\n")}`,
+      },
+    ],
+    { width: 120, height: 24 },
+  )
+  try {
+    expect(findDiffs(viewer.app.renderer.root)).toHaveLength(1)
+  } finally {
+    viewer.app.renderer.destroy()
+  }
+})
+
+test("file navigation and review still work after a virtualized patch", async () => {
+  const additions = Array.from({ length: 6600 }, (_, index) => `+added line ${index}`)
+  const viewer = await renderDiffViewer(
+    [
+      {
+        file: "a-large.txt",
+        status: "added",
+        additions: additions.length,
+        deletions: 0,
+        patch: `--- /dev/null\n+++ b/a-large.txt\n@@ -0,0 +1,${additions.length} @@\n${additions.join("\n")}`,
+      },
+      { ...hunkDiff[0], file: "b-small.txt" },
+    ],
+    { width: 160, height: 24 },
+  )
+  try {
+    const scroll = findScrollBox(viewer.app.renderer.root)!
+    scroll.scrollTo(2700)
+    await viewer.app.flush()
+    viewer.commands.get("diff.previous_hunk")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("added line 0")
+    viewer.commands.get("diff.next_hunk")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("b-small.txt")
+    viewer.commands.get("diff.next_file")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("b-small.txt")
+    expect(viewer.app.captureCharFrame()).toContain("const first")
+    viewer.commands.get("diff.previous_file")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("a-large.txt")
+    viewer.commands.get("diff.mark_reviewed")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).not.toContain("added line 0")
+    viewer.commands.get("diff.mark_reviewed")!.run()
+    await viewer.app.flush()
+    expect(viewer.app.captureCharFrame()).toContain("added line 0")
+  } finally {
+    viewer.app.renderer.destroy()
+  }
+})
+
 function findScrollBox(root: Renderable, patches = true): ScrollBoxRenderable | undefined {
   const node = root.findDescendantById(patches ? "diff-patches" : "diff-files")
   return node instanceof ScrollBoxRenderable ? node : undefined

packages/tui/test/util/added-patch.test.ts

@@ -0,0 +1,44 @@
+import { expect, test } from "bun:test"
+import { splitAddedPatch } from "../../src/util/diff"
+
+test("splits a complete new-file patch into independently numbered chunks", () => {
+  const patch = `diff --git a/new.txt b/new.txt
+new file mode 100644
+--- /dev/null
++++ b/new.txt
+@@ -0,0 +1,5 @@
++one
++++value beginning with plus signs
++three
++four
++five`
+  const chunks = splitAddedPatch(patch, 2)!
+  expect(chunks.map((chunk) => chunk.rows)).toEqual([2, 2, 1])
+  expect(chunks.map((chunk) => chunk.patch.match(/@@ -0,0 \+(\d+),(\d+) @@/)?.slice(1))).toEqual([
+    ["1", "2"],
+    ["3", "2"],
+    ["5", "1"],
+  ])
+  expect(chunks.flatMap((chunk) => chunk.lines)).toEqual([
+    "+one",
+    "+++value beginning with plus signs",
+    "+three",
+    "+four",
+    "+five",
+  ])
+  expect(chunks.every((chunk) => chunk.patch.startsWith("diff --git a/new.txt b/new.txt"))).toBe(true)
+})
+
+test("retains a missing-final-newline marker only on the last chunk", () => {
+  const patch = `--- /dev/null\n+++ b/new.txt\n@@ -0,0 +1,3 @@\n+one\n+two\n+three\n\\ No newline at end of file\n`
+  const chunks = splitAddedPatch(patch, 2)!
+  expect(chunks).toHaveLength(2)
+  expect(chunks[0].patch).not.toContain("No newline")
+  expect(chunks[1].patch).toContain("+three\n\\ No newline at end of file")
+})
+
+test("does not split partial or mixed patches", () => {
+  expect(splitAddedPatch("@@ -1 +1 @@\n-before\n+after", 2)).toBeUndefined()
+  expect(splitAddedPatch("@@ -0,0 +1,3 @@\n+one\n+two", 2)).toBeUndefined()
+  expect(splitAddedPatch("@@ -0,0 +1,2 @@\n+one\n two", 2)).toBeUndefined()
+})