Added Tree sitter query cache (#2711)

Uses the cache added in #2700 to also optimize all tree sitter queries,
not just surrounding pairs.

## 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

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: AndreasArvidsson

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

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

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

@@ -3,7 +3,6 @@ import type {
   ScopeType,
   SimpleScopeType,
   SimpleScopeTypeType,
-  StringRecord,
   TreeSitter,
 } from "@cursorless/common";
 import {
@@ -13,7 +12,6 @@ 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";
@@ -23,18 +21,14 @@ 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
@@ -88,41 +82,24 @@ export class LanguageDefinition {
   }
 
   /**
-   * This is a low-level function that just returns a list of captures of the given
-   * capture name in the document. We use this in our surrounding pair code.
+   * This is a low-level function that just returns a map of all captures in the
+   * document. We use this in our surrounding pair code.
    *
    * @param document The document to search
    * @param captureName The name of a capture to search for
-   * @returns A list of captures of the given capture name in the document
-   */
-  getCaptures(
-    document: TextDocument,
-    captureName: SimpleScopeTypeType,
-  ): QueryCapture[] {
-    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.
+   * @returns A map of captures in the document
    */
-  private getCapturesMap(document: TextDocument): StringRecord<QueryCapture[]> {
+  getCapturesMap(document: TextDocument) {
     const matches = this.query.matches(document);
-    const result: StringRecord<QueryCapture[]> = {};
+    const result: Partial<Record<SimpleScopeTypeType, QueryCapture[]>> = {};
 
     for (const match of matches) {
       for (const capture of match.captures) {
-        if (result[capture.name] == null) {
-          result[capture.name] = [];
+        const name = capture.name as SimpleScopeTypeType;
+        if (result[name] == null) {
+          result[name] = [];
         }
-        result[capture.name]!.push(capture);
+        result[name]!.push(capture);
       }
     }
 

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

@@ -10,6 +10,7 @@ import {
 import { toString } from "lodash-es";
 import type { SyntaxNode } from "web-tree-sitter";
 import { LanguageDefinition } from "./LanguageDefinition";
+import { treeSitterQueryCache } from "./TreeSitterQuery/treeSitterQueryCache";
 
 /**
  * Sentinel value to indicate that a language doesn't have
@@ -32,17 +33,6 @@ 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
    */
@@ -82,7 +72,13 @@ export class LanguageDefinitionsImpl
     private treeSitter: TreeSitter,
     private treeSitterQueryProvider: RawTreeSitterQueryProvider,
   ) {
+    const isTesting = ide.runMode === "test";
+
     ide.onDidOpenTextDocument((document) => {
+      // During testing we open untitled documents that all have the same uri and version which breaks our cache
+      if (isTesting) {
+        treeSitterQueryCache.clear();
+      }
       void this.loadLanguage(document.languageId);
     });
     ide.onDidChangeVisibleTextEditors((editors) => {
@@ -150,6 +146,7 @@ export class LanguageDefinitionsImpl
   private async reloadLanguageDefinitions(): Promise<void> {
     this.languageDefinitions.clear();
     await this.loadAllLanguages();
+    treeSitterQueryCache.clear();
     this.notifier.notifyListeners();
   }
 
@@ -166,14 +163,6 @@ 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/languages/LanguageDefinitions.ts

@@ -14,6 +14,7 @@ import { isContainedInErrorNode } from "./isContainedInErrorNode";
 import { parsePredicates } from "./parsePredicates";
 import { predicateToString } from "./predicateToString";
 import { rewriteStartOfEndOf } from "./rewriteStartOfEndOf";
+import { treeSitterQueryCache } from "./treeSitterQueryCache";
 
 /**
  * Wrapper around a tree-sitter query that provides a more convenient API, and
@@ -70,6 +71,18 @@ export class TreeSitterQuery {
     document: TextDocument,
     start?: Position,
     end?: Position,
+  ): QueryMatch[] {
+    if (!treeSitterQueryCache.isValid(document, start, end)) {
+      const matches = this.getAllMatches(document, start, end);
+      treeSitterQueryCache.update(document, start, end, matches);
+    }
+    return treeSitterQueryCache.get();
+  }
+
+  private getAllMatches(
+    document: TextDocument,
+    start?: Position,
+    end?: Position,
   ): QueryMatch[] {
     return this.query
       .matches(this.treeSitter.getTree(document).rootNode, {

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

@@ -0,0 +1,54 @@
+import type { Position, TextDocument } from "@cursorless/common";
+import type { QueryMatch } from "./QueryCapture";
+
+export class Cache {
+  private documentUri: string = "";
+  private documentVersion: number = -1;
+  private documentLanguageId: string = "";
+  private startPosition: Position | undefined;
+  private endPosition: Position | undefined;
+  private matches: QueryMatch[] = [];
+
+  clear() {
+    this.documentUri = "";
+    this.documentVersion = -1;
+    this.documentLanguageId = "";
+    this.startPosition = undefined;
+    this.endPosition = undefined;
+    this.matches = [];
+  }
+
+  isValid(
+    document: TextDocument,
+    startPosition: Position | undefined,
+    endPosition: Position | undefined,
+  ) {
+    return (
+      this.documentUri === document.uri.toString() &&
+      this.documentVersion === document.version &&
+      this.documentLanguageId === document.languageId &&
+      this.startPosition === startPosition &&
+      this.endPosition === endPosition
+    );
+  }
+
+  update(
+    document: TextDocument,
+    startPosition: Position | undefined,
+    endPosition: Position | undefined,
+    matches: QueryMatch[],
+  ) {
+    this.documentUri = document.uri.toString();
+    this.documentVersion = document.version;
+    this.documentLanguageId = document.languageId;
+    this.startPosition = startPosition;
+    this.endPosition = endPosition;
+    this.matches = matches;
+  }
+
+  get(): QueryMatch[] {
+    return this.matches;
+  }
+}
+
+export const treeSitterQueryCache = new Cache();

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

@@ -1,9 +1,4 @@
-import {
-  matchAllIterator,
-  Range,
-  type SimpleScopeTypeType,
-  type TextDocument,
-} from "@cursorless/common";
+import { matchAllIterator, Range, type TextDocument } from "@cursorless/common";
 import type { LanguageDefinition } from "../../../../languages/LanguageDefinition";
 import type { QueryCapture } from "../../../../languages/TreeSitterQuery/QueryCapture";
 import { getDelimiterRegex } from "./getDelimiterRegex";
@@ -28,12 +23,12 @@ export function getDelimiterOccurrences(
     return [];
   }
 
+  const capturesMap = languageDefinition?.getCapturesMap(document) ?? {};
   const disqualifyDelimiters = new OneWayRangeFinder(
-    getSortedCaptures(languageDefinition, document, "disqualifyDelimiter"),
+    getSortedCaptures(capturesMap.disqualifyDelimiter),
   );
-  // We need a tree for text fragments since they can be nested
   const textFragments = new OneWayNestedRangeFinder(
-    getSortedCaptures(languageDefinition, document, "textFragment"),
+    getSortedCaptures(capturesMap.textFragment),
   );
 
   const delimiterTextToDelimiterInfoMap = Object.fromEntries(
@@ -74,12 +69,10 @@ export function getDelimiterOccurrences(
   return results;
 }
 
-function getSortedCaptures(
-  languageDefinition: LanguageDefinition | undefined,
-  document: TextDocument,
-  captureName: SimpleScopeTypeType,
-): QueryCapture[] {
-  const items = languageDefinition?.getCaptures(document, captureName) ?? [];
+function getSortedCaptures(items?: QueryCapture[]): QueryCapture[] {
+  if (items == null) {
+    return [];
+  }
   items.sort((a, b) => a.range.start.compareTo(b.range.start));
   return items;
 }

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

@@ -71,8 +71,6 @@ export async function runCommand(
     commandRunner = decorator.wrapCommandRunner(readableHatMap, commandRunner);
   }
 
-  languageDefinitions.clearCache();
-
   const response = await commandRunner.run(commandComplete);
 
   return await unwrapLegacyCommandResponse(command, response);