Use narrow search range instead of full document for Tree sitter queries (#2714)

Today we generate all captures in the entire document whenever we do a
Tree sitter query. This is of course quite expensive for large documents
and it would be better if we could narrow the search range. That is
what's implemented in this pull request.

I also found some incorrect scm patterns that broke with this
implementation that I have fixed.

Fixes #1769

## 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: Phil Cohen <[email protected]>

Author: AndreasArvidsson

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/TreeSitterScopeHandler/BaseTreeSitterScopeHandler.ts

@@ -3,12 +3,13 @@ import { showError } from "@cursorless/common";
 import { uniqWith } from "lodash-es";
 import type { TreeSitterQuery } from "../../../../languages/TreeSitterQuery";
 import type { QueryMatch } from "../../../../languages/TreeSitterQuery/QueryCapture";
+import { ide } from "../../../../singletons/ide.singleton";
 import { BaseScopeHandler } from "../BaseScopeHandler";
 import { compareTargetScopes } from "../compareTargetScopes";
 import type { TargetScope } from "../scope.types";
 import type { ScopeIteratorRequirements } from "../scopeHandler.types";
+import { getQuerySearchRange } from "./getQuerySearchRange";
 import { mergeAdjacentBy } from "./mergeAdjacentBy";
-import { ide } from "../../../../singletons/ide.singleton";
 
 /** Base scope handler to use for both tree-sitter scopes and their iteration scopes */
 export abstract class BaseTreeSitterScopeHandler extends BaseScopeHandler {
@@ -20,18 +21,20 @@ export abstract class BaseTreeSitterScopeHandler extends BaseScopeHandler {
     editor: TextEditor,
     position: Position,
     direction: Direction,
-    _hints: ScopeIteratorRequirements,
+    hints: ScopeIteratorRequirements,
   ): Iterable<TargetScope> {
     const { document } = editor;
 
-    // Due to a tree-sitter bug, we generate all scopes from the entire file
-    // instead of using `_hints` to restrict the search range to scopes we care
-    // about. The actual scopes yielded to the client are filtered by
-    // `BaseScopeHandler` anyway, so there's no impact on correctness, just
-    // performance. We'd like to roll this back; see #1769.
+    /** Narrow the range within which tree-sitter searches, for performance */
+    const { start, end } = getQuerySearchRange(
+      document,
+      position,
+      direction,
+      hints,
+    );
 
     const scopes = this.query
-      .matches(document)
+      .matches(document, start, end)
       .map((match) => this.matchToScope(editor, match))
       .filter((scope): scope is ExtendedTargetScope => scope != null)
       .sort((a, b) => compareTargetScopes(direction, position, a, b));

packages/cursorless-engine/src/processTargets/modifiers/scopeHandlers/TreeSitterScopeHandler/getQuerySearchRange.ts

@@ -0,0 +1,89 @@
+import type { Direction, Position, TextDocument } from "@cursorless/common";
+import type {
+  ContainmentPolicy,
+  ScopeIteratorRequirements,
+} from "../scopeHandler.types";
+
+/**
+ * Constructs a range to pass to {@link Query.matches} to find scopes. Note
+ * that {@link Query.matches} will only return scopes that have non-empty
+ * intersection with this range.  Also note that the base
+ * {@link BaseScopeHandler.generateScopes} will filter out any extra scopes
+ * that we yield, so we don't need to be totally precise.
+ *
+ * @returns Range to pass to {@link Query.matches}
+ */
+export function getQuerySearchRange(
+  document: TextDocument,
+  position: Position,
+  direction: Direction,
+  {
+    containment,
+    distalPosition,
+    allowAdjacentScopes,
+  }: ScopeIteratorRequirements,
+) {
+  const { start, end } = getQuerySearchRangeCore(
+    document.offsetAt(position),
+    document.offsetAt(distalPosition),
+    direction,
+    containment,
+    allowAdjacentScopes,
+  );
+
+  return {
+    start: document.positionAt(start),
+    end: document.positionAt(end),
+  };
+}
+
+/** Helper function for {@link getQuerySearchRange} */
+function getQuerySearchRangeCore(
+  offset: number,
+  distalOffset: number,
+  direction: Direction,
+  containment: ContainmentPolicy | null,
+  allowAdjacentScopes: boolean,
+) {
+  /**
+   * If we allow adjacent scopes, we need to shift some of our offsets by one
+   * character
+   */
+  const adjacentShift = allowAdjacentScopes ? 1 : 0;
+
+  if (containment === "required") {
+    // If containment is required, we smear the position left and right by one
+    // character so that we have a non-empty intersection with any scope that
+    // touches position.  Note that we can skip shifting the initial position
+    // if we disallow adjacent scopes.
+    return direction === "forward"
+      ? {
+          start: offset - adjacentShift,
+          end: offset + 1,
+        }
+      : {
+          start: offset - 1,
+          end: offset + adjacentShift,
+        };
+  }
+
+  // If containment is disallowed, we can shift the position forward by a
+  // character to avoid matching scopes that touch position. Otherwise, we
+  // shift the position backward by a character to ensure we get scopes that
+  // touch position, if we allow adjacent scopes.
+  const proximalShift = containment === "disallowed" ? 1 : -adjacentShift;
+
+  // FIXME: Don't go all the way to end of document when there is no
+  // distalOffset? Seems wasteful to query all the way to end of document for
+  // something like "next funk". Might be better to start smaller and
+  // exponentially grow
+  return direction === "forward"
+    ? {
+        start: offset + proximalShift,
+        end: distalOffset + adjacentShift,
+      }
+    : {
+        start: distalOffset - adjacentShift,
+        end: offset - proximalShift,
+      };
+}

queries/c.scm

@@ -182,7 +182,7 @@
 ) @_.domain
 
 ;;!! aaa = 0;
-(
+(_
   (assignment_expression
     left: (_) @name @value.leading.endOf
     right: (_) @value @name.trailing.startOf

queries/javascript.scm

@@ -6,7 +6,7 @@
 
 ;; Define this here because the `field_definition` node type doesn't exist
 ;; in typescript.
-(
+(_
   ;;!! class Foo {
   ;;!!   foo = () => {};
   ;;!    ^^^^^^^^^^^^^^^
@@ -31,7 +31,7 @@
   ";"? @namedFunction.end @functionName.domain.end
 )
 
-(
+(_
   ;;!! foo = ...;
   ;;!  ^^^-------
   (field_definition

queries/typescript.core.scm

@@ -39,7 +39,7 @@
 ) @_.domain
 
 ;; Define these here because these node types don't exist in javascript.
-(
+(_
   [
     ;;!! class Foo { foo() {} }
     ;;!              ^^^^^^^^
@@ -80,7 +80,7 @@
   ";"? @namedFunction.end @functionName.domain.end @name.domain.end
 )
 
-(
+(_
   ;;!! (public | private | protected) foo = ...;
   ;;!  -----------------------------------------
   (public_field_definition
@@ -92,7 +92,7 @@
   ";"? @_.domain.end
 )
 
-(
+(_
   ;;!! (public | private | protected) foo: Bar = ...;
   ;;!  ----------------------------------------------
   (public_field_definition
@@ -340,14 +340,15 @@
 ;;!                   ^^^^
 ;;!                   xxxxxx
 ;;!                   ------------
-(
+(_
   (property_signature
     name: (_) @collectionKey @type.leading.endOf
     type: (_
       ":"
       (_) @type @collectionKey.trailing.startOf
     )
   ) @_.domain.start
+  .
   ";"? @_.domain.end
 )
 
@@ -372,11 +373,12 @@
 )
 
 ;; Statements with optional trailing `;`
-(
+(_
   [
     (property_signature)
     (public_field_definition)
     (abstract_method_signature)
   ] @statement.start
+  .
   ";"? @statement.end
 )