Fix bug in rehype-hljs-var

Sometimes, hljs generates text nodes that are direct descendents of a
`code` element, without being the children of a `span`. rehype-hljs-var
expects all text nodes to be the sole children of an element node.

This change edits rehype-hljs-var to accommodate text nodes. It also
refactors the plugin to clean up the logic for splitting text nodes with
Var component placeholders.

server/fixtures/hcl-addr-var.mdx

@@ -0,0 +1,29 @@
+1. Create a `main.tf` file containing this minimal Terraform code:
+
+   ```hcl
+   terraform {
+     required_providers {
+       teleport = {
+         source  = "terraform.releases.teleport.dev/gravitational/teleport"
+         version = "~> (=teleport.major_version=).0"
+       }
+     }
+   }
+   
+   provider "teleport" {
+     addr               = '<Var name="teleport.example.com:443" />'
+   }
+   ```
+
+1. Then, init your Terraform working directory to download the Teleport provider:
+
+   ```code
+   $ terraform init
+   Initializing the backend...
+   
+   Initializing provider plugins...
+   - Finding terraform.releases.teleport.dev/gravitational/teleport versions matching ...
+   ```
+
+1. Finally, run a Terraform plan:
+

server/fixtures/result/hcl-addr-var.html

@@ -0,0 +1,30 @@
+<ol>
+<li>
+<p>Create a <code>main.tf</code> file containing this minimal Terraform code:</p>
+<pre><code class="hljs language-hcl"><span class="hljs-keyword">terraform</span> {
+  required_providers {
+    teleport = {
+      source  = <span class="hljs-string">"terraform.releases.teleport.dev/gravitational/teleport"</span>
+      version = <span class="hljs-string">"~> (=teleport.major_version=).0"</span>
+    }
+  }
+}
+
+<span class="hljs-keyword">provider</span> <span class="hljs-string">"teleport"</span> {
+  addr               = '<var name="teleport.example.com:443"></var>'
+}
+</code></pre>
+</li>
+<li>
+<p>Then, init your Terraform working directory to download the Teleport provider:</p>
+<pre><code class="hljs language-code">$ terraform init
+Initializing the backend...
+
+Initializing provider plugins...
+- Finding terraform.releases.teleport.dev/gravitational/teleport versions matching ...
+</code></pre>
+</li>
+<li>
+<p>Finally, run a Terraform plan:</p>
+</li>
+</ol>

server/rehype-hljs-var.ts

@@ -1,16 +1,5 @@
 import { unified, Transformer } from "unified";
 import type { VFile } from "vfile";
-import type {
-  MdxJsxFlowElement,
-  MdxJsxTextElement,
-  MdxJsxAttribute,
-  MdxJsxAttributeValueExpression,
-} from "mdast-util-mdx-jsx";
-import type { MDXJSEsm } from "mdast-util-mdxjs-esm";
-import type {
-  MDXFlowExpression,
-  MDXTextExpression,
-} from "mdast-util-mdx-expression";
 import rehypeHighlight, {
   Options as RehypeHighlightOptions,
 } from "rehype-highlight";
