Fix surrounding pair performance

data/fixtures/recorded/surroundingPair/textual/changePair2.yml

@@ -0,0 +1,23 @@
+languageId: typescript
+command:
+  version: 7
+  spokenForm: change pair
+  action:
+    name: clearAndSetSelection
+    target:
+      type: primitive
+      modifiers:
+        - type: containingScope
+          scopeType: {type: surroundingPair, delimiter: any}
+  usePrePhraseSnapshot: true
+initialState:
+  documentContents: "[1, ']', 2]"
+  selections:
+    - anchor: {line: 0, character: 1}
+      active: {line: 0, character: 1}
+  marks: {}
+finalState:
+  documentContents: ""
+  selections:
+    - anchor: {line: 0, character: 0}
+      active: {line: 0, character: 0}

data/fixtures/recorded/surroundingPair/textual/changePair3.yml

@@ -0,0 +1,23 @@
+languageId: typescript
+command:
+  version: 7
+  spokenForm: change pair
+  action:
+    name: clearAndSetSelection
+    target:
+      type: primitive
+      modifiers:
+        - type: containingScope
+          scopeType: {type: surroundingPair, delimiter: any}
+  usePrePhraseSnapshot: true
+initialState:
+  documentContents: "`[1, ${']'}, 3]`"
+  selections:
+    - anchor: {line: 0, character: 1}
+      active: {line: 0, character: 1}
+  marks: {}
+finalState:
+  documentContents: "``"
+  selections:
+    - anchor: {line: 0, character: 1}
+      active: {line: 0, character: 1}

packages/common/src/index.ts

@@ -76,6 +76,7 @@ export * from "./types/Selection";
 export * from "./types/snippet.types";
 export * from "./types/SpokenForm";
 export * from "./types/SpokenFormType";
+export * from "./types/StringRecord";
 export * from "./types/TalonSpokenForms";
 export * from "./types/TestCaseFixture";
 export * from "./types/TestHelpers";

packages/common/src/util/regex.ts

