wip(ui): diff virtualization (#12693)

3 files changed, 112 insertions, 51 deletions

diff --git a/packages/ui/src/components/code.tsx b/packages/ui/src/components/code.tsx
index 4e7c82d78..2fe0e0352 100644
--- a/packages/ui/src/components/code.tsx
+++ b/packages/ui/src/components/code.tsx
@@ -318,7 +318,7 @@ export function Code<T>(props: CodeProps<T>) {
     const needle = query.toLowerCase()
     const out: Range[] = []
 
-    const cols = Array.from(root.querySelectorAll("[data-column-content]")).filter(
+    const cols = Array.from(root.querySelectorAll("[data-content] [data-line], [data-column-content]")).filter(
       (node): node is HTMLElement => node instanceof HTMLElement,
     )
 
@@ -537,17 +537,28 @@ export function Code<T>(props: CodeProps<T>) {
       node.removeAttribute("data-comment-selected")
     }
 
+    const annotations = Array.from(root.querySelectorAll("[data-line-annotation]")).filter(
+      (node): node is HTMLElement => node instanceof HTMLElement,
+    )
+
     for (const range of ranges) {
       const start = Math.max(1, Math.min(range.start, range.end))
       const end = Math.max(range.start, range.end)
 
       for (let line = start; line <= end; line++) {
-        const nodes = Array.from(root.querySelectorAll(`[data-line="${line}"]`))
+        const nodes = Array.from(root.querySelectorAll(`[data-line="${line}"], [data-column-number="${line}"]`))
         for (const node of nodes) {
           if (!(node instanceof HTMLElement)) continue
           node.setAttribute("data-comment-selected", "")
         }
       }
+
+      for (const annotation of annotations) {
+        const line = parseInt(annotation.dataset.lineAnnotation?.split(",")[1] ?? "", 10)
+        if (Number.isNaN(line)) continue
+        if (line < start || line > end) continue
+        annotation.setAttribute("data-comment-selected", "")
+      }
     }
   }
 
diff --git a/packages/ui/src/components/diff-ssr.tsx b/packages/ui/src/components/diff-ssr.tsx
index 602e59a2f..e739afc16 100644
--- a/packages/ui/src/components/diff-ssr.tsx
+++ b/packages/ui/src/components/diff-ssr.tsx
@@ -1,8 +1,9 @@
-import { DIFFS_TAG_NAME, FileDiff, type SelectedLineRange } from "@pierre/diffs"
+import { DIFFS_TAG_NAME, FileDiff, type SelectedLineRange, VirtualizedFileDiff } from "@pierre/diffs"
 import { PreloadMultiFileDiffResult } from "@pierre/diffs/ssr"
 import { createEffect, onCleanup, onMount, Show, splitProps } from "solid-js"
 import { Dynamic, isServer } from "solid-js/web"
 import { createDefaultOptions, styleVariables, type DiffProps } from "../pierre"
+import { acquireVirtualizer, virtualMetrics } from "../pierre/virtualizer"
 import { useWorkerPool } from "../context/worker-pool"
 
 export type SSRDiffProps<T = {}> = DiffProps<T> & {
@@ -24,10 +25,21 @@ export function Diff<T>(props: SSRDiffProps<T>) {
   const workerPool = useWorkerPool(props.diffStyle)
 
   let fileDiffInstance: FileDiff<T> | undefined
+  let sharedVirtualizer: NonNullable<ReturnType<typeof acquireVirtualizer>> | undefined
   const cleanupFunctions: Array<() => void> = []
 
   const getRoot = () => fileDiffRef?.shadowRoot ?? undefined
 
+  const getVirtualizer = () => {
+    if (sharedVirtualizer) return sharedVirtualizer.virtualizer
+
+    const result = acquireVirtualizer(container)
+    if (!result) return
+
+    sharedVirtualizer = result
+    return result.virtualizer
+  }
+
   const applyScheme = () => {
     const scheme = document.documentElement.dataset.colorScheme
     if (scheme === "dark" || scheme === "light") {
@@ -70,10 +82,10 @@ export function Diff<T>(props: SSRDiffProps<T>) {
     const root = getRoot()
     if (!root) return
 
-    const diffs = root.querySelector("[data-diffs]")
+    const diffs = root.querySelector("[data-diff]")
     if (!(diffs instanceof HTMLElement)) return
 
-    const split = diffs.dataset.type === "split"
+    const split = diffs.dataset.diffType === "split"
 
     const start = rowIndex(root, split, range.start, range.side)
     const end = rowIndex(root, split, range.end, range.endSide ?? range.side)
@@ -132,15 +144,19 @@ export function Diff<T>(props: SSRDiffProps<T>) {
       node.removeAttribute("data-comment-selected")
     }
 
-    const diffs = root.querySelector("[data-diffs]")
+    const diffs = root.querySelector("[data-diff]")
     if (!(diffs instanceof HTMLElement)) return
 
-    const split = diffs.dataset.type === "split"
+    const split = diffs.dataset.diffType === "split"
+
+    const rows = Array.from(diffs.querySelectorAll("[data-line-index]")).filter(
+      (node): node is HTMLElement => node instanceof HTMLElement,
+    )
+    if (rows.length === 0) return
 
-    const code = Array.from(diffs.querySelectorAll("[data-code]")).filter(
+    const annotations = Array.from(diffs.querySelectorAll("[data-line-annotation]")).filter(
       (node): node is HTMLElement => node instanceof HTMLElement,
     )
-    if (code.length === 0) return
 
     const lineIndex = (element: HTMLElement) => {
       const raw = element.dataset.lineIndex
@@ -183,19 +199,18 @@ export function Diff<T>(props: SSRDiffProps<T>) {
       const first = Math.min(start, end)
       const last = Math.max(start, end)
 
-      for (const block of code) {
-        for (const element of Array.from(block.children)) {
-          if (!(element instanceof HTMLElement)) continue
-          const idx = lineIndex(element)
-          if (idx === undefined) continue
-          if (idx > last) break
-          if (idx < first) continue
-          element.setAttribute("data-comment-selected", "")
-          const next = element.nextSibling
-          if (next instanceof HTMLElement && next.hasAttribute("data-line-annotation")) {
-            next.setAttribute("data-comment-selected", "")
-          }
-        }
+      for (const row of rows) {
+        const idx = lineIndex(row)
+        if (idx === undefined) continue
+        if (idx < first || idx > last) continue
+        row.setAttribute("data-comment-selected", "")
+      }
+
+      for (const annotation of annotations) {
+        const idx = parseInt(annotation.dataset.lineAnnotation?.split(",")[1] ?? "", 10)
+        if (Number.isNaN(idx)) continue
+        if (idx < first || idx > last) continue
+        annotation.setAttribute("data-comment-selected", "")
       }
     }
   }
@@ -212,14 +227,27 @@ export function Diff<T>(props: SSRDiffProps<T>) {
       onCleanup(() => monitor.disconnect())
     }
 
-    fileDiffInstance = new FileDiff<T>(
-      {
-        ...createDefaultOptions(props.diffStyle),
-        ...others,
-        ...props.preloadedDiff,
-      },
-      workerPool,
-    )
+    const virtualizer = getVirtualizer()
+
+    fileDiffInstance = virtualizer
+      ? new VirtualizedFileDiff<T>(
+          {
+            ...createDefaultOptions(props.diffStyle),
+            ...others,
+            ...props.preloadedDiff,
+          },
+          virtualizer,
+          virtualMetrics,
+          workerPool,
+        )
+      : new FileDiff<T>(
+          {
+            ...createDefaultOptions(props.diffStyle),
+            ...others,
+            ...props.preloadedDiff,
+          },
+          workerPool,
+        )
     // @ts-expect-error - fileContainer is private but needed for SSR hydration
     fileDiffInstance.fileContainer = fileDiffRef
     fileDiffInstance.hydrate({
@@ -273,6 +301,8 @@ export function Diff<T>(props: SSRDiffProps<T>) {
     // Clean up FileDiff event handlers and dispose SolidJS components
     fileDiffInstance?.cleanUp()
     cleanupFunctions.forEach((dispose) => dispose())
+    sharedVirtualizer?.release()
+    sharedVirtualizer = undefined
   })
 
   return (
diff --git a/packages/ui/src/components/diff.tsx b/packages/ui/src/components/diff.tsx
index 21dada535..0966db75e 100644
--- a/packages/ui/src/components/diff.tsx
+++ b/packages/ui/src/components/diff.tsx
@@ -1,8 +1,9 @@
 import { checksum } from "@opencode-ai/util/encode"
-import { FileDiff, type SelectedLineRange } from "@pierre/diffs"
+import { FileDiff, type SelectedLineRange, VirtualizedFileDiff } from "@pierre/diffs"
 import { createMediaQuery } from "@solid-primitives/media"
 import { createEffect, createMemo, createSignal, onCleanup, splitProps } from "solid-js"
 import { createDefaultOptions, type DiffProps, styleVariables } from "../pierre"
+import { acquireVirtualizer, virtualMetrics } from "../pierre/virtualizer"
 import { getWorkerPool } from "../pierre/worker"
 
 type SelectionSide = "additions" | "deletions"
@@ -52,6 +53,7 @@ function findSide(node: Node | null): SelectionSide | undefined {
 export function Diff<T>(props: DiffProps<T>) {
   let container!: HTMLDivElement
   let observer: MutationObserver | undefined
+  let sharedVirtualizer: NonNullable<ReturnType<typeof acquireVirtualizer>> | undefined
   let renderToken = 0
   let selectionFrame: number | undefined
   let dragFrame: number | undefined
@@ -92,6 +94,16 @@ export function Diff<T>(props: DiffProps<T>) {
   const [current, setCurrent] = createSignal<FileDiff<T> | undefined>(undefined)
   const [rendered, setRendered] = createSignal(0)
 
+  const getVirtualizer = () => {
+    if (sharedVirtualizer) return sharedVirtualizer.virtualizer
+
+    const result = acquireVirtualizer(container)
+    if (!result) return
+
+    sharedVirtualizer = result
+    return result.virtualizer
+  }
+
   const getRoot = () => {
     const host = container.querySelector("diffs-container")
     if (!(host instanceof HTMLElement)) return
@@ -147,10 +159,10 @@ export function Diff<T>(props: DiffProps<T>) {
     const root = getRoot()
     if (!root) return
 
-    const diffs = root.querySelector("[data-diffs]")
+    const diffs = root.querySelector("[data-diff]")
     if (!(diffs instanceof HTMLElement)) return
 
-    const split = diffs.dataset.type === "split"
+    const split = diffs.dataset.diffType === "split"
 
     const start = rowIndex(root, split, range.start, range.side)
     const end = rowIndex(root, split, range.end, range.endSide ?? range.side)
@@ -261,15 +273,19 @@ export function Diff<T>(props: DiffProps<T>) {
       node.removeAttribute("data-comment-selected")
     }
 
-    const diffs = root.querySelector("[data-diffs]")
+    const diffs = root.querySelector("[data-diff]")
     if (!(diffs instanceof HTMLElement)) return
 
-    const split = diffs.dataset.type === "split"
+    const split = diffs.dataset.diffType === "split"
+
+    const rows = Array.from(diffs.querySelectorAll("[data-line-index]")).filter(
+      (node): node is HTMLElement => node instanceof HTMLElement,
+    )
+    if (rows.length === 0) return
 
-    const code = Array.from(diffs.querySelectorAll("[data-code]")).filter(
+    const annotations = Array.from(diffs.querySelectorAll("[data-line-annotation]")).filter(
       (node): node is HTMLElement => node instanceof HTMLElement,
     )
-    if (code.length === 0) return
 
     for (const range of ranges) {
       const start = rowIndex(root, split, range.start, range.side)
@@ -285,19 +301,18 @@ export function Diff<T>(props: DiffProps<T>) {
       const first = Math.min(start, end)
       const last = Math.max(start, end)
 
-      for (const block of code) {
-        for (const element of Array.from(block.children)) {
-          if (!(element instanceof HTMLElement)) continue
-          const idx = lineIndex(split, element)
-          if (idx === undefined) continue
-          if (idx > last) break
-          if (idx < first) continue
-          element.setAttribute("data-comment-selected", "")
-          const next = element.nextSibling
-          if (next instanceof HTMLElement && next.hasAttribute("data-line-annotation")) {
-            next.setAttribute("data-comment-selected", "")
-          }
-        }
+      for (const row of rows) {
+        const idx = lineIndex(split, row)
+        if (idx === undefined) continue
+        if (idx < first || idx > last) continue
+        row.setAttribute("data-comment-selected", "")
+      }
+
+      for (const annotation of annotations) {
+        const idx = parseInt(annotation.dataset.lineAnnotation?.split(",")[1] ?? "", 10)
+        if (Number.isNaN(idx)) continue
+        if (idx < first || idx > last) continue
+        annotation.setAttribute("data-comment-selected", "")
       }
     }
   }
@@ -514,12 +529,15 @@ export function Diff<T>(props: DiffProps<T>) {
   createEffect(() => {
     const opts = options()
     const workerPool = getWorkerPool(props.diffStyle)
+    const virtualizer = getVirtualizer()
     const annotations = local.annotations
     const beforeContents = typeof local.before?.contents === "string" ? local.before.contents : ""
     const afterContents = typeof local.after?.contents === "string" ? local.after.contents : ""
 
     instance?.cleanUp()
-    instance = new FileDiff<T>(opts, workerPool)
+    instance = virtualizer
+      ? new VirtualizedFileDiff<T>(opts, virtualizer, virtualMetrics, workerPool)
+      : new FileDiff<T>(opts, workerPool)
     setCurrent(instance)
 
     container.innerHTML = ""
@@ -606,6 +624,8 @@ export function Diff<T>(props: DiffProps<T>) {
 
     instance?.cleanUp()
     setCurrent(undefined)
+    sharedVirtualizer?.release()
+    sharedVirtualizer = undefined
   })
 
   return <div data-component="diff" style={styleVariables} ref={container} />