@@ -35,46 +24,47 @@ export const rehypeVarInHLJS = (
   return (root: Parent, file: VFile) => {
     const highlighter = rehypeHighlight(options);
 
-    let placeholdersToVars: Record<string, MdxJsxFlowElement> = {};
+    let placeholdersToVars: Record<string, Node> = {};
 
     // In a code snippet, Var elements are parsed as text. Replace these with
     // UUID strings to ensure that the parser won't split these up and make
     // them unrecoverable.
     visit(root, undefined, (node: Node, index: number, parent: Parent) => {
+      // We only visit text nodes inside code snippets
       if (
-        node.type === "text" &&
-        parent.hasOwnProperty("tagName") &&
-        (parent as Element).tagName === "code"
+        node.type !== "text" ||
+        !parent.hasOwnProperty("tagName") ||
+        (parent as Element).tagName !== "code"
       ) {
-        const varPattern = new RegExp("<Var [^>]+/>", "g");
-        (node as Text).value = (node as Text).value.replace(
-          varPattern,
-          (match) => {
-            const placeholder = makePlaceholder();
-            // Since the Var element was originally text, parse it so we can recover
-            // its properties. The result should be a small HTML AST with a root
-            // node and one child, the Var node.
-            const varElement = unified()
-              .use(remarkParse)
-              .use(remarkMDX)
-              .parse(match);
-            if (
-              varElement.children.length !== 1 ||
-              (varElement.children[0] as MdxJsxFlowElement).name !== "Var"
-            ) {
-              throw new Error(
-                `Problem parsing file ${file.path}: malformed Var element within a code snippet`
-              );
-            }
-
-            placeholdersToVars[placeholder] = varElement
-              .children[0] as MdxJsxFlowElement;
-            return placeholder;
-          }
-        );
+        return [CONTINUE];
       }
-    });
 
+      const varPattern = new RegExp("<Var [^>]+/>", "g");
+      (node as Text).value = (node as Text).value.replace(
+        varPattern,
+        (match) => {
+          const placeholder = makePlaceholder();
+          // Since the Var element was originally text, parse it so we can recover
+          // its properties. The result should be a small HTML AST with a root
+          // node and one child, the Var node.
+          const varElement = unified()
+            .use(remarkParse)
+            .use(remarkMDX)
+            .parse(match);
+          if (
+            varElement.children.length !== 1 ||
+            (varElement.children[0] as any).name !== "Var"
+          ) {
+            throw new Error(
+              `Problem parsing file ${file.path}: malformed Var element within a code snippet`
+            );
+          }
+
+          placeholdersToVars[placeholder] = varElement.children[0];
+          return placeholder;
+        }
+      );
+    });
     // Apply syntax highlighting
     (highlighter as Function)(root);
 
