1 files changed, 58 insertions, 28 deletions

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 (