Add a new utility for traversing the AST.

lib/types.js

@@ -1732,4 +1732,15 @@ declare module "babel-types" {
   declare function cloneDeep<T>(n: T): T;
   declare function removeProperties<T>(n: T, opts: ?{}): void;
   declare function removePropertiesDeep<T>(n: T, opts: ?{}): T;
+  declare type TraversalAncestors = Array<{
+    node: BabelNode,
+    key: string,
+    index?: number,
+  }>;
+  declare type TraversalHandler<T> = (BabelNode, TraversalAncestors, T) => void;
+  declare type TraversalHandlers<T> = {
+    enter?: TraversalHandler<T>,
+    exit?: TraversalHandler<T>,
+  };
+  declare function traverse<T>(n: BabelNode, TraversalHandler<T> | TraversalHandlers<T>, state?: T): void;
 }

packages/babel-types/src/index.js

@@ -66,6 +66,9 @@ export { VISITOR_KEYS, ALIAS_KEYS, NODE_FIELDS, BUILDER_KEYS, DEPRECATED_KEYS };
 import * as _react from "./react";
 export { _react as react };
 
+import { traverse, traverseFast } from "./traverse";
+export { traverse, traverseFast };
+
 /**
  * Registers `is[Type]` and `assert[Type]` for all types.
  */
@@ -534,36 +537,6 @@ export function isNode(node?): boolean {
 toFastProperties(t);
 toFastProperties(t.VISITOR_KEYS);
 
-/**
- * A prefix AST traversal implementation implementation.
- */
-
-export function traverseFast(
-  node: Node,
-  enter: (node: Node) => void,
-  opts?: Object,
-) {
-  if (!node) return;
-
-  const keys = t.VISITOR_KEYS[node.type];
-  if (!keys) return;
-
-  opts = opts || {};
-  enter(node, opts);
-
-  for (const key of keys) {
-    const subNode = node[key];
-
-    if (Array.isArray(subNode)) {
-      for (const node of subNode) {
-        traverseFast(node, enter, opts);
-      }
-    } else {
-      traverseFast(subNode, enter, opts);
-    }
-  }
-}
-
 const CLEAR_KEYS: Array = ["tokens", "start", "end", "loc", "raw", "rawValue"];
 
 const CLEAR_KEYS_PLUS_COMMENTS: Array = t.COMMENT_KEYS

packages/babel-types/src/traverse.js

@@ -0,0 +1,99 @@
+import { VISITOR_KEYS } from "./index";
+
+/**
+ * A prefix AST traversal implementation meant for simple searching
+ * and processing.
+ */
+export function traverseFast(
+  node: Node,
+  enter: (node: Node) => void,
+  opts?: Object,
+) {
+  if (!node) return;
+
+  const keys = VISITOR_KEYS[node.type];
+  if (!keys) return;
+
+  opts = opts || {};
+  enter(node, opts);
+
+  for (const key of keys) {
+    const subNode = node[key];
+
+    if (Array.isArray(subNode)) {
+      for (const node of subNode) {
+        traverseFast(node, enter, opts);
+      }
+    } else {
+      traverseFast(subNode, enter, opts);
+    }
+  }
+}
+
+export type TraversalAncestors = Array<{
+  node: BabelNode,
+  key: string,
+  index?: number,
+}>;
+export type TraversalHandler<T> = (BabelNode, TraversalAncestors, T) => void;
+export type TraversalHandlers<T> = {
+  enter?: TraversalHandler<T>,
+  exit?: TraversalHandler<T>,
+};
+
+/**
+ * A general AST traversal with both prefix and postfix handlers, and a
+ * state object. Exposes ancestry data to each handler so that more complex
+ * AST data can be taken into account.
+ */
+export function traverse<T>(
+  node: BabelNode,
+  handlers: TraversalHandler<T> | TraversalHandlers<T>,
+  state?: T,
+) {
+  if (typeof handlers === "function") {
+    handlers = { enter: handlers };
+  }
+
+  const { enter, exit } = handlers;
+
+  traverseSimpleImpl(node, enter, exit, state, []);
+}
+function traverseSimpleImpl(node, enter, exit, state, ancestors) {
+  const keys = VISITOR_KEYS[node.type];
+  if (!keys) return;
+
+  if (enter) enter(node, ancestors, state);
+
+  for (const key of keys) {
+    const subNode = node[key];
+
+    if (Array.isArray(subNode)) {
+      for (let i = 0; i < subNode.length; i++) {
+        const child = subNode[i];
+        if (!child) continue;
+
+        ancestors.push({
+          node,
+          key,
+          index: i,
+        });
+
+        traverseSimpleImpl(child, enter, exit, state, ancestors);
+
+        ancestors.pop();
+      }
+    } else if (subNode) {
+      ancestors.push({
+        node,
+        key,
+      });
+
+      traverseSimpleImpl(subNode, enter, exit, state, ancestors);
+
+      ancestors.pop();
+    }
+  }
+
+  if (exit) exit(node, ancestors, state);
+}

scripts/generate-interfaces.js

@@ -140,7 +140,19 @@ lines.push(
   `declare function clone<T>(n: T): T;`,
   `declare function cloneDeep<T>(n: T): T;`,
   `declare function removeProperties<T>(n: T, opts: ?{}): void;`,
-  `declare function removePropertiesDeep<T>(n: T, opts: ?{}): T;`
+  `declare function removePropertiesDeep<T>(n: T, opts: ?{}): T;`,
+  `declare type TraversalAncestors = Array<{
+    node: BabelNode,
+    key: string,
+    index?: number,
+  }>;
+  declare type TraversalHandler<T> = (BabelNode, TraversalAncestors, T) => void;
+  declare type TraversalHandlers<T> = {
+    enter?: TraversalHandler<T>,
+    exit?: TraversalHandler<T>,
+  };`.replace(/(^|\n) {2}/g, "$1"),
+  // eslint-disable-next-line
+  `declare function traverse<T>(n: BabelNode, TraversalHandler<T> | TraversalHandlers<T>, state?: T): void;`
 );
 
 for (const type in t.FLIPPED_ALIAS_KEYS) {