feat(web): find occurrences when idle

9 files changed, 259 insertions, 42 deletions

diff --git a/language-web/src/main/js/editor/EditorParent.ts b/language-web/src/main/js/editor/EditorParent.ts
index ea8c13b6..b890ac3c 100644
--- a/language-web/src/main/js/editor/EditorParent.ts
+++ b/language-web/src/main/js/editor/EditorParent.ts
@@ -190,5 +190,11 @@ export const EditorParent = styled('div')(({ theme }) => {
       textAlign: 'center',
     },
     ...codeMirrorLintStyle,
+    '.cm-problem-write': {
+      background: 'rgba(255, 255, 128, 0.3)',
+    },
+    '.cm-problem-read': {
+      background: 'rgba(255, 255, 255, 0.15)',
+    },
   };
 });
diff --git a/language-web/src/main/js/editor/EditorStore.ts b/language-web/src/main/js/editor/EditorStore.ts
index f47f47a0..8788e00f 100644
--- a/language-web/src/main/js/editor/EditorStore.ts
+++ b/language-web/src/main/js/editor/EditorStore.ts
@@ -30,7 +30,6 @@ import {
   TransactionSpec,
 } from '@codemirror/state';
 import {
-  DecorationSet,
   drawSelection,
   EditorView,
   highlightActiveLine,
@@ -43,8 +42,13 @@ import {
   reaction,
 } from 'mobx';
 
+import { findOccurrences, IOccurrence, setOccurrences } from './findOccurrences';
 import { problemLanguageSupport } from '../language/problemLanguageSupport';
-import { semanticHighlighting, setSemanticHighlighting } from './semanticHighlighting';
+import {
+  IHighlightRange,
+  semanticHighlighting,
+  setSemanticHighlighting,
+} from './semanticHighlighting';
 import type { ThemeStore } from '../theme/ThemeStore';
 import { getLogger } from '../utils/logger';
 import { XtextClient } from '../xtext/XtextClient';
@@ -95,6 +99,7 @@ export class EditorStore {
         EditorView.theme({}, {
           dark: this.themeStore.darkMode,
         }),
+        findOccurrences,
         highlightActiveLine(),
         highlightActiveLineGutter(),
         highlightSpecialChars(),
@@ -204,8 +209,12 @@ export class EditorStore {
     return null;
   }
 
-  updateSemanticHighlighting(decorations: DecorationSet): void {
-    this.dispatch(setSemanticHighlighting(decorations));
+  updateSemanticHighlighting(ranges: IHighlightRange[]): void {
+    this.dispatch(setSemanticHighlighting(ranges));
+  }
+
+  updateOccurrences(write: IOccurrence[], read: IOccurrence[]): void {
+    this.dispatch(setOccurrences(write, read));
   }
 
   /**
diff --git a/language-web/src/main/js/editor/decorationSetExtension.ts b/language-web/src/main/js/editor/decorationSetExtension.ts
new file mode 100644
index 00000000..2d630c20
--- /dev/null
+++ b/language-web/src/main/js/editor/decorationSetExtension.ts
@@ -0,0 +1,39 @@
+import { StateEffect, StateField, TransactionSpec } from '@codemirror/state';
+import { EditorView, Decoration, DecorationSet } from '@codemirror/view';
+
+export type TransactionSpecFactory = (decorations: DecorationSet) => TransactionSpec;
+
+export function decorationSetExtension(): [TransactionSpecFactory, StateField<DecorationSet>] {
+  const setEffect = StateEffect.define<DecorationSet>();
+  const field = StateField.define<DecorationSet>({
+    create() {
+      return Decoration.none;
+    },
+    update(currentDecorations, transaction) {
+      let newDecorations: DecorationSet | null = null;
+      transaction.effects.forEach((effect) => {
+        if (effect.is(setEffect)) {
+          newDecorations = effect.value;
+        }
+      });
+      if (newDecorations === null) {
+        if (transaction.docChanged) {
+          return currentDecorations.map(transaction.changes);
+        }
+        return currentDecorations;
+      }
+      return newDecorations;
+    },
+    provide: (f) => EditorView.decorations.from(f),
+  });
+
+  function transactionSpecFactory(decorations: DecorationSet) {
+    return {
+      effects: [
+        setEffect.of(decorations),
+      ],
+    };
+  }
+
+  return [transactionSpecFactory, field];
+}
diff --git a/language-web/src/main/js/editor/findOccurrences.ts b/language-web/src/main/js/editor/findOccurrences.ts
new file mode 100644
index 00000000..92102746
--- /dev/null
+++ b/language-web/src/main/js/editor/findOccurrences.ts
@@ -0,0 +1,35 @@
+import { Range, RangeSet } from '@codemirror/rangeset';
+import type { TransactionSpec } from '@codemirror/state';
+import { Decoration } from '@codemirror/view';
+
+import { decorationSetExtension } from './decorationSetExtension';
+
+export interface IOccurrence {
+  from: number;
+
+  to: number;
+}
+
+const [setOccurrencesInteral, findOccurrences] = decorationSetExtension();
+
+const writeDecoration = Decoration.mark({
+  class: 'cm-problem-write',
+});
+
+const readDecoration = Decoration.mark({
+  class: 'cm-problem-read',
+});
+
+export function setOccurrences(write: IOccurrence[], read: IOccurrence[]): TransactionSpec {
+  const decorations: Range<Decoration>[] = [];
+  write.forEach(({ from, to }) => {
+    decorations.push(writeDecoration.range(from, to));
+  });
+  read.forEach(({ from, to }) => {
+    decorations.push(readDecoration.range(from, to));
+  });
+  const rangeSet = RangeSet.of(decorations, true);
+  return setOccurrencesInteral(rangeSet);
+}
+
+export { findOccurrences };
diff --git a/language-web/src/main/js/editor/semanticHighlighting.ts b/language-web/src/main/js/editor/semanticHighlighting.ts
index 2d6804f8..2aed421b 100644
--- a/language-web/src/main/js/editor/semanticHighlighting.ts
+++ b/language-web/src/main/js/editor/semanticHighlighting.ts
@@ -1,34 +1,24 @@
-import { StateEffect, StateField, TransactionSpec } from '@codemirror/state';
-import { EditorView, Decoration, DecorationSet } from '@codemirror/view';
+import { RangeSet } from '@codemirror/rangeset';
+import type { TransactionSpec } from '@codemirror/state';
+import { Decoration } from '@codemirror/view';
 
-const setSemanticHighlightingEffect = StateEffect.define<DecorationSet>();
+import { decorationSetExtension } from './decorationSetExtension';
 
-export function setSemanticHighlighting(decorations: DecorationSet): TransactionSpec {
-  return {
-    effects: [
-      setSemanticHighlightingEffect.of(decorations),
-    ],
-  };
+export interface IHighlightRange {
+  from: number;
+
+  to: number;
+
+  classes: string[];
+}
+
+const [setSemanticHighlightingInternal, semanticHighlighting] = decorationSetExtension();
+
+export function setSemanticHighlighting(ranges: IHighlightRange[]): TransactionSpec {
+  const rangeSet = RangeSet.of(ranges.map(({ from, to, classes }) => Decoration.mark({
+    class: classes.map((c) => `cmt-problem-${c}`).join(' '),
+  }).range(from, to)), true);
+  return setSemanticHighlightingInternal(rangeSet);
 }
 
-export const semanticHighlighting = StateField.define<DecorationSet>({
-  create() {
-    return Decoration.none;
-  },
-  update(currentDecorations, transaction) {
-    let newDecorations: DecorationSet | null = null;
-    transaction.effects.forEach((effect) => {
-      if (effect.is(setSemanticHighlightingEffect)) {
-        newDecorations = effect.value;
-      }
-    });
-    if (newDecorations === null) {
-      if (transaction.docChanged) {
-        return currentDecorations.map(transaction.changes);
-      }
-      return currentDecorations;
-    }
-    return newDecorations;
-  },
-  provide: (f) => EditorView.decorations.from(f),
-});
+export { semanticHighlighting };
diff --git a/language-web/src/main/js/xtext/HighlightingService.ts b/language-web/src/main/js/xtext/HighlightingService.ts
index b8ceed20..451a3a52 100644
--- a/language-web/src/main/js/xtext/HighlightingService.ts
+++ b/language-web/src/main/js/xtext/HighlightingService.ts
@@ -1,7 +1,5 @@
-import { Decoration } from '@codemirror/view';
-import { Range, RangeSet } from '@codemirror/rangeset';
-
 import type { EditorStore } from '../editor/EditorStore';
+import type { IHighlightRange } from '../editor/semanticHighlighting';
 import type { UpdateService } from './UpdateService';
 import { getLogger } from '../utils/logger';
 import { isHighlightingResult } from './xtextServiceResults';
@@ -24,7 +22,7 @@ export class HighlightingService {
       return;
     }
     const allChanges = this.updateService.computeChangesSinceLastUpdate();
-    const decorations: Range<Decoration>[] = [];
+    const ranges: IHighlightRange[] = [];
     push.regions.forEach(({ offset, length, styleClasses }) => {
       if (styleClasses.length === 0) {
         return;
@@ -34,10 +32,12 @@ export class HighlightingService {
       if (to <= from) {
         return;
       }
-      decorations.push(Decoration.mark({
-        class: styleClasses.map((c) => `cmt-problem-${c}`).join(' '),
-      }).range(from, to));
+      ranges.push({
+        from,
+        to,
+        classes: styleClasses,
+      });
     });
-    this.store.updateSemanticHighlighting(RangeSet.of(decorations, true));
+    this.store.updateSemanticHighlighting(ranges);
   }
 }
diff --git a/language-web/src/main/js/xtext/OccurrencesService.ts b/language-web/src/main/js/xtext/OccurrencesService.ts
new file mode 100644
index 00000000..804f5ba2
--- /dev/null
+++ b/language-web/src/main/js/xtext/OccurrencesService.ts
@@ -0,0 +1,116 @@
+import { Transaction } from '@codemirror/state';
+
+import type { EditorStore } from '../editor/EditorStore';
+import type { IOccurrence } from '../editor/findOccurrences';
+import type { UpdateService } from './UpdateService';
+import { getLogger } from '../utils/logger';
+import { Timer } from '../utils/Timer';
+import { XtextWebSocketClient } from './XtextWebSocketClient';
+import {
+  isOccurrencesResult,
+  isServiceConflictResult,
+  ITextRegion,
+} from './xtextServiceResults';
+
+const FIND_OCCURRENCES_TIMEOUT_MS = 1000;
+
+// Must clear occurrences asynchronously from `onTransaction`,
+// because we must not emit a conflicting transaction when handling the pending transaction.
+const CLEAR_OCCURRENCES_TIMEOUT_MS = 10;
+
+const log = getLogger('xtext.OccurrencesService');
+
+function transformOccurrences(regions: ITextRegion[]): IOccurrence[] {
+  const occurrences: IOccurrence[] = [];
+  regions.forEach(({ offset, length }) => {
+    if (length > 0) {
+      occurrences.push({
+        from: offset,
+        to: offset + length,
+      });
+    }
+  });
+  return occurrences;
+}
+
+export class OccurrencesService {
+  private store: EditorStore;
+
+  private webSocketClient: XtextWebSocketClient;
+
+  private updateService: UpdateService;
+
+  private hasOccurrences = false;
+
+  private findOccurrencesTimer = new Timer(() => {
+    this.handleFindOccurrences();
+  }, FIND_OCCURRENCES_TIMEOUT_MS);
+
+  private clearOccurrencesTimer = new Timer(() => {
+    this.clearOccurrences();
+  }, CLEAR_OCCURRENCES_TIMEOUT_MS);
+
+  constructor(
+    store: EditorStore,
+    webSocketClient: XtextWebSocketClient,
+    updateService: UpdateService,
+  ) {
+    this.store = store;
+    this.webSocketClient = webSocketClient;
+    this.updateService = updateService;
+  }
+
+  onTransaction(transaction: Transaction): void {
+    if (transaction.docChanged) {
+      this.clearOccurrencesTimer.schedule();
+      this.findOccurrencesTimer.reschedule();
+    }
+    if (transaction.isUserEvent('select')) {
+      this.findOccurrencesTimer.reschedule();
+    }
+  }
+
+  private handleFindOccurrences() {
+    this.clearOccurrencesTimer.cancel();
+    this.updateOccurrences().catch((error) => {
+      log.error('Unexpected error while updating occurrences', error);
+      this.clearOccurrences();
+    });
+  }
+
+  private async updateOccurrences() {
+    await this.updateService.update();
+    const result = await this.webSocketClient.send({
+      resource: this.updateService.resourceName,
+      serviceType: 'occurrences',
+      expectedStateId: this.updateService.xtextStateId,
+      caretOffset: this.store.state.selection.main.head,
+    });
+    const allChanges = this.updateService.computeChangesSinceLastUpdate();
+    if (!allChanges.empty
+      || (isServiceConflictResult(result) && result.conflict === 'canceled')) {
+      // Stale occurrences result, the user already made some changes.
+      // We can safely ignore the occurrences and schedule a new find occurrences call.
+      this.clearOccurrences();
+      this.findOccurrencesTimer.schedule();
+      return;
+    }
+    if (!isOccurrencesResult(result) || result.stateId !== this.updateService.xtextStateId) {
+      log.error('Unexpected occurrences result', result);
+      this.clearOccurrences();
+      return;
+    }
+    const write = transformOccurrences(result.writeRegions);
+    const read = transformOccurrences(result.readRegions);
+    this.hasOccurrences = write.length > 0 || read.length > 0;
+    log.debug('Found', write.length, 'write and', read.length, 'read occurrences');
+    this.store.updateOccurrences(write, read);
+  }
+
+  private clearOccurrences() {
+    if (this.hasOccurrences) {
+      this.store.updateOccurrences([], []);
+      this.hasOccurrences = false;
+    }
+  }
+}
diff --git a/language-web/src/main/js/xtext/XtextClient.ts b/language-web/src/main/js/xtext/XtextClient.ts
index ccb58ab4..03b81b1c 100644
--- a/language-web/src/main/js/xtext/XtextClient.ts
+++ b/language-web/src/main/js/xtext/XtextClient.ts
@@ -7,6 +7,7 @@ import type { Transaction } from '@codemirror/state';
 import type { EditorStore } from '../editor/EditorStore';
 import { ContentAssistService } from './ContentAssistService';
 import { HighlightingService } from './HighlightingService';
+import { OccurrencesService } from './OccurrencesService';
 import { UpdateService } from './UpdateService';
 import { getLogger } from '../utils/logger';
 import { ValidationService } from './ValidationService';
@@ -25,6 +26,8 @@ export class XtextClient {
 
   private validationService: ValidationService;
 
+  private occurrencesService: OccurrencesService;
+
   constructor(store: EditorStore) {
     this.webSocketClient = new XtextWebSocketClient(
       () => this.updateService.onConnect(),
@@ -34,6 +37,11 @@ export class XtextClient {
     this.contentAssistService = new ContentAssistService(this.updateService);
     this.highlightingService = new HighlightingService(store, this.updateService);
     this.validationService = new ValidationService(store, this.updateService);
+    this.occurrencesService = new OccurrencesService(
+      store,
+      this.webSocketClient,
+      this.updateService,
+    );
   }
 
   onTransaction(transaction: Transaction): void {
@@ -41,6 +49,7 @@ export class XtextClient {
     // _before_ the current edit, so we call it before `updateService`.
     this.contentAssistService.onTransaction(transaction);
     this.updateService.onTransaction(transaction);
+    this.occurrencesService.onTransaction(transaction);
   }
 
   private async onPush(resource: string, stateId: string, service: string, push: unknown) {
diff --git a/language-web/src/main/js/xtext/xtextServiceResults.ts b/language-web/src/main/js/xtext/xtextServiceResults.ts
index e32d49c3..b2de1e4a 100644
--- a/language-web/src/main/js/xtext/xtextServiceResults.ts
+++ b/language-web/src/main/js/xtext/xtextServiceResults.ts
@@ -224,3 +224,16 @@ export function isHighlightingResult(result: unknown): result is IHighlightingRe
   return typeof highlightingResult === 'object'
     && isArrayOfType(highlightingResult.regions, isHighlightingRegion);
 }
+
+export interface IOccurrencesResult extends IDocumentStateResult {
+  writeRegions: ITextRegion[];
+
+  readRegions: ITextRegion[];
+}
+
+export function isOccurrencesResult(result: unknown): result is IOccurrencesResult {
+  const occurrencesResult = result as IOccurrencesResult;
+  return isDocumentStateResult(occurrencesResult)
+    && isArrayOfType(occurrencesResult.writeRegions, isTextRegion)
+    && isArrayOfType(occurrencesResult.readRegions, isTextRegion);
+}