Performance improvement (single AST traversal run)

src/action-providers.ts

@@ -1,9 +1,15 @@
 import * as vscode from "vscode";
 
 import { createSelectionFromVSCode } from "./editor/adapters/vscode-editor";
-import { Selection } from "./editor/selection";
-import { RefactoringWithActionProvider } from "./types";
+import {
+  RefactoringWithActionProvider,
+  ActionProvider,
+  LegacyActionProvider,
+  isRefactoringWithActionProvider,
+  isRefactoringWithLegacyActionProvider
+} from "./types";
 import * as t from "./ast";
+import { Selection } from "./editor/selection";
 
 export { RefactoringActionProvider };
 
@@ -21,18 +27,103 @@ class RefactoringActionProvider implements vscode.CodeActionProvider {
     const ast = t.parse(document.getText());
     const selection = createSelectionFromVSCode(range);
 
-    return this.refactorings
-      .filter(refactoring => this.canPerform(refactoring, ast, selection))
-      .map(refactoring => this.buildCodeActionFor(refactoring));
+    return [
+      ...this.findApplicableRefactorings(ast, selection),
+      ...this.findApplicableLegacyRefactorings(ast, selection)
+    ].map(refactoring => this.buildCodeActionFor(refactoring));
+  }
+
+  private findApplicableLegacyRefactorings(ast: t.File, selection: Selection) {
+    const legacyRefactorings = this.refactorings.filter(
+      isRefactoringWithLegacyActionProvider
+    );
+    const applicableLegacyRefactorings = legacyRefactorings.filter(
+      refactoring => this.canPerformLegacy(refactoring, ast, selection)
+    );
+    return applicableLegacyRefactorings;
+  }
+
+  private findApplicableRefactorings(
+    ast: t.File,
+    selection: Selection
+  ): RefactoringWithActionProvider<ActionProvider>[] {
+    const refactorings = this.refactorings.filter(
+      isRefactoringWithActionProvider
+    );
+
+    const applicableRefactorings: RefactoringWithActionProvider<
+      ActionProvider
+    >[] = [];
+
+    t.traverseAST(ast, {
+      enter: (path: t.NodePath<any>) => {
+        refactorings.forEach(refactoring =>
+          this.visitAndCheckApplicability(
+            refactoring,
+            path,
+            selection,
+            (
+              path: t.NodePath,
+              refactoring: RefactoringWithActionProvider<ActionProvider>
+            ) => {
+              if (refactoring.actionProvider.updateMessage) {
+                refactoring.actionProvider.message = refactoring.actionProvider.updateMessage(
+                  path
+                );
+              }
+
+              applicableRefactorings.push(refactoring);
+            }
+          )
+        );
+      }
+    });
+
+    return applicableRefactorings;
+  }
+
+  private visitAndCheckApplicability(
+    refactoring: RefactoringWithActionProvider<ActionProvider>,
+    path: t.NodePath<any>,
+    selection: Selection,
+    whenApplicable: (
+      matchedPath: t.NodePath<any>,
+      refactoring: RefactoringWithActionProvider<ActionProvider>
+    ) => void
+  ) {
+    const visitor: t.Visitor = refactoring.actionProvider.createVisitor(
+      selection,
+      path => whenApplicable(path, refactoring),
+      refactoring
+    );
+
+    this.visit(visitor, path);
+  }
+
+  private visit(visitor: any, path: t.NodePath<any>) {
+    const node: t.Node = path.node;
+
+    try {
+      if (typeof visitor[node.type] === "function") {
+        visitor[node.type](path);
+      } else if (typeof visitor[node.type] === "object") {
+        visitor[node.type].enter(path);
+      }
+    } catch (_) {
+      // Silently fail, we don't care why it failed (e.g. code can't be parsed).
+    }
   }
 
