fix(app): line selection styling

1 files changed, 59 insertions, 14 deletions

diff --git a/packages/ui/src/components/diff.tsx b/packages/ui/src/components/diff.tsx
index 5aff93d1d..2fb2ea3bc 100644
--- a/packages/ui/src/components/diff.tsx
+++ b/packages/ui/src/components/diff.tsx
@@ -191,24 +191,69 @@ export function Diff<T>(props: DiffProps<T>) {
       node.removeAttribute("data-comment-selected")
     }
 
-    for (const range of ranges) {
-      const start = Math.max(1, Math.min(range.start, range.end))
-      const end = Math.max(range.start, range.end)
+    const diffs = root.querySelector("[data-diffs]")
+    if (!(diffs instanceof HTMLElement)) return
+
+    const split = diffs.dataset.type === "split"
+
+    const code = Array.from(diffs.querySelectorAll("[data-code]")).filter(
+      (node): node is HTMLElement => node instanceof HTMLElement,
+    )
+    if (code.length === 0) return
+
+    const lineIndex = (element: HTMLElement) => {
+      const raw = element.dataset.lineIndex
+      if (!raw) return
+      const values = raw
+        .split(",")
+        .map((value) => parseInt(value, 10))
+        .filter((value) => !Number.isNaN(value))
+      if (values.length === 0) return
+      if (!split) return values[0]
+      if (values.length === 2) return values[1]
+      return values[0]
+    }
 
-      for (let line = start; line <= end; line++) {
-        const expectedSide =
-          line === end ? (range.endSide ?? range.side) : line === start ? range.side : (range.side ?? range.endSide)
+    const rowIndex = (line: number, side: SelectionSide | undefined) => {
+      const nodes = Array.from(root.querySelectorAll(`[data-line="${line}"], [data-alt-line="${line}"]`)).filter(
+        (node): node is HTMLElement => node instanceof HTMLElement,
+      )
+      if (nodes.length === 0) return
 
-        const nodes = Array.from(root.querySelectorAll(`[data-line="${line}"], [data-alt-line="${line}"]`))
-        for (const node of nodes) {
-          if (!(node instanceof HTMLElement)) continue
+      const targetSide = side ?? "additions"
 
-          if (expectedSide) {
-            const side = findSide(node)
-            if (side && side !== expectedSide) continue
-          }
+      for (const node of nodes) {
+        if (findSide(node) === targetSide) return lineIndex(node)
+        if (parseInt(node.dataset.altLine ?? "", 10) === line) return lineIndex(node)
+      }
+    }
 
-          node.setAttribute("data-comment-selected", "")
+    for (const range of ranges) {
+      const start = rowIndex(range.start, range.side)
+      if (start === undefined) continue
+
+      const end = (() => {
+        const same = range.end === range.start && (range.endSide == null || range.endSide === range.side)
+        if (same) return start
+        return rowIndex(range.end, range.endSide ?? range.side)
+      })()
+      if (end === undefined) continue
+
+      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", "")
+          }
         }
       }
     }