Add mutable query capture interface

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

@@ -1,9 +1,9 @@
-import { SyntaxNode } from "web-tree-sitter";
 import z from "zod";
 import {
   SchemaInputType,
   SchemaOutputType,
 } from "./operatorArgumentSchemaTypes";
+import { MutableQueryCapture } from "./QueryCapture";
 
 /**
  * A schema used to validate a list of operands for a given predicate operator.
@@ -38,7 +38,7 @@ export type InferSchemaType<T extends HasSchema> = T["schema"];
 type PredicateParameterType<T extends SchemaOutputType> = T extends {
   type: "capture";
 }
-  ? SyntaxNode
+  ? MutableQueryCapture
   : T extends { value: infer V }
   ? V
   : never;

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

@@ -0,0 +1,57 @@
+import { Range } from "@cursorless/common";
+import { SyntaxNode } from "web-tree-sitter";
+
+/**
+ * A capture of a query pattern against a syntax tree.  Often corresponds to a
+ * node within the syntax tree, but can also be a range within a node or be modified
+ * by query predicate operators.
+ */
+export interface QueryCapture {
+  /**
+   * The name of the capture.  Eg for a capture labeled `@foo`, the name is
+   * `foo`.
+   */
+  readonly name: string;
+
+  /** The range of the capture. */
+  readonly range: Range;
+}
+
+/**
+ * A match of a query pattern against a syntax tree.
+ */
+export interface QueryMatch {
+  /**
+   * The captures of the pattern that was matched.
+   */
+  readonly captures: QueryCapture[];
+}
+
+/**
+ * A capture of a query pattern against a syntax tree. This type is used
+ * internally by the query engine to allow operators to modify the capture.
+ */
+export interface MutableQueryCapture extends QueryCapture {
+  /**
+   * The tree-sitter node that was captured. Note that the range may have already
+   * been altered by a prior operator, so please use {@link range} instead of
+   * trying to retrieve the range from the node.
+   */
+  readonly node: SyntaxNode;
+
+  range: Range;
+}
+
+/**
+ * A match of a query pattern against a syntax tree that can be mutated. This
+ * type is used internally by the query engine to allow operators to modify the
+ * match.
+ */
+export interface MutableQueryMatch extends QueryMatch {
+  /**
+   * The index of the pattern that was matched.
+   */
+  readonly patternIdx: number;
+
+  readonly captures: MutableQueryCapture[];
+}

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

@@ -1,10 +1,11 @@
-import { PredicateOperand, QueryMatch } from "web-tree-sitter";
+import { PredicateOperand } from "web-tree-sitter";
 import z from "zod";
 import {
   AcceptFunctionArgs,
   HasSchema,
   InferSchemaType,
 } from "./PredicateOperatorSchemaTypes";
+import { MutableQueryMatch } from "./QueryCapture";
 import { constructZodErrorMessages } from "./constructZodErrorMessages";
 
 /**
@@ -24,7 +25,7 @@ export abstract class QueryPredicateOperator<T extends HasSchema> {
   /**
    * The name of the operator, e.g. `not-type?`
    */
-  abstract readonly name: `${string}?`;
+  abstract readonly name: `${string}${"?" | "!"}`;
 
   /**
    * The schema used to validate the operands to the operator.  This should be a
@@ -42,7 +43,7 @@ export abstract class QueryPredicateOperator<T extends HasSchema> {
    * in the schema.  For example, if the schema is `z.tuple([q.node, q.string])`,
    * then `args` will be `SyntaxNode, string`.
    */
