Fix extract annotations tests

Author: pomber

packages/codehike/src/code/extract-annotations.test.ts

@@ -1,4 +1,4 @@
-import { expect, test } from "vitest"
+import { expect, test, vi } from "vitest"
 import { splitAnnotationsAndCode } from "./extract-annotations.js"
 
 async function t(comment: string) {
@@ -153,6 +153,36 @@ test("multiple start/end pairs of same name", async () => {
   expect(r0.toLineNumber).toBeLessThan(r1.fromLineNumber)
 })
 
+test("same-name nested start/end pairs preserve nesting", async () => {
+  const code = [
+    "// !focus(start)",
+    "const outer = 1",
+    "// !focus(start)",
+    "const inner = 2",
+    "// !focus(end)",
+    "const outerTail = 3",
+    "// !focus(end)",
+    "const after = 4",
+  ].join("\n")
+  const { annotations } = await splitAnnotationsAndCode(
+    code,
+    "javascript",
+    "!",
+  )
+
+  expect(annotations).toHaveLength(2)
+  expect(annotations[0].name).toEqual("focus")
+  expect(annotations[0].ranges[0]).toEqual({
+    fromLineNumber: 1,
+    toLineNumber: 3,
+  })
+  expect(annotations[1].name).toEqual("focus")
+  expect(annotations[1].ranges[0]).toEqual({
+    fromLineNumber: 2,
+    toLineNumber: 2,
+  })
+})
+
 test("different annotation names with start/end", async () => {
   const code = [
     "// !focus(start)",
@@ -195,6 +225,31 @@ test("start/end removes comment lines from code", async () => {
   expect(lines[2]).toContain("let c = 3")
 })
 
+test("adjacent start/end markers are ignored instead of creating empty ranges", async () => {
+  const warn = vi.spyOn(console, "warn").mockImplementation(() => {})
+  try {
+    const code = [
+      "// !focus(start)",
+      "// !focus(end)",
+      "const x = 1",
+    ].join("\n")
+
+    const { code: resultCode, annotations } = await splitAnnotationsAndCode(
+      code,
+      "javascript",
+      "!",
+    )
+
+    expect(resultCode).toEqual("const x = 1")
+    expect(annotations).toHaveLength(0)
+    expect(warn).toHaveBeenCalledWith(
+      "Code Hike warning: Empty !focus start/end annotation range",
+    )
+  } finally {
+    warn.mockRestore()
+  }
+})
+
 test("start/end works with Python comments", async () => {
   const code = [
     "x = 1",

packages/codehike/src/code/extract-annotations.tsx

@@ -22,70 +22,69 @@ function getLineFromRange(range: any): number {
 }
 
 function processStartEndMarkers(annotations: Annotation[]): Annotation[] {
-  const regular: Annotation[] = []
-  const starts: { name: string; query: string; line: number }[] = []
-  const ends: { name: string; query: string; line: number }[] = []
+  const stacks = new Map<
+    string,
+    { name: string; query: string; line: number; order: number }[]
+  >()
+  const orderedAnnotations: { order: number; annotation: Annotation }[] = []
 
-  for (const a of annotations) {
+  for (const [order, a] of annotations.entries()) {
     const q = a.query ?? ""
     if (q.startsWith(START_MARKER)) {
-      starts.push({
+      const start = {
         name: a.name,
         query: q.slice(START_MARKER.length),
         line: getLineFromRange(a.ranges[0]),
-      })
+        order,
+      }
+      const stack = stacks.get(a.name) ?? []
+      stack.push(start)
+      stacks.set(a.name, stack)
     } else if (q.startsWith(END_MARKER)) {
-      ends.push({
-        name: a.name,
-        query: q.slice(END_MARKER.length),
-        line: getLineFromRange(a.ranges[0]),
+      const stack = stacks.get(a.name)
+      const start = stack?.pop()
+      if (!start) {
+        console.warn(
+          `Code Hike warning: Unmatched !${a.name}(end) annotation`,
+        )
+        continue
+      }
+
+      const endLine = getLineFromRange(a.ranges[0]) - 1
+      if (endLine < start.line) {
+        console.warn(
+          `Code Hike warning: Empty !${a.name} start/end annotation range`,
+        )
+        continue
+      }
+
+      orderedAnnotations.push({
+        order: start.order,
+        annotation: {
+          name: start.name,
+          query: start.query,
+          ranges: [{ fromLineNumber: start.line, toLineNumber: endLine }],
+        },
       })
     } else {
-      regular.push(a)
+      orderedAnnotations.push({ order, annotation: a })
     }
   }
 
-  if (starts.length === 0 && ends.length === 0) {
+  if (orderedAnnotations.length === annotations.length) {
     return annotations
   }
 
-  const paired: Annotation[] = []
-  const usedEnds = new Set<number>()
-
-  for (const start of starts) {
-    // find the first unused end with the same name that comes after the start
-    const endIndex = ends.findIndex(
-      (e, i) =>
-        !usedEnds.has(i) &&
-        e.name === start.name &&
-        e.line >= start.line,
-    )
-    if (endIndex === -1) {
+  for (const stack of stacks.values()) {
+    for (const start of stack) {
       console.warn(
         `Code Hike warning: Unmatched !${start.name}(start) annotation`,
       )
-      continue
-    }
-    usedEnds.add(endIndex)
-    const end = ends[endIndex]
-    paired.push({
-      name: start.name,
-      query: start.query,
-      ranges: [
-        { fromLineNumber: start.line, toLineNumber: end.line - 1 },
-      ],
-    })
-  }
-
-  for (let i = 0; i < ends.length; i++) {
-    if (!usedEnds.has(i)) {
-      console.warn(
-        `Code Hike warning: Unmatched !${ends[i].name}(end) annotation`,
-      )
     }
   }
 
-  return [...regular, ...paired]
+  orderedAnnotations.sort((a, b) => a.order - b.order)
+  return orderedAnnotations.map((entry) => entry.annotation)
 }
 
 async function extractCommentAnnotations(