fix: forward stop signal to parent path (#14105)

packages/babel-traverse/src/index.ts

@@ -1,4 +1,3 @@
-import TraversalContext from "./context";
 import * as visitors from "./visitors";
 import { VISITOR_KEYS, removeProperties, traverseFast } from "@babel/types";
 import type * as t from "@babel/types";
@@ -6,6 +5,7 @@ import * as cache from "./cache";
 import type NodePath from "./path";
 import type { default as Scope, Binding } from "./scope";
 import type { Visitor } from "./types";
+import { traverseNode } from "./traverse-node";
 
 export type { Visitor, Binding };
 export { default as NodePath } from "./path";
@@ -64,7 +64,7 @@ function traverse(
 
   visitors.explode(opts);
 
-  traverse.node(parent, opts, scope, state, parentPath);
+  traverseNode(parent, opts, scope, state, parentPath);
 }
 
 export default traverse;
@@ -82,17 +82,11 @@ traverse.node = function (
   opts: TraverseOptions,
   scope?: Scope,
   state?: any,
-  parentPath?: NodePath,
+  path?: NodePath,
-  skipKeys?,
+  skipKeys?: string[],
 ) {
-  const keys = VISITOR_KEYS[node.type];
+  traverseNode(node, opts, scope, state, path, skipKeys);
-  if (!keys) return;
+  // traverse.node always returns undefined
-
-  const context = new TraversalContext(scope, opts, state, parentPath);
-  for (const key of keys) {
-    if (skipKeys && skipKeys[key]) continue;
-    if (context.visit(node, key)) return;
-  }
 };
 
 traverse.clearNode = function (node: t.Node, opts?) {

packages/babel-traverse/src/path/context.ts

@@ -1,6 +1,6 @@
 // This file contains methods responsible for maintaining a TraversalContext.
 
-import traverse from "../index";
+import { traverseNode } from "../traverse-node";
 import { SHOULD_SKIP, SHOULD_STOP } from "./index";
 import type TraversalContext from "../context";
 import type NodePath from "./index";
@@ -95,7 +95,7 @@ export function visit(this: NodePath): boolean {
   restoreContext(this, currentContext);
 
   this.debug("Recursing into...");
-  traverse.node(
+  this.shouldStop = traverseNode(
     this.node,
     this.opts,
     this.scope,

packages/babel-traverse/src/traverse-node.ts

@@ -0,0 +1,40 @@
+import TraversalContext from "./context";
+import type { TraverseOptions } from "./index";
+import type NodePath from "./path";
+import type Scope from "./scope";
+import type * as t from "@babel/types";
+import { VISITOR_KEYS } from "@babel/types";
+
+/**
+ * Traverse the children of given node
+ * @param {Node} node
+ * @param {TraverseOptions} opts The traverse options used to create a new traversal context
+ * @param {scope} scope A traversal scope used to create a new traversal context. When opts.noScope is true, scope should not be provided
+ * @param {any} state A user data storage provided as the second callback argument for traversal visitors
+ * @param {NodePath} path A NodePath of given node
+ * @param {string[]} skipKeys A list of key names that should be skipped during traversal. The skipKeys are applied to every descendants
+ * @returns {boolean} Whether the traversal stops early
+
+ * @note This function does not visit the given `node`.
+ */
+export function traverseNode(
+  node: t.Node,
+  opts: TraverseOptions,
+  scope?: Scope,
+  state?: any,
+  path?: NodePath,
+  skipKeys?: string[],
+): boolean {
+  const keys = VISITOR_KEYS[node.type];
+  if (!keys) return false;
+
+  const context = new TraversalContext(scope, opts, state, path);
+  for (const key of keys) {
+    if (skipKeys && skipKeys[key]) continue;
+    if (context.visit(node, key)) {
+      return true;
+    }
+  }
+
+  return false;
+}

packages/babel-traverse/test/traverse.js

@@ -277,4 +277,64 @@ describe("traverse", function () {
       expect(blockStatementVisitedCounter).toBe(1);
     });
   });
+  describe("path.stop()", () => {
+    it("should stop the traversal when a grand child is stopped", () => {
+      const ast = parse("f;g;");
+
+      let visitedCounter = 0;
+      traverse(ast, {
+        noScope: true,
+        Identifier(path) {
+          visitedCounter += 1;
+          path.stop();
+        },
+      });
+
+      expect(visitedCounter).toBe(1);
+    });
+
+    it("can be reverted in the exit listener of the parent whose child is stopped", () => {
+      const ast = parse("f;g;");
+
+      let visitedCounter = 0;
+      traverse(ast, {
+        noScope: true,
+        Identifier(path) {
+          visitedCounter += 1;
+          path.stop();
+        },
+        ExpressionStatement: {
+          exit(path) {
+            path.shouldStop = false;
+            path.shouldSkip = false;
+          },
+        },
+      });
+
+      expect(visitedCounter).toBe(2);
+    });
+
+    it("should not affect root traversal", () => {
+      const ast = parse("f;g;");
+
+      let visitedCounter = 0;
+      let programShouldStop;
+      traverse(ast, {
+        noScope: true,
+        Program(path) {
+          path.traverse({
+            noScope: true,
+            Identifier(path) {
+              visitedCounter += 1;
+              path.stop();
+            },
+          });
+          programShouldStop = path.shouldStop;
+        },
+      });
+
+      expect(visitedCounter).toBe(1);
+      expect(programShouldStop).toBe(false);
+    });
+  });
 });