-  private canPerform(
-    refactoring: RefactoringWithActionProvider,
+  private canPerformLegacy(
+    refactoring: RefactoringWithActionProvider<LegacyActionProvider>,
     ast: t.AST,
     selection: Selection
   ) {
     try {
-      return refactoring.actionProvider.canPerform(ast, selection);
+      return (
+        typeof refactoring.actionProvider.canPerform === "function" &&
+        refactoring.actionProvider.canPerform(ast, selection)
+      );
     } catch (_) {
       // Silently fail, we don't care why it failed (e.g. code can't be parsed).
       return false;

src/refactorings/simplify-ternary/index.ts

@@ -1,4 +1,8 @@
-import { simplifyTernary, hasTernaryToSimplify } from "./simplify-ternary";
+import {
+  simplifyTernary,
+  createTernaryToSimplifyVisitor
+} from "./simplify-ternary";
+import * as t from "../../ast";
 
 import { RefactoringWithActionProvider } from "../../types";
 
@@ -10,7 +14,10 @@ const config: RefactoringWithActionProvider = {
   },
   actionProvider: {
     message: "Simplify ternary",
-    canPerform: hasTernaryToSimplify
+    createVisitor: createTernaryToSimplifyVisitor,
+    updateMessage(path: t.NodePath): string {
+      return `Simplify ternary (for demo purposes only: ${path.node.type})`;
+    }
   }
 };
 

src/refactorings/simplify-ternary/simplify-ternary.ts

@@ -2,7 +2,7 @@ import { Editor, Code, ErrorReason } from "../../editor/editor";
 import { Selection } from "../../editor/selection";
 import * as t from "../../ast";
 
-export { simplifyTernary, hasTernaryToSimplify };
+export { simplifyTernary, createVisitor as createTernaryToSimplifyVisitor };
 
 async function simplifyTernary(
   code: Code,
@@ -19,14 +19,6 @@ async function simplifyTernary(
   await editor.write(updatedCode.code);
 }
 
-function hasTernaryToSimplify(ast: t.AST, selection: Selection): boolean {
-  let result = false;
-
-  t.traverseAST(ast, createVisitor(selection, () => (result = true)));
-
-  return result;
-}
-
 function updateCode(ast: t.AST, selection: Selection): t.Transformed {
   return t.transformAST(
     ast,

src/types.ts

@@ -1,8 +1,16 @@
 import { Code, Editor } from "./editor/editor";
 import { Selection } from "./editor/selection";
-import { AST } from "./ast";
+import { AST, Visitor, NodePath } from "./ast";
 
-export { Refactoring, RefactoringWithActionProvider, Operation };
+export {
+  Refactoring,
+  RefactoringWithActionProvider,
+  Operation,
+  ActionProvider,
+  LegacyActionProvider,
+  isRefactoringWithActionProvider,
+  isRefactoringWithLegacyActionProvider
+};
 
 interface Refactoring {
   command: {
@@ -11,21 +19,54 @@ interface Refactoring {
   };
 }
 
-interface RefactoringWithActionProvider extends Refactoring {
+interface ActionProvider {
+  message: string;
+  isPreferred?: boolean;
+  createVisitor: (
+    selection: Selection,
+    onMatch: (path: NodePath) => void,
+    refactoring: RefactoringWithActionProvider
+  ) => Visitor;
+  updateMessage?: (path: NodePath) => string;
+}
+
+interface LegacyActionProvider {
+  message: string;
+  isPreferred?: boolean;
+  canPerform: (ast: AST, selection: Selection) => boolean;
+}
+
+interface RefactoringWithActionProvider<
+  ActionProviderType = ActionProvider | LegacyActionProvider
+> extends Refactoring {
   command: {
     key: string;
     title: string;
     operation: Operation;
   };
-  actionProvider: {
-    message: string;
-    canPerform: (ast: AST, selection: Selection) => boolean;
-    isPreferred?: boolean;
-  };
+  actionProvider: ActionProviderType;
 }
 
 type Operation = (
   code: Code,
   selection: Selection,
   write: Editor
 ) => Promise<void>;
+
+function isLegacyActionProvider(
+  actionProvider: ActionProvider | LegacyActionProvider
+): actionProvider is LegacyActionProvider {
+  return (actionProvider as LegacyActionProvider).canPerform !== undefined;
+}
+
+function isRefactoringWithLegacyActionProvider(
+  refactoring: RefactoringWithActionProvider
+): refactoring is RefactoringWithActionProvider<LegacyActionProvider> {
+  return isLegacyActionProvider(refactoring.actionProvider);
+}
+
+function isRefactoringWithActionProvider(
+  refactoring: RefactoringWithActionProvider
+): refactoring is RefactoringWithActionProvider<ActionProvider> {
+  return !isLegacyActionProvider(refactoring.actionProvider);
+}