Improve performance for multiple cursors by adding scope handler cache (#3101)

Expensive textual scopes like surrounding pair and collection item had
unreasonable performance with large files and multiple cursors. I ran
into this with a large json file and Cursorless stopped responding.

This pull request introduces a cache for scope handlers as well as new
tests

Author: AndreasArvidsson

packages/cursorless-engine/src/index.ts

@@ -9,6 +9,8 @@ export * from "./cursorlessEngine";
 export * from "./customCommandGrammar/parseCommand";
 export * from "./generateSpokenForm/defaultSpokenForms/surroundingPairsDelimiters";
 export * from "./generateSpokenForm/generateSpokenForm";
+export * from "./languages/TreeSitterQuery/TreeSitterQueryCache";
+export * from "./processTargets/modifiers/scopeHandlers/ScopeHandlerCache";
 export * from "./singletons/ide.singleton";
 export * from "./spokenForms/defaultSpokenFormMap";
 export * from "./testUtil/extractTargetKeys";

packages/cursorless-engine/src/languages/LanguageDefinitions.ts

@@ -8,7 +8,7 @@ import type {
 import { Notifier, showError } from "@cursorless/common";
 import { toString } from "lodash-es";
 import { LanguageDefinition } from "./LanguageDefinition";
-import { treeSitterQueryCache } from "./TreeSitterQuery/treeSitterQueryCache";
+import { treeSitterQueryCache } from "./TreeSitterQuery/TreeSitterQueryCache";
 
 /**
  * Sentinel value to indicate that a language doesn't have

packages/cursorless-engine/src/languages/TreeSitterQuery/TreeSitterQuery.ts

@@ -16,7 +16,7 @@ import {
   getStartOfEndOfRange,
   rewriteStartOfEndOf,
 } from "./rewriteStartOfEndOf";
-import { treeSitterQueryCache } from "./treeSitterQueryCache";
+import { treeSitterQueryCache } from "./TreeSitterQueryCache";
 
 /**
  * Wrapper around a tree-sitter query that provides a more convenient API, and

packages/cursorless-engine/src/languages/TreeSitterQuery/treeSitterQueryCache.ts renamed to packages/cursorless-engine/src/languages/TreeSitterQuery/TreeSitterQueryCache.ts

@@ -1,7 +1,7 @@
 import type { Position, TextDocument } from "@cursorless/common";
 import type { QueryMatch } from "./QueryCapture";
 
-export class Cache {
+export class TreeSitterQueryCache {
   private documentVersion: number = -1;
   private documentUri: string = "";
   private documentLanguageId: string = "";
@@ -58,4 +58,4 @@ function positionsEqual(a: Position | undefined, b: Position | undefined) {
   return a.isEqual(b);
 }
 
-export const treeSitterQueryCache = new Cache();
+export const treeSitterQueryCache = new TreeSitterQueryCache();

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/CollectionItemScopeHandler/CollectionItemTextualScopeHandler.ts

@@ -13,6 +13,7 @@ import type {
   ComplexScopeType,
   ScopeIteratorRequirements,
 } from "../scopeHandler.types";
+import { scopeHandlerCache } from "../ScopeHandlerCache";
 import type { ScopeHandlerFactory } from "../ScopeHandlerFactory";
 import { isEveryScopeModifier } from "../util/isHintsEveryScope";
 import { OneWayNestedRangeFinder } from "../util/OneWayNestedRangeFinder";
@@ -44,6 +45,25 @@ export class CollectionItemTextualScopeHandler extends BaseScopeHandler {
     hints: ScopeIteratorRequirements,
   ): Iterable<TargetScope> {
     const isEveryScope = isEveryScopeModifier(hints);
+    const cacheKey = "CollectionItemTextualScopeHandler_" + isEveryScope;
+
+    if (!scopeHandlerCache.isValid(cacheKey, editor.document)) {
+      const scopes = this.getsScopes(editor, direction, isEveryScope);
+      scopeHandlerCache.update(cacheKey, editor.document, scopes);
+    }
+
+    const scopes = scopeHandlerCache.get<TargetScope>();
+
+    scopes.sort((a, b) => compareTargetScopes(direction, position, a, b));
+
+    yield* scopes;
+  }
+
+  private getsScopes(
+    editor: TextEditor,
+    direction: Direction,
+    isEveryScope: boolean,
+  ) {
     const separatorRanges = getSeparatorOccurrences(editor.document);
     const interiorRanges = getInteriorRanges(
       this.scopeHandlerFactory,
@@ -134,9 +154,7 @@ export class CollectionItemTextualScopeHandler extends BaseScopeHandler {
       }
     }
 
-    scopes.sort((a, b) => compareTargetScopes(direction, position, a, b));
-
-    yield* scopes;
+    return scopes;
   }
 
   private addScopes(scopes: TargetScope[], state: IterationState) {

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/ScopeHandlerCache.ts

@@ -0,0 +1,40 @@
+import type { TextDocument } from "@cursorless/common";
+
+export class ScopeHandlerCache {
+  private key: string = "";
+  private documentVersion: number = -1;
+  private documentUri: string = "";
+  private documentLanguageId: string = "";
+  private matches: any[] = [];
+
+  clear() {
+    this.key = "";
+    this.documentUri = "";
+    this.documentVersion = -1;
+    this.documentLanguageId = "";
+    this.matches = [];
+  }
+
+  isValid(key: string, document: TextDocument) {
+    return (
+      this.key === key &&
+      this.documentVersion === document.version &&
+      this.documentUri === document.uri.toString() &&
+      this.documentLanguageId === document.languageId
+    );
+  }
+
+  update(key: string, document: TextDocument, matches: any[]) {
+    this.key = key;
+    this.documentVersion = document.version;
+    this.documentUri = document.uri.toString();
+    this.documentLanguageId = document.languageId;
+    this.matches = matches;
+  }
+
+  get<T>(): T[] {
+    return this.matches;
+  }
+}
+
+export const scopeHandlerCache = new ScopeHandlerCache();

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/SurroundingPairScopeHandler/SurroundingPairScopeHandler.ts

@@ -19,6 +19,7 @@ import { getDelimiterOccurrences } from "./getDelimiterOccurrences";
 import { getIndividualDelimiters } from "./getIndividualDelimiters";
 import { getSurroundingPairOccurrences } from "./getSurroundingPairOccurrences";
 import type { SurroundingPairOccurrence } from "./types";
+import { scopeHandlerCache } from "../ScopeHandlerCache";
 
 export class SurroundingPairScopeHandler extends BaseScopeHandler {
   public readonly iterationScopeType: ConditionalScopeType = {
@@ -52,22 +53,31 @@ export class SurroundingPairScopeHandler extends BaseScopeHandler {
       return;
     }
 
-    const delimiterOccurrences = getDelimiterOccurrences(
-      this.languageDefinitions.get(this.languageId),
-      editor.document,
-      getIndividualDelimiters(this.scopeType.delimiter, this.languageId),
-    );
+    const cacheKey = "SurroundingPairScopeHandler_" + this.scopeType.delimiter;
+
+    if (!scopeHandlerCache.isValid(cacheKey, editor.document)) {
+      const delimiterOccurrences = getDelimiterOccurrences(
+        this.languageDefinitions.get(this.languageId),
+        editor.document,
+        getIndividualDelimiters(this.scopeType.delimiter, this.languageId),
+      );
+
+      const surroundingPairs =
+        getSurroundingPairOccurrences(delimiterOccurrences);
+
+      scopeHandlerCache.update(cacheKey, editor.document, surroundingPairs);
+    }
 
-    let surroundingPairs = getSurroundingPairOccurrences(delimiterOccurrences);
+    const surroundingPairs = scopeHandlerCache.get<SurroundingPairOccurrence>();
 
-    surroundingPairs = maybeApplyEmptyTargetHack(
+    const updatedSurroundingPairs = maybeApplyEmptyTargetHack(
       direction,
       hints,
       position,
       surroundingPairs,
     );
 
-    yield* surroundingPairs
+    yield* updatedSurroundingPairs
       .map((pair) =>
         createTargetScope(
           editor,

packages/cursorless-vscode-e2e/src/suite/performance.vscode.test.ts

@@ -5,7 +5,7 @@ import {
   type ScopeType,
   type SimpleScopeTypeType,
 } from "@cursorless/common";
-import { openNewEditor, runCursorlessCommand } from "@cursorless/vscode-common";
+import { openNewEditor, runCursorlessAction } from "@cursorless/vscode-common";
 import assert from "assert";
 import * as vscode from "vscode";
 import { endToEndTestSetup } from "../endToEndTestSetup";
@@ -51,7 +51,7 @@ suite("Performance", async function () {
     ["paragraph", smallThresholdMs],
     ["document", smallThresholdMs],
     ["nonWhitespaceSequence", smallThresholdMs],
-    // Parse tree based, containing/every scope
+    // Parse tree based, containing / every scope
     ["string", smallThresholdMs],
     ["map", smallThresholdMs],
     ["collectionKey", smallThresholdMs],
@@ -65,9 +65,11 @@ suite("Performance", async function () {
     ["boundedParagraph", largeThresholdMs],
     ["boundedNonWhitespaceSequence", largeThresholdMs],
     ["collectionItem", largeThresholdMs],
+    ["collectionItem", largeThresholdMs, "every"],
+    ["collectionItem", largeThresholdMs, "previous"],
     // Surrounding pair
-    [{ type: "surroundingPair", delimiter: "any" }, largeThresholdMs],
     [{ type: "surroundingPair", delimiter: "curlyBrackets" }, largeThresholdMs],
+    [{ type: "surroundingPair", delimiter: "any" }, largeThresholdMs],
     [{ type: "surroundingPair", delimiter: "any" }, largeThresholdMs, "every"],
     [
       { type: "surroundingPair", delimiter: "any" },
@@ -86,10 +88,29 @@ suite("Performance", async function () {
       asyncSafety(() => selectScopeType(scopeType, threshold, modifierType)),
     );
   }
+
+  test(
+    "Select surroundingPair with multiple cursors",
+    asyncSafety(() =>
+      selectWithMultipleCursors(largeThresholdMs, {
+        type: "surroundingPair",
+        delimiter: "any",
+      }),
+    ),
+  );
+
+  test(
+    "Select collectionItem with multiple cursors",
+    asyncSafety(() =>
+      selectWithMultipleCursors(largeThresholdMs, {
+        type: "collectionItem",
+      }),
+    ),
+  );
 });
 
-async function removeToken(thresholdMs: number) {
-  await testPerformance(thresholdMs, {
+function removeToken(thresholdMs: number) {
+  return testPerformance(thresholdMs, {
     name: "remove",
     target: {
       type: "primitive",
@@ -98,12 +119,36 @@ async function removeToken(thresholdMs: number) {
   });
 }
 
-async function selectScopeType(
+function selectWithMultipleCursors(thresholdMs: number, scopeType: ScopeType) {
+  return testPerformanceCallback(
+    thresholdMs,
+    () => {
+      return runCursorlessAction({
+        name: "setSelectionBefore",
+        target: {
+          type: "primitive",
+          modifiers: [getModifier({ type: "collectionItem" }, "every")],
+        },
+      });
+    },
+    () => {
+      return runCursorlessAction({
+        name: "setSelection",
+        target: {
+          type: "primitive",
+          modifiers: [getModifier(scopeType)],
+        },
+      });
+    },
+  );
+}
+
+function selectScopeType(
   scopeType: ScopeType,
   thresholdMs: number,
   modifierType?: ModifierType,
 ) {
-  await testPerformance(thresholdMs, {
+  return testPerformance(thresholdMs, {
     name: "setSelection",
     target: {
       type: "primitive",
@@ -112,41 +157,31 @@ async function selectScopeType(
   });
 }
 
-function getModifier(
-  scopeType: ScopeType,
-  modifierType: ModifierType = "containing",
-): Modifier {
-  switch (modifierType) {
-    case "containing":
-      return { type: "containingScope", scopeType };
-    case "every":
-      return { type: "everyScope", scopeType };
-    case "previous":
-      return {
-        type: "relativeScope",
-        direction: "backward",
-        offset: 1,
-        length: 1,
-        scopeType,
-      };
-  }
+function testPerformance(thresholdMs: number, action: ActionDescriptor) {
+  return testPerformanceCallback(thresholdMs, () => {
+    return runCursorlessAction(action);
+  });
 }
 
-async function testPerformance(thresholdMs: number, action: ActionDescriptor) {
+async function testPerformanceCallback(
+  thresholdMs: number,
+  callback: () => Promise<unknown>,
+  beforeCallback?: () => Promise<unknown>,
+) {
   const editor = await openNewEditor(testData, { languageId: "json" });
   // This is the position of the last json key in the document
   const position = new vscode.Position(editor.document.lineCount - 3, 5);
   const selection = new vscode.Selection(position, position);
   editor.selections = [selection];
   editor.revealRange(selection);
 
+  if (beforeCallback != null) {
+    await beforeCallback();
+  }
+
   const start = performance.now();
 
-  await runCursorlessCommand({
-    version: 7,
-    usePrePhraseSnapshot: false,
-    action,
-  });
+  await callback();
 
   const duration = Math.round(performance.now() - start);
 
@@ -158,6 +193,26 @@ async function testPerformance(thresholdMs: number, action: ActionDescriptor) {
   );
 }
 
+function getModifier(
+  scopeType: ScopeType,
+  modifierType: ModifierType = "containing",
+): Modifier {
+  switch (modifierType) {
+    case "containing":
+      return { type: "containingScope", scopeType };
+    case "every":
+      return { type: "everyScope", scopeType };
+    case "previous":
+      return {
+        type: "relativeScope",
+        direction: "backward",
+        offset: 1,
+        length: 1,
+        scopeType,
+      };
+  }
+}
+
 function getScopeTypeAndTitle(
   scope: SimpleScopeTypeType | ScopeType,
 ): [ScopeType, string] {

packages/cursorless-vscode/src/constructTestHelpers.ts

@@ -13,13 +13,17 @@ import type {
   TextEditor,
 } from "@cursorless/common";
 import type { StoredTargetMap } from "@cursorless/cursorless-engine";
-import { plainObjectToTarget } from "@cursorless/cursorless-engine";
+import {
+  plainObjectToTarget,
+  scopeHandlerCache,
+  treeSitterQueryCache,
+} from "@cursorless/cursorless-engine";
+import { takeSnapshot } from "@cursorless/test-case-recorder";
 import type { VscodeTestHelpers } from "@cursorless/vscode-common";
 import type * as vscode from "vscode";
-import { takeSnapshot } from "@cursorless/test-case-recorder";
+import { toVscodeEditor } from "./ide/vscode/toVscodeEditor";
 import type { VscodeFileSystem } from "./ide/vscode/VscodeFileSystem";
 import type { VscodeIDE } from "./ide/vscode/VscodeIDE";
-import { toVscodeEditor } from "./ide/vscode/toVscodeEditor";
 import { vscodeApi } from "./vscodeApi";
 import type { VscodeTutorial } from "./VscodeTutorial";
 
@@ -45,6 +49,11 @@ export function constructTestHelpers(
       return vscodeIDE.fromVscodeEditor(editor);
     },
 
+    clearCache() {
+      scopeHandlerCache.clear();
+      treeSitterQueryCache.clear();
+    },
+
     // FIXME: Remove this once we have a better way to get this function
     // accessible from our tests
     takeSnapshot(