wip(ui): diff virtualization (#12693)

3 files changed, 99 insertions, 48 deletions

diff --git a/packages/ui/src/pierre/index.ts b/packages/ui/src/pierre/index.ts
index f6446f3cc..dc9d857bf 100644
--- a/packages/ui/src/pierre/index.ts
+++ b/packages/ui/src/pierre/index.ts
@@ -13,7 +13,7 @@ export type DiffProps<T = {}> = FileDiffOptions<T> & {
 }
 
 const unsafeCSS = `
-[data-diffs] {
+[data-diff] {
   --diffs-bg: light-dark(var(--diffs-light-bg), var(--diffs-dark-bg));
   --diffs-bg-buffer: var(--diffs-bg-buffer-override, light-dark( color-mix(in lab, var(--diffs-bg) 92%, var(--diffs-mixer)), color-mix(in lab, var(--diffs-bg) 92%, var(--diffs-mixer))));
   --diffs-bg-hover: var(--diffs-bg-hover-override, light-dark( color-mix(in lab, var(--diffs-bg) 97%, var(--diffs-mixer)), color-mix(in lab, var(--diffs-bg) 91%, var(--diffs-mixer))));
@@ -44,7 +44,7 @@ const unsafeCSS = `
   --diffs-bg-selection-text: rgb(from var(--surface-warning-strong) r g b / 0.2);
 }
 
