wip(ui): diff virtualization (#12693)

1 files changed, 13 insertions, 2 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", "")
+      }
     }
   }