Use narrow search range instead of full document for Tree sitter queries

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
 )