fix(highlight): parse yaml diff fragments per indentation segment

src/background.ts

@@ -120,6 +120,10 @@ function highlight(langName: string, text: string) {
     return Promise.resolve({ rows: hljsRows(langName.slice(5), text) });
   if (langName === "gotmpl") return highlightGotmpl(text);
   if (langName === "markdown") return highlightMarkdown(text);
+  if (langName === "yaml")
+    return init()
+      .then(() => loadLang("yaml"))
+      .then((yaml: Lang) => ({ rows: yamlFragmentRows(yaml, text) }));
   return init()
     .then(() => loadLang(langName))
     .then(({ language, query }: Lang) => {
@@ -134,6 +138,36 @@ function highlight(langName: string, text: string) {
     });
 }
 
+// yaml roots its document at the first line's indent, so a diff fragment that
+// later dedents below it — any hunk that starts deep inside a document — turns
+// the rest of the parse into one ERROR node with no captures. Parse each
+// maximal run of lines that never dedents below its opening indent as its own
+// document and shift the captured rows back into place; an ordinary full file
+// is a single segment, i.e. exactly one parse like before.
+function yamlFragmentRows(yaml: Lang, text: string): Rows {
+  const lines = text.split("\n");
+  const starts: number[] = [];
+  let open = Infinity;
+  for (let i = 0; i < lines.length; i++) {
+    const line = lines[i];
+    if (!line.trim()) continue;
+    const indent = line.length - line.trimStart().length;
+    if (indent < open) {
+      starts.push(i);
+      open = indent;
+    }
+  }
+  const rows: Rows = {};
+  for (let k = 0; k < starts.length; k++) {
+    const from = starts[k];
+    const to = k + 1 < starts.length ? starts[k + 1] : lines.length;
+    const nested: Rows = {};
+    captureInto(nested, yaml, lines.slice(from, to).join("\n"));
+    mergeRows(rows, nested, from, 0);
+  }
+  return rows;
+}
+
 // Helm/Go templates: highlight as yaml with the {{ … }} actions on top
 // (like nvim-treesitter's gotmpl setup). The yaml grammar chokes on inline
 // template syntax, so it parses a copy with every action blanked out — same
@@ -146,11 +180,9 @@ function highlightGotmpl(text: string) {
       const lineStarts = lineStartsOf(text);
       const rows: Rows = {};
 
+      // masking keeps every line length, so the fragment rows line up 1:1
       const masked = text.replace(/\{\{[\s\S]*?\}\}/g, (m) => m.replace(/[^\n]/g, " "));
-      const yamlTree = parseWith(yaml.language, masked);
-      for (const { name, node } of yaml.query.captures(yamlTree.rootNode).slice(0, MAX_CAPTURES))
-        pushNode(rows, node, cssClass(name), lineStarts, text);
-      yamlTree.delete();
+      mergeRows(rows, yamlFragmentRows(yaml, masked), 0, 0);
 
       const gtTree = parseWith(gt.language, text);
       for (const { name, node } of gt.query.captures(gtTree.rootNode).slice(0, MAX_CAPTURES))

test/e2e/grammars.spec.ts

@@ -54,3 +54,22 @@ test("added tree-sitter grammars and the markdown fallback produce tokens", asyn
   const md = page.locator("section.pt-file", { hasText: "NOTES.md" });
   await expect(md.locator(".pt-keyword").first()).toBeVisible({ timeout: 20000 });
 });
+
+// a hunk from deep inside a yaml document dedents below its first line; the
+// per-segment parse must still colour the keys after the dedent
+test("a yaml hunk that dedents below its opening line still highlights", async ({
+  context,
+  page,
+}) => {
+  const body = readFileSync(path.join(__dirname, "../fixtures/yaml-fragment.diff"), "utf8");
+  const url = "https://example.com/yaml-fragment.diff";
+  await context.route(url, (route) =>
+    route.fulfill({ contentType: "text/plain; charset=utf-8", body })
+  );
+  await page.goto(url);
+  const section = page.locator("section.pt-file", { hasText: "crds/machines.yaml" });
+  await expect(section.locator(".pt-property", { hasText: "gpus" }).first()).toBeVisible({
+    timeout: 20000,
+  });
+  await expect(section.locator(".pt-property", { hasText: "gpuClassName" }).first()).toBeVisible();
+});

test/fixtures/yaml-fragment.diff

@@ -0,0 +1,13 @@
+diff --git a/crds/machines.yaml b/crds/machines.yaml
+--- a/crds/machines.yaml
++++ b/crds/machines.yaml
+@@ -1067,3 +1067,9 @@ spec:
+                       name:
+                         description: |
+                           The name of the device.
++                gpus:
++                  type: array
++                  items:
++                    properties:
++                      gpuClassName:
++                        type: string