diff --git a/lib/types.js b/lib/types.js index e4003bffb9..929c37d111 100644 --- a/lib/types.js +++ b/lib/types.js @@ -1732,4 +1732,15 @@ declare module "babel-types" { declare function cloneDeep(n: T): T; declare function removeProperties(n: T, opts: ?{}): void; declare function removePropertiesDeep(n: T, opts: ?{}): T; + declare type TraversalAncestors = Array<{ + node: BabelNode, + key: string, + index?: number, + }>; + declare type TraversalHandler = (BabelNode, TraversalAncestors, T) => void; + declare type TraversalHandlers = { + enter?: TraversalHandler, + exit?: TraversalHandler, + }; + declare function traverse(n: BabelNode, TraversalHandler | TraversalHandlers, state?: T): void; } diff --git a/packages/babel-types/src/index.js b/packages/babel-types/src/index.js index 72bdb99560..172e2d4027 100644 --- a/packages/babel-types/src/index.js +++ b/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 diff --git a/packages/babel-types/src/traverse.js b/packages/babel-types/src/traverse.js new file mode 100644 index 0000000000..c4983d1685 --- /dev/null +++ b/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 = (BabelNode, TraversalAncestors, T) => void; +export type TraversalHandlers = { + enter?: TraversalHandler, + exit?: TraversalHandler, +}; + +/** + * 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( + node: BabelNode, + handlers: TraversalHandler | TraversalHandlers, + 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); +} diff --git a/scripts/generate-interfaces.js b/scripts/generate-interfaces.js index 5eeaac1d65..88552d5611 100644 --- a/scripts/generate-interfaces.js +++ b/scripts/generate-interfaces.js @@ -140,7 +140,19 @@ lines.push( `declare function clone(n: T): T;`, `declare function cloneDeep(n: T): T;`, `declare function removeProperties(n: T, opts: ?{}): void;`, - `declare function removePropertiesDeep(n: T, opts: ?{}): T;` + `declare function removePropertiesDeep(n: T, opts: ?{}): T;`, + `declare type TraversalAncestors = Array<{ + node: BabelNode, + key: string, + index?: number, + }>; + declare type TraversalHandler = (BabelNode, TraversalAncestors, T) => void; + declare type TraversalHandlers = { + enter?: TraversalHandler, + exit?: TraversalHandler, + };`.replace(/(^|\n) {2}/g, "$1"), + // eslint-disable-next-line + `declare function traverse(n: BabelNode, TraversalHandler | TraversalHandlers, state?: T): void;` ); for (const type in t.FLIPPED_ALIAS_KEYS) {