From fce2aa8fa30a05c9db0259eba1c59d74bd587108 Mon Sep 17 00:00:00 2001
From: Sebastian McKenzie <sebmck@gmail.com>
Date: Tue, 5 May 2015 03:21:31 +0100
Subject: [PATCH] add support for esquery

---
 package.json                            |  3 +-
 src/babel/transformation/transformer.js |  2 +-
 src/babel/traversal/path/index.js       | 18 +++++++-
 src/babel/traversal/visitors.js         | 58 +++++++++++++++++++------
 4 files changed, 65 insertions(+), 16 deletions(-)

diff --git a/package.json b/package.json
index 8f1796c9dc..6db88afcfc 100644
--- a/package.json
+++ b/package.json
@@ -35,7 +35,8 @@
     "core-js": "^0.9.0",
     "debug": "^2.1.1",
     "detect-indent": "^3.0.0",
-    "estraverse": "^3.0.0",
+    "esquery": "^0.3.0",
+    "estraverse": "^4.0.0",
     "esutils": "^2.0.0",
     "fs-readdir-recursive": "^0.1.0",
     "globals": "^6.4.0",
diff --git a/src/babel/transformation/transformer.js b/src/babel/transformation/transformer.js
index 723a5c8c97..3e12abcf1d 100644
--- a/src/babel/transformation/transformer.js
+++ b/src/babel/transformation/transformer.js
@@ -44,7 +44,7 @@ export default class Transformer {
 
     //
 
-    if (!this.shouldVisit) {
+    if (!this.shouldVisit && !this.handlers.enter && !this.handlers.exit) {
       var types = Object.keys(this.handlers);
       this.shouldVisit = function (node) {
         for (var i = 0; i < types.length; i++) {
diff --git a/src/babel/traversal/path/index.js b/src/babel/traversal/path/index.js
index 8970437f56..14b801de53 100644
--- a/src/babel/traversal/path/index.js
+++ b/src/babel/traversal/path/index.js
@@ -91,6 +91,22 @@ export default class TraversalPath {
     return ourScope;
   }
 
+  /**
+   * Description
+   */
+
+  getAncestry() {
+    var ancestry = [];
+
+    var path = this.parentPath;
+    while (path) {
+      ancestry.push(path.node);
+      path = path.parentPath;
+    }
+
+    return ancestry;
+  }
+
   /**
    * Description
    */
@@ -664,7 +680,7 @@ export default class TraversalPath {
     if (!node) return;
 
     var opts = this.opts;
-    var fns  = [opts[key]];
+    var fns  = [].concat(opts[key]);
 
     if (opts[node.type]) {
       fns = fns.concat(opts[node.type][key]);
diff --git a/src/babel/traversal/visitors.js b/src/babel/traversal/visitors.js
index 4bc805f62f..8e994c1ccb 100644
--- a/src/babel/traversal/visitors.js
+++ b/src/babel/traversal/visitors.js
@@ -1,23 +1,22 @@
 import * as virtualTypes from "./path/virtual-types";
 import * as messages from "../messages";
 import * as t from "../types";
+import esquery from "esquery";
 
 export function explode(visitor, mergeConflicts) {
   // make sure there's no __esModule type since this is because we're using loose mode
   // and it sets __esModule to be enumerable on all modules :(
   delete visitor.__esModule;
 
-  // ensure visitors are objects
-  for (let nodeType in visitor) {
-    if (shouldIgnoreKey(nodeType)) continue;
-
-    var fns = visitor[nodeType];
-
-    if (typeof fns === "function") {
-      visitor[nodeType] = { enter: fns };
-    }
+  if (visitor.queries) {
+    ensureEntranceObjects(visitor.queries);
+    addQueries(visitor);
+    delete visitor.queries;
   }
 
+  // ensure visitors are objects
+  ensureEntranceObjects(visitor);
+
   // add type wrappers
   for (let nodeType in visitor) {
     if (shouldIgnoreKey(nodeType)) continue;
@@ -62,7 +61,7 @@ export function explode(visitor, mergeConflicts) {
       var existing = visitor[alias];
       if (existing) {
         if (mergeConflicts) {
-          merge(fns, existing);
+          merge(existing, fns);
         }
       } else {
         visitor[alias] = fns;
@@ -106,6 +105,39 @@ export function verify(visitor) {
   visitor._verified = true;
 }
 
+function ensureEntranceObjects(obj) {
+  for (let key in obj) {
+    if (shouldIgnoreKey(key)) continue;
+
+    var fns = obj[key];
+    if (typeof fns === "function") {
+      obj[key] = { enter: fns };
+    }
+  }
+}
+
+function addQueries(visitor) {
+  for (var selector in visitor.queries) {
+    var fns = visitor.queries[selector];
+    addSelector(visitor, selector, fns);
+  }
+}
+
+function addSelector(visitor, selector, fns) {
+  selector = esquery.parse(selector);
+
+  for (var key in fns) {
+    let fn = fns[key];
+    fns[key] = function (node) {
+      if (esquery.matches(node, selector, this.getAncestry())) {
+        return fn.apply(this, arguments);
+      }
+    };
+  }
+
+  merge(visitor, fns);
+}
+
 function wrapCheck(wrapper, fn) {
   return function () {
     if (wrapper.checkPath(this)) {
@@ -127,8 +159,8 @@ function shouldIgnoreKey(key) {
   return false;
 }
 
-function merge(visitor1, visitor2) {
-  for (var key in visitor1) {
-    visitor2[key] = (visitor2[alias] || []).concat(visitor1[key]);
+function merge(dest, src) {
+  for (var key in src) {
+    dest[key] = (dest[key] || []).concat(src[key]);
   }
 }