Add CodeMirror extension for Operations

src/editor/index.ts

@@ -41,6 +41,7 @@ import { onMouseDragRegex, onMouseDragMakeANewNumber } from '@src/lib/utils'
 import { themeCompartment } from '@src/editor/plugins/theme'
 import { kclAstExtension } from '@src/editor/plugins/ast'
 import { localHistoryTarget } from '@src/editor/HistoryView'
+import { operationsExtension } from '@src/editor/plugins/operations'
 
 export const lineWrappingCompartment = new Compartment()
 export const cursorBlinkingCompartment = new Compartment()
@@ -64,6 +65,7 @@ export function baseEditorExtensions() {
     historyCompartment.of(history()),
     localHistoryTarget.of([]),
     kclAstExtension(),
+    operationsExtension(),
     closeBrackets(),
     codeFolding(),
     keymap.of([

src/editor/plugins/operations/index.ts

@@ -0,0 +1,89 @@
+import type { Extension, Range } from '@codemirror/state'
+import { StateEffect, StateField, Annotation } from '@codemirror/state'
+import { Decoration, EditorView } from '@codemirror/view'
+import type { Operation } from '@rust/kcl-lib/bindings/Operation'
+
+/** Transaction annotation to identify operation list updates */
+export const operationsAnnotation = Annotation.define<boolean>()
+
+/** Effect used to replace the current operations list in state */
+export const setOperationsEffect = StateEffect.define<Operation[] | null>()
+
+/** StateField to hold the current operations list reference on the EditorState */
+export const operationsStateField = StateField.define<Operation[]>({
+  create() {
+    return []
+  },
+  update(value, tr) {
+    for (const e of tr.effects) {
+      if (e.is(setOperationsEffect)) {
+        return e.value ?? []
+      }
+    }
+    return value
+  },
+})
+
+/**
+ * Decorations field that stores ranges annotated with operations metadata
+ * TODO: Unused until we can find a more responsible way to turn on decoration fields. It's easy to absolutely
+ * tank the editor with nested decorations.
+ */
+export const operationDecorationsField = StateField.define<
+  ReturnType<typeof Decoration.set>
+>({
+  create() {
+    return Decoration.none
+  },
+  update(decorations, tr) {
+    // Recompute all decorations when the artifact graph effect is applied
+    for (const e of tr.effects) {
+      if (e.is(setOperationsEffect)) {
+        const operationsList = e.value ?? []
+        return buildOperationDecorations(operationsList, tr.state)
+      }
+    }
+    // Map existing decorations through document changes
+    if (tr.docChanged) {
+      return decorations.map(tr.changes)
+    }
+    return decorations
+  },
+  provide: (f) => EditorView.decorations.from(f),
+})
+
+/**
+ * Given an ArtifactGraph, apply decorations to the corresponding code ranges with
+ * each artifact's metadata.
+ */
+function buildOperationDecorations(
+  operations: Operation[],
+  state: EditorView['state']
+) {
+  const widgets: Range<Decoration>[] = []
+  const docLen = state.doc.length
+
+  for (const op of operations) {
+    if (op.type === 'GroupEnd') continue
+    const sr = op.sourceRange
+    const from = Math.max(0, Math.min(sr[0], docLen))
+    const to = Math.max(0, Math.min(sr[1], docLen))
+    if (to <= from) continue
+    widgets.push(
+      Decoration.mark({
+        class: 'cm-operation-range',
+        attributes: {
+          'data-operation-type': op.type,
+          'data-operation-json': JSON.stringify(op),
+        },
+        inclusive: false,
+      }).range(from, to)
+    )
+  }
+
+  return Decoration.set(widgets, true)
+}
+
+export function operationsExtension(): Extension {
+  return [operationsStateField]
+}

src/index.css

@@ -143,6 +143,25 @@ code,
 
 #code-mirror-override {
   @apply text-xs;
+
+  /* Operation highlighting CSS. TODO: turn on when we have more rich debug behavior toggling */
+  /* .cm-operation-range { */
+  /* --base-color: oklch(60% .2 15 / .3); */
+
+  /* .dark & { */
+  /* --base-color: oklch(50% .2 15 / .5); */
+  /* } */
+
+  /* background-color: var(--base-color); */
+
+  /* &[data-operation-type="StdLibCall"] { */
+  /* background-color: oklch(from var(--base-color) l c calc(h + 120)); */
+  /* } */
+
+  /* &[data-operation-type="VariableDeclaration"] { */
+  /* background-color: oklch(from var(--base-color) l c calc(h + 240)); */
+  /* } */
+  /* } */
 }
 
 /*

src/lang/KclManager.ts

@@ -100,6 +100,10 @@ import {
   setAstEffect,
   updateAstAnnotation,
 } from '@src/editor/plugins/ast'
+import {
+  operationsAnnotation,
+  setOperationsEffect,
+} from '@src/editor/plugins/operations'
 import { setKclVersion } from '@src/lib/kclVersion'
 import {
   baseEditorExtensions,
@@ -739,6 +743,18 @@ export class KclManager extends EventTarget {
     }
   }
 
+  /**
+   */
+  private dispatchUpdateOperations(newOperations: Operation[]) {
+    this.editorView.dispatch({
+      effects: [setOperationsEffect.of(newOperations)],
+      annotations: [
+        operationsAnnotation.of(true),
+        Transaction.addToHistory.of(false),
+      ],
+    })
+  }
+
   /**
    * Dispatches a CodeMirror state effect to update the AST
    * stored on the current EditorView.
@@ -937,6 +953,8 @@ export class KclManager extends EventTarget {
     this.ast = structuredClone(ast)
     // updateArtifactGraph relies on updated executeState/variables
     await this.updateArtifactGraph(execState.artifactGraph)
+    this.dispatchUpdateOperations(execState.operations)
+
     if (!isInterrupted) {
       this.singletons.sceneInfra.modelingSend({
         type: 'code edit during sketch',