Increase surrounding pair performance

packages/common/src/types/command/PartialTargetDescriptor.types.ts

@@ -230,7 +230,7 @@ export interface SimpleScopeType {
   type: SimpleScopeTypeType;
 }
 
-export type ScopeTypeType = SimpleScopeTypeType | ScopeType["type"];
+export type ScopeTypeType = ScopeType["type"];
 
 export interface CustomRegexScopeType {
   type: "customRegex";

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

@@ -87,16 +87,22 @@ export class LanguageDefinition {
    * 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
+   * @param captureNames Optional capture names to include
    * @returns A map of captures in the document
    */
-  getCapturesMap(document: TextDocument) {
-    const matches = this.query.matches(document);
-    const result: Partial<Record<SimpleScopeTypeType, QueryCapture[]>> = {};
+  getCapturesMap<T extends SimpleScopeTypeType>(
+    document: TextDocument,
+    captureNames: readonly T[],
+  ) {
+    const matches = this.query.matchesForCaptures(
+      document,
+      new Set(captureNames),
+    );
+    const result: Partial<Record<T, QueryCapture[]>> = {};
 
     for (const match of matches) {
       for (const capture of match.captures) {
-        const name = capture.name as SimpleScopeTypeType;
+        const name = capture.name as T;
         if (result[name] == null) {
           result[name] = [];
         }

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

@@ -1,10 +1,10 @@
 import type { Position, TextDocument, TreeSitter } from "@cursorless/common";
 import type * as treeSitter from "web-tree-sitter";
 import { ide } from "../../singletons/ide.singleton";
+import { getNormalizedCaptureName } from "./captureNames";
 import { checkCaptureStartEnd } from "./checkCaptureStartEnd";
 import { getNodeRange } from "./getNodeRange";
 import { isContainedInErrorNode } from "./isContainedInErrorNode";
-import { normalizeCaptureName } from "./normalizeCaptureName";
 import { parsePredicatesWithErrorHandling } from "./parsePredicatesWithErrorHandling";
 import { positionToPoint } from "./positionToPoint";
 import type {
@@ -55,7 +55,7 @@ export class TreeSitterQuery {
 
   hasCapture(name: string): boolean {
     return this.query.captureNames.some(
-      (n) => normalizeCaptureName(n) === name,
+      (n) => getNormalizedCaptureName(n) === name,
     );
   }
 
@@ -64,29 +64,86 @@ export class TreeSitterQuery {
     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 this.getMatches(document, start, end, undefined);
+  }
+
+  matchesForCaptures(
+    document: TextDocument,
+    captureNames: Set<string>,
+  ): QueryMatch[] {
+    return this.getMatches(document, undefined, undefined, captureNames);
+  }
+
+  private getMatches(
+    document: TextDocument,
+    start: Position | undefined,
+    end: Position | undefined,
+    captureNameFilter: Set<string> | undefined,
+  ): QueryMatch[] {
+    if (
+      !treeSitterQueryCache.isValid(document, start, end, captureNameFilter)
+    ) {
+      const matches = this.calculateMatches(
+        document,
+        start,
+        end,
+        captureNameFilter,
+      );
+      treeSitterQueryCache.update(
+        document,
+        start,
+        end,
+        captureNameFilter,
+        matches,
+      );
     }
     return treeSitterQueryCache.get();
   }
 
-  private getAllMatches(
+  private calculateMatches(
     document: TextDocument,
-    start?: Position,
-    end?: Position,
+    start: Position | undefined,
+    end: Position | undefined,
+    captureNameFilter: Set<string> | undefined,
   ): QueryMatch[] {
     const matches = this.getTreeMatches(document, start, end);
     const results: QueryMatch[] = [];
 
     for (const match of matches) {
-      const mutableMatch = this.createMutableQueryMatch(document, match);
+      if (
+        captureNameFilter != null &&
+        !match.captures.some((capture) =>
+          captureNameFilter.has(getNormalizedCaptureName(capture.name)),
+        )
+      ) {
+        continue;
+      }
+
+      const hasPatternPredicates =
+        this.patternPredicates[match.patternIndex].length > 0;
+
+      const mutableMatch = this.createMutableQueryMatch(
+        document,
+        match,
+        // If there are pattern predicates, we need to include all captures when
+        // creating the mutable match, since the predicates may depend on any of
+        // the captures.
+        !hasPatternPredicates ? captureNameFilter : undefined,
+      );
 
       if (!this.runPredicates(mutableMatch)) {
         continue;
       }
 
-      results.push(this.createQueryMatch(mutableMatch));
+      const queryMatch = this.createQueryMatch(
+        mutableMatch,
+        // We only need to filter here if we didn't filter in createMutableQueryMatch()
+        hasPatternPredicates ? captureNameFilter : undefined,
+      );
+
+      if (queryMatch != null) {
+        results.push(queryMatch);
+      }
     }
 
     return results;
@@ -107,18 +164,32 @@ export class TreeSitterQuery {
   private createMutableQueryMatch(
     document: TextDocument,
     match: treeSitter.QueryMatch,
+    captureNameFilter?: Set<string>,
   ): MutableQueryMatch {
-    return {
-      patternIdx: match.patternIndex,
-      captures: match.captures.map(({ name, node }) => ({
+    const captures: MutableQueryCapture[] = [];
+
+    for (const { name, node } of match.captures) {
+      if (
+        captureNameFilter != null &&
+        !captureNameFilter.has(getNormalizedCaptureName(name))
+      ) {
+        continue;
+      }
+
+      captures.push({
         name,
         node,
         document,
         range: getNodeRange(node),
         insertionDelimiter: undefined,
         allowMultiple: false,
         hasError: () => isContainedInErrorNode(node),
-      })),
+      });
+    }
+
+    return {
+      patternIdx: match.patternIndex,
+      captures,
     };
   }
 
@@ -131,7 +202,10 @@ export class TreeSitterQuery {
     return true;
   }
 
-  private createQueryMatch(match: MutableQueryMatch): QueryMatch {
+  private createQueryMatch(
+    match: MutableQueryMatch,
+    captureNameFilter?: Set<string>,
+  ): QueryMatch | undefined {
     const result: MutableQueryCapture[] = [];
     const map = new Map<
       string,
@@ -144,7 +218,10 @@ export class TreeSitterQuery {
     // name, for which we'd return a capture with name `foo`.
 
     for (const capture of match.captures) {
-      const name = normalizeCaptureName(capture.name);
+      const name = getNormalizedCaptureName(capture.name);
+      if (captureNameFilter != null && !captureNameFilter.has(name)) {
+        continue;
+      }
       const range = getStartOfEndOfRange(capture);
       const existing = map.get(name);
 
@@ -168,6 +245,10 @@ export class TreeSitterQuery {
       }
     }
 
+    if (result.length === 0) {
+      return undefined;
+    }
+
     if (this.shouldCheckCaptures) {
       this.checkCaptures(Array.from(map.values()));
     }

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

@@ -8,41 +8,47 @@ export class TreeSitterQueryCache {
   private startPosition: Position | undefined;
   private endPosition: Position | undefined;
   private matches: QueryMatch[] = [];
+  private captureNames: Set<string> | undefined;
 
   clear() {
     this.documentUri = "";
     this.documentVersion = -1;
     this.documentLanguageId = "";
     this.startPosition = undefined;
     this.endPosition = undefined;
+    this.captureNames = undefined;
     this.matches = [];
   }
 
   isValid(
     document: TextDocument,
     startPosition: Position | undefined,
     endPosition: Position | undefined,
+    captureNames: Set<string> | undefined,
   ) {
     return (
       this.documentVersion === document.version &&
       this.documentUri === document.uri.toString() &&
       this.documentLanguageId === document.languageId &&
       positionsEqual(this.startPosition, startPosition) &&
-      positionsEqual(this.endPosition, endPosition)
+      positionsEqual(this.endPosition, endPosition) &&
+      setEqual(this.captureNames, captureNames)
     );
   }
 
   update(
     document: TextDocument,
     startPosition: Position | undefined,
     endPosition: Position | undefined,
+    captureNames: Set<string> | undefined,
     matches: QueryMatch[],
   ) {
     this.documentVersion = document.version;
     this.documentUri = document.uri.toString();
     this.documentLanguageId = document.languageId;
     this.startPosition = startPosition;
     this.endPosition = endPosition;
+    this.captureNames = captureNames;
     this.matches = matches;
   }
 
@@ -58,4 +64,19 @@ function positionsEqual(a: Position | undefined, b: Position | undefined) {
   return a.isEqual(b);
 }
 
+function setEqual(a: Set<string> | undefined, b: Set<string> | undefined) {
+  if (a == null || b == null) {
+    return a === b;
+  }
+  if (a.size !== b.size) {
+    return false;
+  }
+  for (const item of a) {
+    if (!b.has(item)) {
+      return false;
+    }
+  }
+  return true;
+}
+
 export const treeSitterQueryCache = new TreeSitterQueryCache();

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

@@ -0,0 +1,86 @@
+import { pseudoScopes, simpleScopeTypeTypes } from "@cursorless/common";
+
+const wildcard = "_";
+const captureNames = [
+  ...simpleScopeTypeTypes.filter((s) => !pseudoScopes.has(s)),
+  wildcard,
+  "interior",
+];
+
+const positionRelationships = ["prefix", "leading", "trailing"];
+const positionSuffixes = [
+  "startOf",
+  "endOf",
+  "start.startOf",
+  "start.endOf",
+  "end.startOf",
+  "end.endOf",
+];
+
+const rangeRelationships = [
+  "domain",
+  "removal",
+  "iteration",
+  "iteration.domain",
+];
+const rangeSuffixes = [
+  "start",
+  "end",
+  "start.startOf",
+  "start.endOf",
+  "end.startOf",
+  "end.endOf",
+];
+
+const allowedCaptures = new Set<string>();
+
+for (const captureName of captureNames) {
+  // Wildcard is not allowed by itself without a relationship
+  if (captureName !== wildcard) {
+    // eg: statement
+    allowedCaptures.add(captureName);
+
+    // eg: statement.start | statement.start.endOf
+    for (const suffix of rangeSuffixes) {
+      allowedCaptures.add(`${captureName}.${suffix}`);
+    }
+  }
+
+  for (const relationship of positionRelationships) {
+    // eg: statement.leading
+    allowedCaptures.add(`${captureName}.${relationship}`);
+
+    for (const suffix of positionSuffixes) {
+      // eg: statement.leading.endOf
+      allowedCaptures.add(`${captureName}.${relationship}.${suffix}`);
+    }
+  }
+
+  for (const relationship of rangeRelationships) {
+    // eg: statement.domain
+    allowedCaptures.add(`${captureName}.${relationship}`);
+
+    for (const suffix of rangeSuffixes) {
+      // eg: statement.domain.start | statement.domain.start.endOf
+      allowedCaptures.add(`${captureName}.${relationship}.${suffix}`);
+    }
+  }
+}
+
+const normalizedCaptureNamesMap = new Map<string, string>();
+
+for (const captureName of allowedCaptures) {
+  normalizedCaptureNamesMap.set(captureName, normalizeCaptureName(captureName));
+}
+
+function normalizeCaptureName(name: string): string {
+  return name.replace(/(\.(start|end))?(\.(startOf|endOf))?$/, "");
+}
+
+export function isCaptureAllowed(captureName: string): boolean {
+  return allowedCaptures.has(captureName);
+}
+
+export function getNormalizedCaptureName(captureName: string): string {
+  return normalizedCaptureNamesMap.get(captureName) ?? captureName;
+}