-  protected abstract accept(
+  protected abstract run(
     ...args: AcceptFunctionArgs<z.infer<InferSchemaType<T>>>
   ): boolean;
 
@@ -61,8 +62,8 @@ export abstract class QueryPredicateOperator<T extends HasSchema> {
     return result.success
       ? {
           success: true,
-          predicate: (match: QueryMatch) =>
-            this.accept(...this.constructAcceptArgs(result.data, match)),
+          predicate: (match: MutableQueryMatch) =>
+            this.run(...this.constructAcceptArgs(result.data, match)),
         }
       : {
           success: false,
@@ -79,25 +80,25 @@ export abstract class QueryPredicateOperator<T extends HasSchema> {
    */
   private constructAcceptArgs(
     rawOutput: z.output<InferSchemaType<T>>,
-    match: QueryMatch,
+    match: MutableQueryMatch,
   ): AcceptFunctionArgs<z.infer<InferSchemaType<T>>> {
     return rawOutput.map((operand) => {
       if (operand.type === "capture") {
-        const node = match.captures.find(
+        const capture = match.captures.find(
           (capture) => capture.name === operand.name,
-        )?.node;
+        );
 
-        if (node == null) {
+        if (capture == null) {
           // FIXME: We could allow some predicates to be forgiving,
-          // because it's possible to have a capture on an optional node.
+          // because it's possible to have a capture on an optional nodeInfo.
           // In that case we'd prob just return `true` if any capture was
           // `null`, but we should check that the given capture name
           // appears statically in the given pattern.  But we don't yet
           // have a use case so let's leave it for now.
           throw new Error(`Could not find capture ${operand.name}`);
         }
 
-        return node;
+        return capture;
       } else {
         return operand.value;
       }
@@ -107,7 +108,7 @@ export abstract class QueryPredicateOperator<T extends HasSchema> {
 
 interface SuccessfulPredicateResult {
   success: true;
-  predicate: (match: QueryMatch) => boolean;
+  predicate: (match: MutableQueryMatch) => boolean;
 }
 
 interface FailedPredicateResult {

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

@@ -1,9 +1,11 @@
 import { Position, TextDocument, showError } from "@cursorless/common";
-import { Point, Query, QueryMatch } from "web-tree-sitter";
+import { Point, Query } from "web-tree-sitter";
 import { ide } from "../../singletons/ide.singleton";
 import { TreeSitter } from "../../typings/TreeSitter";
+import { getNodeRange } from "../../util/nodeSelectors";
 import { parsePredicates } from "./parsePredicates";
 import { predicateToString } from "./predicateToString";
+import { MutableQueryMatch, QueryMatch } from "./QueryCapture";
 
 /**
  * Wrapper around a tree-sitter query that provides a more convenient API, and
@@ -23,7 +25,7 @@ export class TreeSitterQuery {
      * array corresponds to a pattern, and each element of the inner array
      * corresponds to a predicate for that pattern.
      */
-    private patternPredicates: ((match: QueryMatch) => boolean)[][],
+    private patternPredicates: ((match: MutableQueryMatch) => boolean)[][],
   ) {}
 
   static create(languageId: string, treeSitter: TreeSitter, query: Query) {
@@ -56,16 +58,30 @@ export class TreeSitterQuery {
     return new TreeSitterQuery(treeSitter, query, predicates);
   }
 
-  matches(document: TextDocument, start: Position, end: Position) {
+  matches(
+    document: TextDocument,
+    start: Position,
+    end: Position,
+  ): QueryMatch[] {
     return this.query
       .matches(
         this.treeSitter.getTree(document).rootNode,
         positionToPoint(start),
         positionToPoint(end),
       )
-      .filter((rawMatch) =>
-        this.patternPredicates[rawMatch.pattern].every((predicate) =>
-          predicate(rawMatch),
+      .map(
+        ({ pattern, captures }): MutableQueryMatch => ({
+          patternIdx: pattern,
+          captures: captures.map(({ name, node }) => ({
+            name,
+            node,
+            range: getNodeRange(node),
+          })),
+        }),
+      )
+      .filter((match) =>
+        this.patternPredicates[match.patternIdx].every((predicate) =>
+          predicate(match),
         ),
       );
   }

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

@@ -1,4 +1,5 @@
-import { PredicateResult, QueryMatch } from "web-tree-sitter";
+import { PredicateResult } from "web-tree-sitter";
+import { MutableQueryMatch } from "./QueryCapture";
 import { queryPredicateOperators } from "./queryPredicateOperators";
 
 /**
@@ -15,11 +16,11 @@ import { queryPredicateOperators } from "./queryPredicateOperators";
  */
 export function parsePredicates(predicateDescriptors: PredicateResult[][]) {
   const errors: PredicateError[] = [];
-  const predicates: ((match: QueryMatch) => boolean)[][] = [];
+  const predicates: ((match: MutableQueryMatch) => boolean)[][] = [];
 
   predicateDescriptors.forEach((patternPredicateDescriptors, patternIdx) => {
     /** The predicates for a given pattern */
-    const patternPredicates: ((match: QueryMatch) => boolean)[] = [];
+    const patternPredicates: ((match: MutableQueryMatch) => boolean)[] = [];
 
     patternPredicateDescriptors.forEach((predicateDescriptor, predicateIdx) => {
       const operator = queryPredicateOperators.find(

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

@@ -1,8 +1,8 @@
-import { SyntaxNode } from "web-tree-sitter";
 import z from "zod";
+import { HasSchema } from "./PredicateOperatorSchemaTypes";
+import { MutableQueryCapture } from "./QueryCapture";
 import { QueryPredicateOperator } from "./QueryPredicateOperator";
 import { q } from "./operatorArgumentSchemaTypes";
-import { HasSchema } from "./PredicateOperatorSchemaTypes";
 
 /**
  * A predicate operator that returns true if the node is not of the given type.
@@ -13,7 +13,7 @@ import { HasSchema } from "./PredicateOperatorSchemaTypes";
 class NotType extends QueryPredicateOperator<NotType> {
   name = "not-type?" as const;
   schema = z.tuple([q.node, q.string]).rest(q.string);
-  accept(node: SyntaxNode, ...types: string[]) {
+  run({ node }: MutableQueryCapture, ...types: string[]) {
     return !types.includes(node.type);
   }
 }
@@ -27,7 +27,7 @@ class NotType extends QueryPredicateOperator<NotType> {
 class NotParentType extends QueryPredicateOperator<NotParentType> {
   name = "not-parent-type?" as const;
   schema = z.tuple([q.node, q.string]).rest(q.string);
-  accept(node: SyntaxNode, ...types: string[]) {
+  run({ node }: MutableQueryCapture, ...types: string[]) {
     return node.parent == null || !types.includes(node.parent.type);
   }
 }
@@ -40,7 +40,7 @@ class NotParentType extends QueryPredicateOperator<NotParentType> {
 class IsNthChild extends QueryPredicateOperator<IsNthChild> {
   name = "is-nth-child?" as const;
   schema = z.tuple([q.node, q.integer]);
-  accept(node: SyntaxNode, n: number) {
+  run({ node }: MutableQueryCapture, n: number) {
     return node.parent?.children.indexOf(node) === n;
   }
 }

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

@@ -4,8 +4,8 @@ import {
   TextDocument,
   TextEditor,
 } from "@cursorless/common";
-import { QueryMatch } from "web-tree-sitter";
 import { TreeSitterQuery } from "../../../../languages/TreeSitterQuery";
+import { QueryMatch } from "../../../../languages/TreeSitterQuery/QueryCapture";
 import BaseScopeHandler from "../BaseScopeHandler";
 import { compareTargetScopes } from "../compareTargetScopes";
 import { TargetScope } from "../scope.types";

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

@@ -1,10 +1,10 @@
 import { ScopeType, SimpleScopeType, TextEditor } from "@cursorless/common";
-import { QueryMatch } from "web-tree-sitter";
 import { TreeSitterQuery } from "../../../../languages/TreeSitterQuery";
 import { PlainTarget } from "../../../targets";
 import { TargetScope } from "../scope.types";
 import { BaseTreeSitterScopeHandler } from "./BaseTreeSitterScopeHandler";
 import { getCaptureRangeByName, getRelatedRange } from "./captureUtils";
+import { QueryMatch } from "../../../../languages/TreeSitterQuery/QueryCapture";
 
 /** Scope handler to be used for iteration scopes of tree-sitter scope types */
 export class TreeSitterIterationScopeHandler extends BaseTreeSitterScopeHandler {

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

@@ -1,7 +1,6 @@
 import { SimpleScopeType, TextEditor } from "@cursorless/common";
-
-import { QueryMatch } from "web-tree-sitter";
 import { TreeSitterQuery } from "../../../../languages/TreeSitterQuery";
+import { QueryMatch } from "../../../../languages/TreeSitterQuery/QueryCapture";
 import ScopeTypeTarget from "../../../targets/ScopeTypeTarget";
 import { TargetScope } from "../scope.types";
 import { CustomScopeType } from "../scopeHandler.types";

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

@@ -1,6 +1,6 @@
 import { ScopeType, TextEditor } from "@cursorless/common";
-import { QueryMatch } from "web-tree-sitter";
 import { TreeSitterQuery } from "../../../../languages/TreeSitterQuery";
+import { QueryMatch } from "../../../../languages/TreeSitterQuery/QueryCapture";
 import { TEXT_FRAGMENT_CAPTURE_NAME } from "../../../../languages/captureNames";
 import { PlainTarget } from "../../../targets";
 import { TargetScope } from "../scope.types";

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

@@ -1,5 +1,4 @@
-import { QueryMatch } from "web-tree-sitter";
-import { getNodeRange } from "../../../../util/nodeSelectors";
+import { QueryMatch } from "../../../../languages/TreeSitterQuery/QueryCapture";
 
 /**
  * Gets the range of a node that is related to the scope.  For example, if the
@@ -32,9 +31,7 @@ export function getRelatedRange(
  * @returns A range or undefined if no matching capture was found
  */
 export function getCaptureRangeByName(match: QueryMatch, ...names: string[]) {
-  const relatedNode = match.captures.find((capture) =>
+  return match.captures.find((capture) =>
     names.some((name) => capture.name === name),
-  )?.node;
-
-  return relatedNode == null ? undefined : getNodeRange(relatedNode);
+  )?.range;
 }