@@ -83,91 +73,102 @@ export const rehypeVarInHLJS = (
     // placeholder UUIDs and replace them with their original Var elements,
     // inserting these as HTML AST nodes.
     visit(root, undefined, (node: Node, index: number, parent: Parent) => {
-      const el = node as Element;
+      const el = node as Element | Text;
+      // We expect the element to have a single text node or be a single text
+      // node.
       if (
-        el.type === "element" &&
-        el.children.length === 1 &&
-        el.children[0].type === "text"
+        !(
+          el.type === "element" &&
+          el.children.length === 1 &&
+          el.children[0].type === "text"
+        ) &&
+        !(el.type === "text")
       ) {
-        const hljsSpanValue = (el.children[0] as Text).value;
+        return [CONTINUE];
+      }
 
-        // This is an hljs span with only the placeholder as its child.
-        // We don't need the span, so replace it with the original Var.
-        if (placeholdersToVars[hljsSpanValue]) {
-          (parent as any).children[index] = placeholdersToVars[hljsSpanValue];
-          return [CONTINUE];
-        }
+      let hljsSpanValue = "";
+      if (el.type === "text") {
+        hljsSpanValue = el.value;
+      } else {
+        hljsSpanValue = (el.children[0] as Text).value;
+      }
 
-        const placeholders = Array.from(
-          hljsSpanValue.matchAll(new RegExp(placeholderPattern, "g"))
-        );
+      // This is an hljs span with only the placeholder as its child.
+      // We don't need the span, so replace it with the original Var.
+      if (placeholdersToVars[hljsSpanValue]) {
+        (parent as any).children[index] = placeholdersToVars[hljsSpanValue];
+        return [CONTINUE];
+      }
 
-        // No placeholders to recover, so there's nothing more to do.
-        if (placeholders.length == 0) {
-          return [CONTINUE];
-        }
+      const placeholders = Array.from(
+        hljsSpanValue.matchAll(new RegExp(placeholderPattern, "g"))
+      );
 
-        // An hljs span's text includes one or more Vars among other content, so
-        // we need to replace the span with a series of spans separated by
-        // Vars.
-        let lastIndex = 0;
-        let newChildren: Array<MdxJsxFlowElement | Element> = [];
-        // If there is content before the first Var, separate it into a new hljs
-        // span.
-        if (placeholders[0].index > 0) {
+      // No placeholders to recover, so there's nothing more to do.
+      if (placeholders.length == 0) {
+        return [CONTINUE];
+      }
+
+      // The element's text includes one or more Vars among other content, so
+      // we need to replace the span with a series of spans separated by
+      // Vars.
+      let newChildren: Array<Text | Element> = [];
+
+      // Assemble a map of indexes to their corresponding placeholders so we
+      // can tell whether a given index falls within a placeholder.
+      const placeholderIndices = new Map();
+      placeholders.forEach((p) => {
+        placeholderIndices.set(p.index, p[0]);
+      });
+
+      let valueIdx = 0;
+      while (valueIdx < hljsSpanValue.length) {
+        // The current index is in a placeholder, so add the original Var
+        // component to newChildren.
+        if (placeholderIndices.has(valueIdx)) {
+          const placeholder = placeholderIndices.get(valueIdx);
+          valueIdx += placeholder.length;
+          newChildren.push(placeholdersToVars[placeholder] as Element);
+          continue;
+        }
+        // The current index is outside a placeholder, so assemble a text or
+        // span node and push that to newChildren.
+        let textVal = "";
+        while (
+          !placeholderIndices.has(valueIdx) &&
+          valueIdx < hljsSpanValue.length
+        ) {
+          textVal += hljsSpanValue[valueIdx];
+          valueIdx++;
+        }
+        if (el.type === "text") {
+          newChildren.push({
+            type: "text",
+            value: textVal,
+          });
+        } else {
           newChildren.push({
             tagName: "span",
             type: "element",
             properties: el.properties,
             children: [
               {
                 type: "text",
-                value: hljsSpanValue.substring(
-                  lastIndex,
-                  placeholders[0].index
-                ),
+                value: textVal,
               },
             ],
           });
-          lastIndex = placeholders[0].index;
         }
-        placeholders.forEach((ph, i) => {
-          const placeholderValue = ph[0];
-          newChildren.push(placeholdersToVars[placeholderValue]);
-          lastIndex += placeholderValue.length;
-
-          // Check if there is some non-Var text between either (a) this and the
-          // next Var or (b) between this Var and the end of the content. If
-          // so, add another span and advance the last index.
-          let nextIndex = 0;
-          if (i < placeholders.length - 1) {
-            nextIndex = placeholders[i + 1].index;
-          } else if (i == placeholders.length - 1) {
-            nextIndex = hljsSpanValue.length;
-          }
-          if (lastIndex < nextIndex) {
-            newChildren.push({
-              tagName: "span",
-              type: "element",
-              properties: el.properties,
-              children: [
-                {
-                  type: "text",
-                  value: hljsSpanValue.substring(lastIndex, nextIndex),
-                },
-              ],
-            });
-            lastIndex = nextIndex;
-          }
-        });
-        // Delete the current span and replace it with the new children.
-        (parent.children as Array<MdxJsxFlowElement | Element>).splice(
-          index,
-          1,
-          ...newChildren
-        );
-        return [SKIP, index + newChildren.length];
       }
+
+      // Delete the current span and replace it with the new children.
+      (parent.children as Array<Text | Element>).splice(
+        index,
+        1,
+        ...newChildren
+      );
+      return [SKIP, index + newChildren.length];
     });
   };
 };

uvu-tests/rehype-hljs-var.test.ts

@@ -145,4 +145,19 @@ Suite("Ignore VarList in code snippet components", () => {
   });
 });
 
+Suite("Next node as one of several code node children", () => {
+  const result = transformer({
+    value: readFileSync(resolve("server/fixtures/hcl-addr-var.mdx"), "utf-8"),
+    path: "/docs/index.mdx",
+  });
+
+  assert.equal(
+    (result.value as string).trim(),
+    readFileSync(
+      resolve("server/fixtures/result/hcl-addr-var.html"),
+      "utf-8"
+    ).trim()
+  );
+});
+
 Suite.run();