Highlight instanceReference (#2154)

## Checklist

- [-] I have added
[tests](https://www.cursorless.org/docs/contributing/test-case-recorder/)
- [-] I have updated the
[docs](https://github.com/cursorless-dev/cursorless/tree/main/docs) and
[cheatsheet](https://github.com/cursorless-dev/cursorless/tree/main/cursorless-talon/src/cheatsheet)
- [-] I have not broken the cheatsheet

packages/common/src/StoredTargetKey.ts

@@ -0,0 +1,6 @@
+export const storedTargetKeys = [
+  "that",
+  "source",
+  "instanceReference",
+] as const;
+export type StoredTargetKey = (typeof storedTargetKeys)[number];

packages/common/src/index.ts

@@ -99,3 +99,4 @@ export * from "./scopeSupportFacets/scopeSupportFacets.types";
 export * from "./scopeSupportFacets/scopeSupportFacetInfos";
 export * from "./scopeSupportFacets/textualScopeSupportFacetInfos";
 export * from "./scopeSupportFacets/getLanguageScopeSupport";
+export * from "./StoredTargetKey";

packages/cursorless-engine/src/core/StoredTargets.ts

@@ -1,19 +1,31 @@
+import { Notifier } from "@cursorless/common";
 import { Target } from "../typings/target.types";
-
-export type StoredTargetKey = "that" | "source" | "instanceReference";
+import { StoredTargetKey, storedTargetKeys } from "@cursorless/common";
 
 /**
  * Used to store targets between commands.  This is used by marks like `that`
  * and `source`.
  */
 export class StoredTargetMap {
   private targetMap: Map<StoredTargetKey, Target[] | undefined> = new Map();
+  private notifier = new Notifier<[StoredTargetKey, Target[] | undefined]>();
 
   set(key: StoredTargetKey, targets: Target[] | undefined) {
     this.targetMap.set(key, targets);
+    this.notifier.notifyListeners(key, targets);
   }
 
   get(key: StoredTargetKey) {
     return this.targetMap.get(key);
   }
+
+  onStoredTargets(
+    callback: (key: StoredTargetKey, targets: Target[] | undefined) => void,
+  ) {
+    for (const key of storedTargetKeys) {
+      callback(key, this.get(key));
+    }
+
+    return this.notifier.registerListener(callback);
+  }
 }

packages/cursorless-engine/src/processTargets/marks/StoredTargetStage.ts

@@ -1,4 +1,5 @@
-import { StoredTargetKey, StoredTargetMap } from "../../core/StoredTargets";
+import { StoredTargetKey } from "@cursorless/common";
+import { StoredTargetMap } from "../../core/StoredTargets";
 import { Target } from "../../typings/target.types";
 import { MarkStage } from "../PipelineStages.types";
 

packages/cursorless-vscode/src/constructTestHelpers.ts

@@ -7,12 +7,12 @@ import {
   NormalizedIDE,
   ScopeProvider,
   SerializedMarks,
+  StoredTargetKey,
   TargetPlainObject,
   TestCaseSnapshot,
   TextEditor,
 } from "@cursorless/common";
 import {
-  StoredTargetKey,
   StoredTargetMap,
   plainObjectToTarget,
   takeSnapshot,

packages/cursorless-vscode/src/extension.ts

@@ -50,6 +50,7 @@ import {
 } from "./ScopeVisualizerCommandApi";
 import { StatusBarItem } from "./StatusBarItem";
 import { vscodeApi } from "./vscodeApi";
+import { storedTargetHighlighter } from "./storedTargetHighlighter";
 
 /**
  * Extension entrypoint called by VSCode on Cursorless startup.
@@ -127,6 +128,8 @@ export async function activate(
     commandServerApi != null,
   );
 
+  context.subscriptions.push(storedTargetHighlighter(vscodeIDE, storedTargets));
+
   registerCommands(
     context,
     vscodeIDE,

packages/cursorless-vscode/src/storedTargetHighlighter.ts

@@ -0,0 +1,75 @@
+import { StoredTargetKey, groupBy, toCharacterRange } from "@cursorless/common";
+import { StoredTargetMap } from "@cursorless/cursorless-engine";
+import {
+  ScopeRangeType,
+  ScopeVisualizerColorConfig,
+} from "@cursorless/vscode-common";
+import { VscodeIDE } from "./ide/vscode/VscodeIDE";
+import { VscodeFancyRangeHighlighter } from "./ide/vscode/VSCodeScopeVisualizer/VscodeFancyRangeHighlighter";
+import { getColorsFromConfig } from "./ide/vscode/VSCodeScopeVisualizer/getColorsFromConfig";
+import { mapValues } from "lodash";
+import { usingSetting } from "./usingSetting";
+
+const targetColorMap: Partial<Record<StoredTargetKey, ScopeRangeType>> = {
+  instanceReference: "domain",
+};
+
+/**
+ * Constructs the stored target highlighter and listens for changes to stored
+ * targets, highlighting them in the editor.
+ * @param ide The ide object
+ * @param storedTargets Keeps track of stored targets
+ * @returns A disposable that disposes of the stored target highlighter
+ */
+export function storedTargetHighlighter(
+  ide: VscodeIDE,
+  storedTargets: StoredTargetMap,
+) {
+  return usingSetting<ScopeVisualizerColorConfig>(
+    "cursorless.scopeVisualizer",
+    "colors",
+    (colorConfig) => {
+      const highlighters = mapValues(targetColorMap, (type) =>
+        type == null
+          ? undefined
+          : new VscodeFancyRangeHighlighter(
+              getColorsFromConfig(colorConfig, type),
+            ),
+      );
+
+      const storedTargetsDisposable = storedTargets.onStoredTargets(
+        (key, targets) => {
+          const highlighter = highlighters[key];
+
+          if (highlighter == null) {
+            return;
+          }
+
+          const editorRangeMap = groupBy(
+            targets ?? [],
+            ({ editor }) => editor.id,
+          );
+
+          ide.visibleTextEditors.forEach((editor) => {
+            highlighter.setRanges(
+              editor,
+              (editorRangeMap.get(editor.id) ?? []).map(({ contentRange }) =>
+                toCharacterRange(contentRange),
+              ),
+            );
+          });
+        },
+      );
+
+      return {
+        dispose: () => {
+          for (const highlighter of Object.values(highlighters)) {
+            highlighter?.dispose();
+          }
+
+          storedTargetsDisposable.dispose();
+        },
+      };
+    },
+  );
+}

packages/cursorless-vscode/src/usingSetting.ts

@@ -0,0 +1,38 @@
+import { vscodeApi } from "./vscodeApi";
+import { Disposable } from "@cursorless/common";
+
+/**
+ * Watches for changes to a setting and calls a factory function whenever the
+ * setting changes, disposing of any disposables created by the factory. On the
+ * initial call, the factory function is called immediately.
+ *
+ * @param section The section of the setting
+ * @param setting The setting
+ * @param factory A function that takes the setting value and returns a disposable
+ * @returns A disposable that disposes of the setting listener and any disposables created by the factory
+ */
+export function usingSetting<T>(
+  section: string,
+  setting: string,
+  factory: (value: T) => Disposable,
+): Disposable {
+  const runFactoryWithLatestConfig = () =>
+    factory(vscodeApi.workspace.getConfiguration(section).get<T>(setting)!);
+
+  let disposable = runFactoryWithLatestConfig();
+  const configurationDisposable = vscodeApi.workspace.onDidChangeConfiguration(
+    ({ affectsConfiguration }) => {
+      if (affectsConfiguration(`${section}.${setting}`)) {
+        disposable.dispose();
+        disposable = runFactoryWithLatestConfig();
+      }
+    },
+  );
+
+  return {
+    dispose: () => {
+      disposable.dispose();
+      configurationDisposable.dispose();
+    },
+  };
+}