[codex] Add MDX source highlighting

packages/cspec-editor/src/index.ts

@@ -2,7 +2,7 @@ import { autocompletion, type Completion, type CompletionContext, type Completio
 import { defaultKeymap, history, historyKeymap } from "@codemirror/commands";
 import { markdown } from "@codemirror/lang-markdown";
 import { linter, type Diagnostic as CmDiagnostic } from "@codemirror/lint";
-import { EditorSelection, EditorState, StateEffect, StateField, type Extension } from "@codemirror/state";
+import { EditorSelection, EditorState, StateEffect, StateField, type Extension, type Range } from "@codemirror/state";
 import { Decoration, EditorView, keymap, ViewPlugin, type DecorationSet, type ViewUpdate, WidgetType } from "@codemirror/view";
 import { tags } from "@lezer/highlight";
 import { HighlightStyle, syntaxHighlighting } from "@codemirror/language";
@@ -98,6 +98,7 @@ export function createCspecEditorExtensions(options: CspecEditorOptions): Extens
     keymap.of([...defaultKeymap, ...historyKeymap]),
     autocompletion({ override: [(context) => cspecCompletions(context, options.file)] }),
     linter((view) => cspecLint(view, options.file)),
+    mdxSourceSyntaxHighlighting(),
     cspecDecorations(options.file, options.onJumpToDefinition),
     cspecJumpHandler(options.file, options.onJumpToDefinition),
     cspecTheme
@@ -240,6 +241,103 @@ function cspecDecorations(file: string, onJumpToDefinition?: (target: CspecJumpT
   });
 }
 