@@ -30,16 +30,20 @@ function makeCache<T, U>(func: (arg: T) => U) {
 export const rightAnchored = makeCache(_rightAnchored);
 export const leftAnchored = makeCache(_leftAnchored);
 
+export function matchAllIterator(text: string, regex: RegExp) {
+  // Reset the regex to start at the beginning of string, in case the regex has
+  // been used before.
+  // See https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/RegExp/exec#finding_successive_matches
+  regex.lastIndex = 0;
+  return text.matchAll(regex);
+}
+
 export function matchAll<T>(
   text: string,
   regex: RegExp,
   mapfn: (v: RegExpMatchArray, k: number) => T,
 ) {
-  // Reset the regex to start at the beginning of string, in case the regex has
-  // been used before.
-  // See https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/RegExp/exec#finding_successive_matches
-  regex.lastIndex = 0;
-  return Array.from(text.matchAll(regex), mapfn);
+  return Array.from(matchAllIterator(text, regex), mapfn);
 }
 
 export function testRegex(regex: RegExp, text: string): boolean {

packages/cursorless-engine/src/disabledComponents/DisabledLanguageDefinitions.ts

@@ -16,6 +16,10 @@ export class DisabledLanguageDefinitions implements LanguageDefinitions {
     return undefined;
   }
 
+  clearCache(): void {
+    // Do nothing
+  }
+
   getNodeAtLocation(
     _document: TextDocument,
     _range: Range,

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

@@ -3,6 +3,7 @@ import type {
   ScopeType,
   SimpleScopeType,
   SimpleScopeTypeType,
+  StringRecord,
   TreeSitter,
 } from "@cursorless/common";
 import {
@@ -12,6 +13,7 @@ import {
   type TextDocument,
 } from "@cursorless/common";
 import { TreeSitterScopeHandler } from "../processTargets/modifiers/scopeHandlers";
+import { LanguageDefinitionCache } from "./LanguageDefinitionCache";
 import { TreeSitterQuery } from "./TreeSitterQuery";
 import type { QueryCapture } from "./TreeSitterQuery/QueryCapture";
 import { validateQueryCaptures } from "./TreeSitterQuery/validateQueryCaptures";
@@ -21,14 +23,18 @@ import { validateQueryCaptures } from "./TreeSitterQuery/validateQueryCaptures";
  * tree-sitter query used to extract scopes for the given language
  */
 export class LanguageDefinition {
+  private cache: LanguageDefinitionCache;
+
   private constructor(
     /**
      * The tree-sitter query used to extract scopes for the given language.
      * Note that this query contains patterns for all scope types that the
      * language supports using new-style tree-sitter queries
      */
     private query: TreeSitterQuery,
-  ) {}
+  ) {
+    this.cache = new LanguageDefinitionCache();
+  }
 
   /**
    * Construct a language definition for the given language id, if the language
@@ -93,10 +99,34 @@ export class LanguageDefinition {
     document: TextDocument,
     captureName: SimpleScopeTypeType,
   ): QueryCapture[] {
-    return this.query
-      .matches(document)
-      .map((match) => match.captures.find(({ name }) => name === captureName))
-      .filter((capture) => capture != null);
+    if (!this.cache.isValid(document)) {
+      this.cache.update(document, this.getCapturesMap(document));
+    }
+
+    return this.cache.get(captureName);
+  }
+
+  clearCache(): void {
+    this.cache.clear();
+  }
+
+  /**
+   * This is a low level function that returns a map of all captures in the document.
+   */
+  private getCapturesMap(document: TextDocument): StringRecord<QueryCapture[]> {
+    const matches = this.query.matches(document);
+    const result: StringRecord<QueryCapture[]> = {};
+
+    for (const match of matches) {
+      for (const capture of match.captures) {
+        if (result[capture.name] == null) {
+          result[capture.name] = [];
+        }
+        result[capture.name]!.push(capture);
+      }
+    }
+
+    return result;
   }
 }
 

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

@@ -0,0 +1,35 @@
+import type {
+  SimpleScopeTypeType,
+  StringRecord,
+  TextDocument,
+} from "@cursorless/common";
+import type { QueryCapture } from "./TreeSitterQuery/QueryCapture";
+
+export class LanguageDefinitionCache {
+  private documentUri: string = "";
+  private documentVersion: number = -1;
+  private captures: StringRecord<QueryCapture[]> = {};
+
+  clear() {
+    this.documentUri = "";
+    this.documentVersion = -1;
+    this.captures = {};
+  }
+
+  isValid(document: TextDocument) {
+    return (
+      this.documentUri === document.uri.toString() &&
+      this.documentVersion === document.version
+    );
+  }
+
+  update(document: TextDocument, captures: StringRecord<QueryCapture[]>) {
+    this.documentUri = document.uri.toString();
+    this.documentVersion = document.version;
+    this.captures = captures;
+  }
+
+  get(captureName: SimpleScopeTypeType): QueryCapture[] {
+    return this.captures[captureName] ?? [];
+  }
+}

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

@@ -32,6 +32,11 @@ export interface LanguageDefinitions {
    */
   get(languageId: string): LanguageDefinition | undefined;
 
+  /**
+   * Clear the cache of all language definitions. This is run at the start of a command.
+   */
+  clearCache(): void;
+
   /**
    * @deprecated Only for use in legacy containing scope stage
    */
@@ -155,6 +160,14 @@ export class LanguageDefinitionsImpl
     return definition === LANGUAGE_UNDEFINED ? undefined : definition;
   }
 
+  clearCache(): void {
+    for (const definition of this.languageDefinitions.values()) {
+      if (definition !== LANGUAGE_UNDEFINED) {
+        definition.clearCache();
+      }
+    }
+  }
+
   public getNodeAtLocation(document: TextDocument, range: Range): SyntaxNode {
     return this.treeSitter.getNodeAtLocation(document, range);
   }

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

@@ -0,0 +1,49 @@
+import type { Range } from "@cursorless/common";
+
+/**
+ * An iterator that allows for efficient lookup of ranges that contain a search item.
+ * The items must be sorted in document order.
+ */
+export class RangeLookupList<T extends { range: Range }> {
+  private index = 0;
+
+  /**
+   * @param items The items to search in. Must be sorted in document order.
+   */
+  constructor(private items: T[]) {}
+
+  add(item: T) {
+    this.items.push(item);
+  }
+
+  contains(searchItem: Range): boolean {
+    return this.advance(searchItem);
+  }
+
+  getContaining(searchItem: Range): T | undefined {
+    if (this.advance(searchItem)) {
+      return this.items[this.index];
+    }
+
+    return undefined;
+  }
+
+  private advance(searchItem: Range): boolean {
+    while (this.index < this.items.length) {
+      const range = this.items[this.index].range;
+
+      if (range.contains(searchItem)) {
+        return true;
+      }
+
+      // Search item is before the range. Since the ranges are sorted, we can stop here.
+      if (searchItem.end.isBeforeOrEqual(range.start)) {
+        return false;
+      }
+
+      this.index++;
+    }
+
+    return false;
+  }
+}

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

@@ -0,0 +1,73 @@
+import type { Range } from "@cursorless/common";
+import { RangeLookupList } from "./RangeLookupList";
+
+/**
+ * A tree of ranges that allows for efficient lookup of ranges that contain a search item.
+ * The items must be sorted in document order.
+ */
+export class RangeLookupTree<T extends { range: Range }> {
+  private children: RangeLookupList<RangeLookupTreeNode<T>>;
+
+  /**
+   * @param items The items to search in. Must be sorted in document order.
+   */
+  constructor(items: T[]) {
+    this.children = createNodes(items);
+  }
+
+  getSmallLestContaining(separator: Range): T | undefined {
+    return this.children
+      .getContaining(separator)
+      ?.getSmallLestContaining(separator);
+  }
+}
+
+function createNodes<T extends { range: Range }>(
+  items: T[],
+): RangeLookupList<RangeLookupTreeNode<T>> {
+  const results: RangeLookupTreeNode<T>[] = [];
+  const parents: RangeLookupTreeNode<T>[] = [];
+
+  for (const item of items) {
+    const node = new RangeLookupTreeNode(item);
+
+    while (
+      parents.length > 0 &&
+      !parents[parents.length - 1].range.contains(item.range)
+    ) {
+      parents.pop();
+    }
+
+    const parent = parents[parents.length - 1];
+
+    if (parent != null) {
+      parent.children.add(node);
+    } else {
+      results.push(node);
+    }
+
+    parents.push(node);
+  }
+
+  return new RangeLookupList(results);
+}
+
+class RangeLookupTreeNode<T extends { range: Range }> {
+  public children: RangeLookupList<RangeLookupTreeNode<T>>;
+
+  constructor(private item: T) {
+    this.children = new RangeLookupList([]);
+  }
+
+  get range(): Range {
+    return this.item.range;
+  }
+
+  getSmallLestContaining(range: Range): T {
+    const child = this.children
+      .getContaining(range)
+      ?.getSmallLestContaining(range);
+
+    return child ?? this.item;
+  }
+}