-:host([data-color-scheme='dark']) [data-diffs] {
+:host([data-color-scheme='dark']) [data-diff] {
   --diffs-selection-number-fg: #fdfbfb;
   --diffs-bg-selection: var(--diffs-bg-selection-override, rgb(from var(--solaris-dark-6) r g b / 0.65));
   --diffs-bg-selection-number: var(
@@ -53,7 +53,7 @@ const unsafeCSS = `
   );
 }
 
-[data-diffs] ::selection {
+[data-diff] ::selection {
   background-color: var(--diffs-bg-selection-text);
 }
 
@@ -65,61 +65,48 @@ const unsafeCSS = `
   background-color: rgb(from var(--surface-warning-strong) r g b / 0.55);
 }
 
-[data-diffs] [data-comment-selected]:not([data-selected-line]) [data-column-content] {
+[data-diff] [data-line][data-comment-selected]:not([data-selected-line]) {
   box-shadow: inset 0 0 0 9999px var(--diffs-bg-selection);
 }
 
-[data-diffs] [data-comment-selected]:not([data-selected-line]) [data-column-number] {
+[data-diff] [data-column-number][data-comment-selected]:not([data-selected-line]) {
   box-shadow: inset 0 0 0 9999px var(--diffs-bg-selection-number);
   color: var(--diffs-selection-number-fg);
 }
 
-[data-diffs] [data-selected-line] {
+[data-diff] [data-line-annotation][data-comment-selected]:not([data-selected-line]) [data-annotation-content] {
+  box-shadow: inset 0 0 0 9999px var(--diffs-bg-selection);
+}
+
+[data-diff] [data-line][data-selected-line] {
   background-color: var(--diffs-bg-selection);
   box-shadow: inset 2px 0 0 var(--diffs-selection-border);
 }
 
-[data-diffs] [data-selected-line] [data-column-number] {
+[data-diff] [data-column-number][data-selected-line] {
   background-color: var(--diffs-bg-selection-number);
   color: var(--diffs-selection-number-fg);
 }
 
-[data-diffs] [data-line-type='context'][data-selected-line] [data-column-number],
-[data-diffs] [data-line-type='context-expanded'][data-selected-line] [data-column-number],
-[data-diffs] [data-line-type='change-addition'][data-selected-line] [data-column-number],
-[data-diffs] [data-line-type='change-deletion'][data-selected-line] [data-column-number] {
+[data-diff] [data-column-number][data-line-type='context'][data-selected-line],
+[data-diff] [data-column-number][data-line-type='context-expanded'][data-selected-line],
+[data-diff] [data-column-number][data-line-type='change-addition'][data-selected-line],
+[data-diff] [data-column-number][data-line-type='change-deletion'][data-selected-line] {
   color: var(--diffs-selection-number-fg);
 }
 
 /* The deletion word-diff emphasis is stronger than additions; soften it while selected so the selection highlight reads consistently. */
-[data-diffs] [data-line-type='change-deletion'][data-selected-line] {
+[data-diff] [data-line][data-line-type='change-deletion'][data-selected-line] {
   --diffs-bg-deletion-emphasis: light-dark(
     rgb(from var(--diffs-deletion-base) r g b / 0.07),
     rgb(from var(--diffs-deletion-base) r g b / 0.1)
   );
 }
 
-[data-diffs-header],
-[data-diffs] {
-  [data-separator-wrapper] {
-    margin: 0 !important;
-    border-radius: 0 !important;
-  }
-  [data-expand-button] {
-    width: 6.5ch !important;
-    height: 24px !important;
-    justify-content: end !important;
-    padding-left: 3ch !important;
-    padding-inline: 1ch !important;
-  }
-  [data-separator-multi-button] {
-    grid-template-rows: 10px 10px !important;
-    [data-expand-button] {
-      height: 12px !important;
-    }
-  }
-  [data-separator-content] {
-    height: 24px !important;
+[data-diff-header],
+[data-diff] {
+  [data-separator] {
+    height: 24px;
   }
   [data-column-number] {
     background-color: var(--background-stronger);
@@ -146,28 +133,15 @@ export function createDefaultOptions<T>(style: FileDiffOptions<T>["diffStyle"])
     overflow: "wrap",
     diffStyle: style ?? "unified",
     diffIndicators: "bars",
+    lineHoverHighlight: "both",
     disableBackground: false,
     expansionLineCount: 20,
+    hunkSeparators: "line-info-basic",
     lineDiffType: style === "split" ? "word-alt" : "none",
     maxLineDiffLength: 1000,
     maxLineLengthForHighlighting: 1000,
     disableFileHeader: true,
     unsafeCSS,
-    // hunkSeparators(hunkData: HunkData) {
-    //   const fragment = document.createDocumentFragment()
-    //   const numCol = document.createElement("div")
-    //   numCol.innerHTML = `<svg data-slot="diff-hunk-separator-line-number-icon" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M8.97978 14.0204L8.62623 13.6668L9.33334 12.9597L9.68689 13.3133L9.33333 13.6668L8.97978 14.0204ZM12 16.3335L12.3535 16.6871L12 17.0406L11.6464 16.687L12 16.3335ZM14.3131 13.3133L14.6667 12.9597L15.3738 13.6668L15.0202 14.0204L14.6667 13.6668L14.3131 13.3133ZM12.5 16.0002V16.5002H11.5V16.0002H12H12.5ZM9.33333 13.6668L9.68689 13.3133L12.3535 15.9799L12 16.3335L11.6464 16.687L8.97978 14.0204L9.33333 13.6668ZM12 16.3335L11.6464 15.9799L14.3131 13.3133L14.6667 13.6668L15.0202 14.0204L12.3535 16.6871L12 16.3335ZM6.5 8.00016V7.50016H8.5V8.00016V8.50016H6.5V8.00016ZM9.5 8.00016V7.50016H11.5V8.00016V8.50016H9.5V8.00016ZM12.5 8.00016V7.50016H14.5V8.00016V8.50016H12.5V8.00016ZM15.5 8.00016V7.50016H17.5V8.00016V8.50016H15.5V8.00016ZM12 10.5002H12.5V16.0002H12H11.5V10.5002H12Z" fill="currentColor"/></svg> `
-    //   numCol.dataset["slot"] = "diff-hunk-separator-line-number"
-    //   fragment.appendChild(numCol)
-    //   const contentCol = document.createElement("div")
-    //   contentCol.dataset["slot"] = "diff-hunk-separator-content"
-    //   const span = document.createElement("span")
-    //   span.dataset["slot"] = "diff-hunk-separator-content-span"
-    //   span.textContent = `${hunkData.lines} unmodified lines`
-    //   contentCol.appendChild(span)
-    //   fragment.appendChild(contentCol)
-    //   return fragment
-    // },
   } as const
 }
 
diff --git a/packages/ui/src/pierre/virtualizer.ts b/packages/ui/src/pierre/virtualizer.ts
new file mode 100644
index 000000000..4957afc12
--- /dev/null
+++ b/packages/ui/src/pierre/virtualizer.ts
@@ -0,0 +1,76 @@
+import { type VirtualFileMetrics, Virtualizer } from "@pierre/diffs"
+
+type Target = {
+  key: Document | HTMLElement
+  root: Document | HTMLElement
+  content: HTMLElement | undefined
+}
+
+type Entry = {
+  virtualizer: Virtualizer
+  refs: number
+}
+
+const cache = new WeakMap<Document | HTMLElement, Entry>()
+
+export const virtualMetrics: Partial<VirtualFileMetrics> = {
+  lineHeight: 24,
+  hunkSeparatorHeight: 24,
+  fileGap: 0,
+}
+
+function target(container: HTMLElement): Target | undefined {
+  if (typeof document === "undefined") return
+
+  const root = container.closest("[data-component='session-review']")
+  if (root instanceof HTMLElement) {
+    const content = root.querySelector("[data-slot='session-review-container']")
+    return {
+      key: root,
+      root,
+      content: content instanceof HTMLElement ? content : undefined,
+    }
+  }
+
+  return {
+    key: document,
+    root: document,
+    content: undefined,
+  }
+}
+
+export function acquireVirtualizer(container: HTMLElement) {
+  const resolved = target(container)
+  if (!resolved) return
+
+  let entry = cache.get(resolved.key)
+  if (!entry) {
+    const virtualizer = new Virtualizer()
+    virtualizer.setup(resolved.root, resolved.content)
+    entry = {
+      virtualizer,
+      refs: 0,
+    }
+    cache.set(resolved.key, entry)
+  }
+
+  entry.refs += 1
+  let done = false
+
+  return {
+    virtualizer: entry.virtualizer,
+    release() {
+      if (done) return
+      done = true
+
+      const current = cache.get(resolved.key)
+      if (!current) return
+
+      current.refs -= 1
+      if (current.refs > 0) return
+
+      current.virtualizer.cleanUp()
+      cache.delete(resolved.key)
+    },
+  }
+}
diff --git a/packages/ui/src/pierre/worker.ts b/packages/ui/src/pierre/worker.ts
index 0d117c368..1993ad7aa 100644
--- a/packages/ui/src/pierre/worker.ts
+++ b/packages/ui/src/pierre/worker.ts
@@ -21,6 +21,7 @@ function createPool(lineDiffType: "none" | "word-alt") {
     {
       theme: "OpenCode",
       lineDiffType,
+      preferredHighlighter: "shiki-wasm",
     },
   )