From 2e0afe9af19b56d3bfb055696c4d337ea615a0e3 Mon Sep 17 00:00:00 2001
From: Pokey Rule <pokey.rule@gmail.com>
Date: Tue, 27 Apr 2021 17:50:08 +0100
Subject: [PATCH] Simple node matcher improvements

---
 src/Types.ts          |  1 +
 src/actions.ts        |  2 +-
 src/nodeMatchers.ts   | 56 ++++++++++++++++---------------------------
 src/processTargets.ts | 15 ++++++------
 4 files changed, 30 insertions(+), 44 deletions(-)

diff --git a/src/Types.ts b/src/Types.ts
index cac193fea..d4fa6c370 100644
--- a/src/Types.ts
+++ b/src/Types.ts
@@ -50,6 +50,7 @@ export type ScopeType =
   | "class"
   | "arrowFunction"
   | "namedFunction"
+  | "pair"
   | "ifStatement";
 export type PieceType = "subtoken" | "character";
 
diff --git a/src/actions.ts b/src/actions.ts
index f2bb2a958..1c71bd381 100644
--- a/src/actions.ts
+++ b/src/actions.ts
@@ -45,10 +45,10 @@ class Actions {
     }
     const editor = editors[0];
 
-    editor.selections = targets.map((target) => target.selection.selection);
     if (editor.viewColumn != null) {
       await commands.executeCommand(columnFocusCommands[editor.viewColumn]);
     }
+    editor.selections = targets.map((target) => target.selection.selection);
     editor.revealRange(editor.selections[0]);
   };
 
diff --git a/src/nodeMatchers.ts b/src/nodeMatchers.ts
index 786b4ebdf..6fcc614b1 100644
--- a/src/nodeMatchers.ts
+++ b/src/nodeMatchers.ts
@@ -1,28 +1,14 @@
 import { SyntaxNode } from "tree-sitter";
 import { Position, Selection } from "vscode";
 
-type NodeMatcher = (node: SyntaxNode) => boolean;
-
-type NodeSelectionExtractor = (
-  node: SyntaxNode,
-  isInside: boolean
-) => Selection | null;
-
-interface NodeDescriptor {
-  matcher: NodeMatcher;
-  extractor: NodeSelectionExtractor;
-}
+type NodeMatcher = (node: SyntaxNode, isInside: boolean) => Selection | null;
 
 function hasType(...typeNames: string[]): NodeMatcher {
-  return (node: SyntaxNode) => typeNames.includes(node.type);
+  return (node: SyntaxNode, isInside: boolean) =>
+    typeNames.includes(node.type) ? simpleSelectionExtractor(node) : null;
 }
 
-const makeSimpleNodeDescriptor = (...typeNames: string[]) => ({
-  matcher: hasType("pair"),
-  extractor: simpleSelectionExtractor,
-});
-
-function simpleSelectionExtractor(node: SyntaxNode, isInside: boolean) {
+function simpleSelectionExtractor(node: SyntaxNode) {
   return new Selection(
     new Position(node.startPosition.row, node.startPosition.column),
     new Position(node.endPosition.row, node.endPosition.column)
@@ -30,25 +16,25 @@ function simpleSelectionExtractor(node: SyntaxNode, isInside: boolean) {
 }
 
 const nodeMatchers = {
-  class: makeSimpleNodeDescriptor("class_declaration"),
-  arrowFunction: makeSimpleNodeDescriptor("arrow_function"),
-  pair: {
-    matcher: hasType("pair"),
-    extractor: (node: SyntaxNode, isInside: boolean) => {
-      return new Selection(
-        new Position(node.startPosition.row, node.startPosition.column),
-        new Position(node.endPosition.row, node.endPosition.column)
-      );
-    },
+  class: hasType("class_declaration"),
+  arrowFunction: hasType("arrow_function"),
+  pair: (node: SyntaxNode, isInside: boolean) => {
+    if (node.type !== "pair") {
+      return null;
+    }
+    return new Selection(
+      new Position(node.startPosition.row, node.startPosition.column),
+      new Position(node.endPosition.row, node.endPosition.column)
+    );
   },
-  namedFunction(node: SyntaxNode) {
+  namedFunction(node: SyntaxNode, isInside: boolean) {
     // Simple case, eg
     // function foo() {}
     if (
       node.type === "function_declaration" ||
       node.type === "method_definition"
     ) {
-      return true;
+      return simpleSelectionExtractor(node);
     }
 
     // Class property defined as field definition with arrow
@@ -61,14 +47,14 @@ const nodeMatchers = {
       // @ts-ignore
       node.valueNode.type === "arrow_function"
     ) {
-      return true;
+      return simpleSelectionExtractor(node);
     }
 
     // eg:
     // const foo = () => "hello"
     if (node.type === "lexical_declaration") {
       if (node.namedChildCount !== 1) {
-        return false;
+        return null;
       }
 
       const child = node.firstNamedChild!;
@@ -78,13 +64,13 @@ const nodeMatchers = {
         // @ts-ignore
         child.valueNode.type === "arrow_function"
       ) {
-        return true;
+        return simpleSelectionExtractor(node);
       }
     }
 
-    return false;
+    return null;
   },
-  ifStatement: makeSimpleNodeDescriptor("if_statement"),
+  ifStatement: hasType("if_statement"),
 };
 
 export default nodeMatchers;
diff --git a/src/processTargets.ts b/src/processTargets.ts
index c75e82169..49e396fe7 100644
--- a/src/processTargets.ts
+++ b/src/processTargets.ts
@@ -119,8 +119,12 @@ function getSelectionsFromMark(
 function transformSelection(
   context: ProcessedTargetsContext,
   transformation: Transformation,
+  position: Position,
   selection: SelectionWithEditor
 ): SelectionWithEditor {
+  const isInside =
+    position === "end" || position === "start" || position === "inside";
+
   switch (transformation.type) {
     case "identity":
       return selection;
@@ -132,16 +136,11 @@ function transformSelection(
       const nodeMatcher = nodeMatchers[transformation.scopeType];
 
       while (node != null) {
-        if (nodeMatcher(node)) {
+        const matchedSelection = nodeMatcher(node, isInside);
+        if (matchedSelection != null) {
           return {
             editor: selection.editor,
-            selection: new Selection(
-              new vscode.Position(
-                node.startPosition.row,
-                node.startPosition.column
-              ),
-              new vscode.Position(node.endPosition.row, node.endPosition.column)
-            ),
+            selection: matchedSelection,
           };
         }
         console.log(node.type);