+function mdxSourceSyntaxHighlighting() {
+  return ViewPlugin.fromClass(class {
+    decorations: DecorationSet;
+    constructor(view: EditorView) {
+      this.decorations = buildMdxSourceSyntaxDecorations(view);
+    }
+    update(update: ViewUpdate) {
+      if (update.docChanged || update.viewportChanged || update.transactions.some((tr) => tr.effects.length)) {
+        this.decorations = buildMdxSourceSyntaxDecorations(update.view);
+      }
+    }
+  }, {
+    decorations: (value) => value.decorations
+  });
+}
+
+type DecorationRange = Range<Decoration>;
+
+function buildMdxSourceSyntaxDecorations(view: EditorView): DecorationSet {
+  if (!view.state.field(sourceModeField)) return Decoration.none;
+  const source = view.state.doc.toString();
+  const ranges = [
+    ...mdxImportExportSyntax(source),
+    ...mdxJsxSyntax(source),
+    ...mdxEmbedSyntax(source)
+  ];
+  return Decoration.set(ranges.sort((a, b) => a.from - b.from || a.to - b.to), true);
+}
+
+function mdxImportExportSyntax(source: string): DecorationRange[] {
+  const ranges: DecorationRange[] = [];
+  const linePattern = /^[^\S\n]*(?:import|export)\b[^\n]*/gm;
+  for (const match of source.matchAll(linePattern)) {
+    const line = match[0] ?? "";
+    const start = match.index ?? 0;
+    const keyword = line.match(/\b(?:import|export)\b/);
+    if (keyword?.index !== undefined) {
+      ranges.push(mdxMark("keyword").range(start + keyword.index, start + keyword.index + keyword[0].length));
+    }
+    addStringMarks(line, start, ranges);
+  }
+  return ranges;
+}
+
+function mdxJsxSyntax(source: string): DecorationRange[] {
+  const ranges: DecorationRange[] = [];
+  const tagPattern = /<\/?([A-Za-z][\w.$:-]*)([^<>]*?)\/?>/g;
+  for (const match of source.matchAll(tagPattern)) {
+    const text = match[0] ?? "";
+    const start = match.index ?? 0;
+    const name = match[1] ?? "";
+    const attrs = match[2] ?? "";
+    const openingLength = text.startsWith("</") ? 2 : 1;
+    const nameStart = start + openingLength;
+    const closingLength = text.endsWith("/>") ? 2 : 1;
+    ranges.push(mdxMark("punctuation").range(start, start + openingLength));
+    ranges.push(mdxMark("tagName").range(nameStart, nameStart + name.length));
+    ranges.push(mdxMark("punctuation").range(start + text.length - closingLength, start + text.length));
+
+    const attrsStart = nameStart + name.length;
+    const attrPattern = /\b([A-Za-z_:][\w:.-]*)(?=\s*=)/g;
+    for (const attr of attrs.matchAll(attrPattern)) {
+      const attrName = attr[1] ?? "";
+      const attrStart = attrsStart + (attr.index ?? 0);
+      ranges.push(mdxMark("attribute").range(attrStart, attrStart + attrName.length));
+    }
+    addStringMarks(attrs, attrsStart, ranges);
+  }
+  return ranges;
+}
+
+function mdxEmbedSyntax(source: string): DecorationRange[] {
+  const ranges: DecorationRange[] = [];
+  const embedPattern = /\{(?:(?:ref\((?:"(?:[^"\\]|\\.)*"|'(?:[^'\\]|\\.)*')\))|(?:[A-Za-z_$][\w$]*(?:(?:\.[A-Za-z_$][\w$]*)|(?:\["(?:[^"\\]|\\.)*"\]))*))\.\$\}/g;
+  for (const match of source.matchAll(embedPattern)) {
+    const text = match[0] ?? "";
+    const start = match.index ?? 0;
+    ranges.push(mdxMark("punctuation").range(start, start + 1));
+    ranges.push(mdxMark("embed").range(start + 1, start + text.length - 1));
+    ranges.push(mdxMark("punctuation").range(start + text.length - 1, start + text.length));
+    addStringMarks(text.slice(1, -1), start + 1, ranges);
+  }
+  return ranges;
+}
+
+function addStringMarks(text: string, offset: number, ranges: DecorationRange[]): void {
+  const stringPattern = /"([^"\\]|\\.)*"|'([^'\\]|\\.)*'/g;
+  for (const match of text.matchAll(stringPattern)) {
+    const start = offset + (match.index ?? 0);
+    ranges.push(mdxMark("string").range(start, start + (match[0] ?? "").length));
+  }
+}
+
+function mdxMark(name: "attribute" | "embed" | "keyword" | "punctuation" | "string" | "tagName"): Decoration {
+  return Decoration.mark({ class: `cm-cspec-mdx-${name}` });
+}
+
 function buildDecorations(view: EditorView, file: string, onJumpToDefinition?: (target: CspecJumpTarget) => void): DecorationSet {
   if (view.state.field(sourceModeField)) return Decoration.none;
   const source = view.state.doc.toString();
@@ -614,5 +712,26 @@ const cspecTheme = EditorView.theme({
     background: "#eef1f4",
     borderRadius: "4px",
     padding: "0 3px"
+  },
+  ".cm-cspec-mdx-keyword": {
+    color: "#8a3ffc",
+    fontWeight: "700"
+  },
+  ".cm-cspec-mdx-tagName": {
+    color: "#0f766e",
+    fontWeight: "700"
+  },
+  ".cm-cspec-mdx-attribute": {
+    color: "#8a4b0f"
+  },
+  ".cm-cspec-mdx-string": {
+    color: "#0b6bcb"
+  },
+  ".cm-cspec-mdx-punctuation": {
+    color: "#6b7280"
+  },
+  ".cm-cspec-mdx-embed": {
+    color: "#b42318",
+    fontWeight: "700"
   }
 });

packages/cspec-editor/test.js

@@ -165,6 +165,10 @@ test("decorations and source mode preserve the canonical document", async () =>
   assert.equal(view.state.doc.toString(), source);
   assert.equal(parent.querySelector(".cm-cspec-blockHeader"), null);
   assert.equal(parent.querySelector(".cm-cspec-embed"), null);
+  assert.ok(parent.querySelector(".cm-cspec-mdx-keyword"));
+  assert.ok(parent.querySelector(".cm-cspec-mdx-tagName"));
+  assert.ok(parent.querySelector(".cm-cspec-mdx-attribute"));
+  assert.ok(parent.querySelector(".cm-cspec-mdx-embed"));
 
   view.dispatch({ effects: setSourceMode.of(false) });
   view.dispatch({ selection: { anchor: source.indexOf("Prompt text.") } });