refactor(frontend): simplify diagnostic tracking

5 files changed, 98 insertions, 38 deletions

diff --git a/subprojects/frontend/src/editor/DiagnosticValue.ts b/subprojects/frontend/src/editor/DiagnosticValue.ts
index ad23c467..b4e0b165 100644
--- a/subprojects/frontend/src/editor/DiagnosticValue.ts
+++ b/subprojects/frontend/src/editor/DiagnosticValue.ts
@@ -1,14 +1,16 @@
 import type { Diagnostic } from '@codemirror/lint';
 import { RangeValue } from '@codemirror/state';
 
+export type Severity = Diagnostic['severity'];
+
 export default class DiagnosticValue extends RangeValue {
-  static VALUES: Record<Diagnostic['severity'], DiagnosticValue> = {
+  static VALUES: Record<Severity, DiagnosticValue> = {
     error: new DiagnosticValue('error'),
     warning: new DiagnosticValue('warning'),
     info: new DiagnosticValue('info'),
   };
 
-  private constructor(public readonly severity: Diagnostic['severity']) {
+  private constructor(public readonly severity: Severity) {
     super();
   }
 
diff --git a/subprojects/frontend/src/editor/EditorStore.ts b/subprojects/frontend/src/editor/EditorStore.ts
index d966690c..1f301c31 100644
--- a/subprojects/frontend/src/editor/EditorStore.ts
+++ b/subprojects/frontend/src/editor/EditorStore.ts
@@ -14,7 +14,6 @@ import {
   type Transaction,
   type TransactionSpec,
   type EditorState,
-  RangeSet,
 } from '@codemirror/state';
 import { type Command, EditorView } from '@codemirror/view';
 import { makeAutoObservable, observable } from 'mobx';
@@ -24,10 +23,10 @@ import type PWAStore from '../PWAStore';
 import getLogger from '../utils/getLogger';
 import XtextClient from '../xtext/XtextClient';
 
-import DiagnosticValue from './DiagnosticValue';
 import LintPanelStore from './LintPanelStore';
 import SearchPanelStore from './SearchPanelStore';
 import createEditorState from './createEditorState';
+import { countDiagnostics } from './exposeDiagnostics';
 import { type IOccurrence, setOccurrences } from './findOccurrences';
 import {
   type IHighlightRange,
@@ -51,8 +50,6 @@ export default class EditorStore {
 
   showLineNumbers = false;
 
-  diagnostics: RangeSet<DiagnosticValue> = RangeSet.of([]);
-
   constructor(initialValue: string, pwaStore: PWAStore) {
     this.id = nanoid();
     this.state = createEditorState(initialValue, this);
@@ -162,9 +159,6 @@ export default class EditorStore {
     log.trace('Editor transaction', tr);
     this.state = tr.state;
     this.client.onTransaction(tr);
-    if (tr.docChanged) {
-      this.diagnostics = this.diagnostics.map(tr.changes);
-    }
   }
 
   doCommand(command: Command): boolean {
@@ -182,31 +176,19 @@ export default class EditorStore {
   }
 
   updateDiagnostics(diagnostics: Diagnostic[]): void {
-    diagnostics.sort((a, b) => a.from - b.from);
     this.dispatch(setDiagnostics(this.state, diagnostics));
-    this.diagnostics = RangeSet.of(
-      diagnostics.map(({ severity, from, to }) =>
-        DiagnosticValue.VALUES[severity].range(from, to),
-      ),
-    );
-  }
-
-  countDiagnostics(severity: Diagnostic['severity']): number {
-    return this.diagnostics.update({
-      filter: (_from, _to, value) => value.eq(DiagnosticValue.VALUES[severity]),
-    }).size;
   }
 
   get errorCount(): number {
-    return this.countDiagnostics('error');
+    return countDiagnostics(this.state, 'error');
   }
 
   get warningCount(): number {
-    return this.countDiagnostics('warning');
+    return countDiagnostics(this.state, 'warning');
   }
 
   get infoCount(): number {
-    return this.countDiagnostics('info');
+    return countDiagnostics(this.state, 'info');
   }
 
   nextDiagnostic(): void {
diff --git a/subprojects/frontend/src/editor/createEditorState.ts b/subprojects/frontend/src/editor/createEditorState.ts
index 05028fcc..ce1efa4f 100644
--- a/subprojects/frontend/src/editor/createEditorState.ts
+++ b/subprojects/frontend/src/editor/createEditorState.ts
@@ -36,6 +36,7 @@ import problemLanguageSupport from '../language/problemLanguageSupport';
 
 import type EditorStore from './EditorStore';
 import SearchPanel from './SearchPanel';
+import exposeDiagnostics from './exposeDiagnostics';
 import findOccurrences from './findOccurrences';
 import indentationMarkerViewPlugin from './indentationMarkerViewPlugin';
 import scrollbarViewPlugin from './scrollbarViewPlugin';
@@ -56,6 +57,7 @@ export default function createEditorState(
       bracketMatching(),
       drawSelection(),
       EditorState.allowMultipleSelections.of(true),
+      exposeDiagnostics,
       findOccurrences,
       highlightActiveLine(),
       highlightActiveLineGutter(),
diff --git a/subprojects/frontend/src/editor/exposeDiagnostics.ts b/subprojects/frontend/src/editor/exposeDiagnostics.ts
new file mode 100644
index 00000000..82f24c93
--- /dev/null
+++ b/subprojects/frontend/src/editor/exposeDiagnostics.ts
@@ -0,0 +1,80 @@
+import { setDiagnosticsEffect } from '@codemirror/lint';
+import {
+  StateField,
+  RangeSet,
+  type Extension,
+  type EditorState,
+} from '@codemirror/state';
+
+import DiagnosticValue, { type Severity } from './DiagnosticValue';
+
+type SeverityCounts = Partial<Record<Severity, number>>;
+
+interface ExposedDiagnostics {
+  readonly diagnostics: RangeSet<DiagnosticValue>;
+
+  readonly severityCounts: SeverityCounts;
+}
+
+function countSeverities(
+  diagnostics: RangeSet<DiagnosticValue>,
+): SeverityCounts {
+  const severityCounts: SeverityCounts = {};
+  const iter = diagnostics.iter();
+  while (iter.value !== null) {
+    const {
+      value: { severity },
+    } = iter;
+    severityCounts[severity] = (severityCounts[severity] ?? 0) + 1;
+    iter.next();
+  }
+  return severityCounts;
+}
+
+const exposedDiagnosticsState = StateField.define<ExposedDiagnostics>({
+  create() {
+    return {
+      diagnostics: RangeSet.of([]),
+      severityCounts: {},
+    };
+  },
+
+  update({ diagnostics: diagnosticsSet, severityCounts }, transaction) {
+    let newDiagnosticsSet = diagnosticsSet;
+    if (transaction.docChanged) {
+      newDiagnosticsSet = newDiagnosticsSet.map(transaction.changes);
+    }
+    transaction.effects.forEach((effect) => {
+      if (effect.is(setDiagnosticsEffect)) {
+        const diagnostics = effect.value.map(({ severity, from, to }) =>
+          DiagnosticValue.VALUES[severity].range(from, to),
+        );
+        diagnostics.sort(({ from: a }, { from: b }) => a - b);
+        newDiagnosticsSet = RangeSet.of(diagnostics);
+      }
+    });
+    return {
+      diagnostics: newDiagnosticsSet,
+      severityCounts:
+        // Only recompute if the diagnostics were changed.
+        diagnosticsSet === newDiagnosticsSet
+          ? severityCounts
+          : countSeverities(newDiagnosticsSet),
+    };
+  },
+});
+
+const exposeDiagnostics: Extension = [exposedDiagnosticsState];
+
+export default exposeDiagnostics;
+
+export function getDiagnostics(state: EditorState): RangeSet<DiagnosticValue> {
+  return state.field(exposedDiagnosticsState).diagnostics;
+}
+
+export function countDiagnostics(
+  state: EditorState,
+  severity: Severity,
+): number {
+  return state.field(exposedDiagnosticsState).severityCounts[severity] ?? 0;
+}
diff --git a/subprojects/frontend/src/editor/scrollbarViewPlugin.ts b/subprojects/frontend/src/editor/scrollbarViewPlugin.ts
index 0edaeb70..b0ea769f 100644
--- a/subprojects/frontend/src/editor/scrollbarViewPlugin.ts
+++ b/subprojects/frontend/src/editor/scrollbarViewPlugin.ts
@@ -2,6 +2,7 @@ import { type PluginValue, ViewPlugin } from '@codemirror/view';
 import { reaction } from 'mobx';
 
 import type EditorStore from './EditorStore';
+import { getDiagnostics } from './exposeDiagnostics';
 import findOccurrences from './findOccurrences';
 
 export const HOLDER_CLASS = 'cm-scroller-holder';
@@ -105,11 +106,12 @@ export default function scrollbarViewPlugin(
     const annotations: HTMLDivElement[] = [];
 
     function rebuildAnnotations(trackYHeight: number) {
+      const { state } = view;
       const annotationOverlayHeight = Math.min(
         view.contentHeight,
         trackYHeight,
       );
-      const lineHeight = annotationOverlayHeight / editorStore.state.doc.lines;
+      const lineHeight = annotationOverlayHeight / state.doc.lines;
 
       let i = 0;
 
@@ -117,11 +119,9 @@ export default function scrollbarViewPlugin(
         from: number,
         to?: number,
       ): HTMLDivElement {
-        const startLine = editorStore.state.doc.lineAt(from).number;
+        const startLine = state.doc.lineAt(from).number;
         const endLine =
-          to === undefined
-            ? startLine
-            : editorStore.state.doc.lineAt(to).number;
+          to === undefined ? startLine : state.doc.lineAt(to).number;
         const top = (startLine - 1) * lineHeight;
         const height = Math.max(
           MIN_ANNOTATION_HEIGHT,
@@ -145,13 +145,13 @@ export default function scrollbarViewPlugin(
         return annotation;
       }
 
-      editorStore.state.selection.ranges.forEach(({ head }) => {
+      state.selection.ranges.forEach(({ head }) => {
         const selectionAnnotation = getOrCreateAnnotation(head);
         selectionAnnotation.className = ANNOTATION_SELECTION_CLASS;
         selectionAnnotation.style.width = `${SCROLLBAR_WIDTH}px`;
       });
 
-      const diagnosticsIter = editorStore.diagnostics.iter();
+      const diagnosticsIter = getDiagnostics(state).iter();
       while (diagnosticsIter.value !== null) {
         const diagnosticAnnotation = getOrCreateAnnotation(
           diagnosticsIter.from,
@@ -162,7 +162,7 @@ export default function scrollbarViewPlugin(
         diagnosticsIter.next();
       }
 
-      const occurrences = editorStore.state.field(findOccurrences);
+      const occurrences = view.state.field(findOccurrences);
       const occurrencesIter = occurrences.iter();
       while (occurrencesIter.value !== null) {
         const occurrenceAnnotation = getOrCreateAnnotation(
@@ -259,15 +259,9 @@ export default function scrollbarViewPlugin(
 
     requestRebuild();
 
-    const disposeRebuildReaction = reaction(
-      () => editorStore.diagnostics,
-      requestRebuild,
-    );
-
     return {
       update: requestRebuild,
       destroy() {
-        disposeRebuildReaction();
         disposePanelReaction();
         observer?.disconnect();
         scrollDOM.removeEventListener('scroll', requestUpdate);