Fix surrounding pair performance (#2700)

With a 10k line json. `"take pair"` went from about 12 seconds to 200 ms
on my computer

## Checklist

- [x] 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

---------

Co-authored-by: Pokey Rule <755842+pokey@users.noreply.github.com>

Author: AndreasArvidsson

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/actions/ExecuteCommand.ts

@@ -14,6 +14,7 @@ import type { ActionReturnValue } from "./actions.types";
  */
 export default class ExecuteCommand {
   private callbackAction: CallbackAction;
+
   constructor(rangeUpdater: RangeUpdater) {
     this.callbackAction = new CallbackAction(rangeUpdater);
     this.run = this.run.bind(this);

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 = new LanguageDefinitionCache();
+  }
+
+  /**
+   * 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,29 @@
+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[]> = {};
+
+  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,17 @@ export interface LanguageDefinitions {
    */
   get(languageId: string): LanguageDefinition | undefined;
 
+  /**
+   * Clear the cache of all language definitions. This is run at the start of a command.
+   * This isn't strict necessary for normal user operations since whenever the user
+   * makes a change to the document the document version is updated. When
+   * running our test though we keep closing and reopening an untitled document.
+   * That test document will have the same uri and version unfortunately. Also
+   * to be completely sure there isn't some extension doing similar trickery
+   * it's just good hygiene to clear the cache before every command.
+   */
+  clearCache(): void;
+
   /**
    * @deprecated Only for use in legacy containing scope stage
    */
@@ -155,6 +166,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/OneWayNestedRangeFinder.test.ts

@@ -0,0 +1,30 @@
+import { Range } from "@cursorless/common";
+import { OneWayNestedRangeFinder } from "./OneWayNestedRangeFinder";
+import assert from "assert";
+
+const items = [
+  { range: new Range(0, 0, 0, 5) },
+  { range: new Range(0, 1, 0, 4) },
+  { range: new Range(0, 2, 0, 3) },
+  { range: new Range(0, 6, 0, 8) },
+];
+
+suite("OneWayNestedRangeFinder", () => {
+  test("getSmallestContaining 1", () => {
+    const finder = new OneWayNestedRangeFinder(items);
+    const actual = finder.getSmallestContaining(new Range(0, 2, 0, 2));
+    assert.equal(actual?.range.toString(), new Range(0, 2, 0, 3).toString());
+  });
+
+  test("getSmallestContaining 2", () => {
+    const finder = new OneWayNestedRangeFinder(items);
+    const actual = finder.getSmallestContaining(new Range(0, 7, 0, 7));
+    assert.equal(actual?.range.toString(), new Range(0, 6, 0, 8).toString());
+  });
+
+  test("getSmallestContaining 3", () => {
+    const finder = new OneWayNestedRangeFinder(items);
+    const actual = finder.getSmallestContaining(new Range(0, 0, 0, 0));
+    assert.equal(actual?.range.toString(), new Range(0, 0, 0, 5).toString());
+  });
+});

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

@@ -0,0 +1,76 @@
+import type { Range } from "@cursorless/common";
+import { OneWayRangeFinder } from "./OneWayRangeFinder";
+
+/**
+ * Given a list of ranges (the haystack), allows the client to search for smallest range containing a range (the needle).
+ * Has the following requirements:
+ * - the haystack must be sorted in document order
+ * - **the needles must be in document order as well**. This enables us to avoid backtracking as you search for a sequence of items.
+ * - the haystack entries **may** be nested, but one haystack entry cannot partially contain another
+ */
+export class OneWayNestedRangeFinder<T extends { range: Range }> {
+  private children: OneWayRangeFinder<RangeLookupTreeNode<T>>;
+
+  /**
+   * @param items The items to search in. Must be sorted in document order.
+   */
+  constructor(items: T[]) {
+    this.children = createNodes(items);
+  }
+
+  getSmallestContaining(separator: Range): T | undefined {
+    return this.children
+      .getContaining(separator)
+      ?.getSmallestContaining(separator);
+  }
+}
+
+function createNodes<T extends { range: Range }>(
+  items: T[],
+): OneWayRangeFinder<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 OneWayRangeFinder(results);
+}
+
+class RangeLookupTreeNode<T extends { range: Range }> {
+  public children: OneWayRangeFinder<RangeLookupTreeNode<T>>;
+
+  constructor(private item: T) {
+    this.children = new OneWayRangeFinder([]);
+  }
+
+  get range(): Range {
+    return this.item.range;
+  }
+
+  getSmallestContaining(range: Range): T {
+    const child = this.children
+      .getContaining(range)
+      ?.getSmallestContaining(range);
+
+    return child ?? this.item;
+  }
+}

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

@@ -0,0 +1,52 @@
+import type { Range } from "@cursorless/common";
+
+/**
+ * Given a list of ranges (the haystack), allows the client to search for a sequence of ranges (the needles).
+ * Has the following requirements:
+ * - the haystack must be sorted in document order
+ * - **the needles must be in document order as well**. This enables us to avoid backtracking as you search for a sequence of items.
+ * - the haystack entries must not overlap. Adjacent is fine
+ */
+export class OneWayRangeFinder<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;
+  }
+}