Restore traversal context after enter / traverse

packages/babel-plugin-proposal-async-generator-functions/test/fixtures/regression/13801/input.js

@@ -0,0 +1,11 @@
+async function main() {
+  async () => {
+    // IIFE: required for babel to crash
+    for await (const string of async_iterable) {
+      // for await: required for babel to crash
+      console.log(string);
+    }
+  };
+
+  const [one] = [1]; // array destructuring: required for babel to crash
+}

packages/babel-plugin-proposal-async-generator-functions/test/fixtures/regression/13801/options.json

@@ -0,0 +1,8 @@
+{
+  "plugins": [
+    "transform-destructuring",
+    "transform-regenerator",
+    "proposal-async-generator-functions"
+  ],
+  "externalHelpers": false
+}

packages/babel-plugin-proposal-async-generator-functions/test/fixtures/regression/13801/output.js

@@ -0,0 +1,91 @@
+function _asyncIterator(iterable) { var method; if (typeof Symbol !== "undefined") { if (Symbol.asyncIterator) method = iterable[Symbol.asyncIterator]; if (method == null && Symbol.iterator) method = iterable[Symbol.iterator]; } if (method == null) method = iterable["@@asyncIterator"]; if (method == null) method = iterable["@@iterator"]; if (method == null) throw new TypeError("Object is not async iterable"); return method.call(iterable); }
+
+function main() {
+  var one;
+  return regeneratorRuntime.async(function main$(_context2) {
+    while (1) switch (_context2.prev = _context2.next) {
+      case 0:
+        () => {
+          var _iteratorAbruptCompletion, _didIteratorError, _iteratorError, _iterator, _step, string;
+
+          return regeneratorRuntime.async(function _callee$(_context) {
+            while (1) switch (_context.prev = _context.next) {
+              case 0:
+                // IIFE: required for babel to crash
+                _iteratorAbruptCompletion = false;
+                _didIteratorError = false;
+                _context.prev = 2;
+                _iterator = _asyncIterator(async_iterable);
+
+              case 4:
+                _context.next = 6;
+                return regeneratorRuntime.awrap(_iterator.next());
+
+              case 6:
+                if (!(_iteratorAbruptCompletion = !(_step = _context.sent).done)) {
+                  _context.next = 12;
+                  break;
+                }
+
+                string = _step.value;
+                // for await: required for babel to crash
+                console.log(string);
+
+              case 9:
+                _iteratorAbruptCompletion = false;
+                _context.next = 4;
+                break;
+
+              case 12:
+                _context.next = 18;
+                break;
+
+              case 14:
+                _context.prev = 14;
+                _context.t0 = _context["catch"](2);
+                _didIteratorError = true;
+                _iteratorError = _context.t0;
+
+              case 18:
+                _context.prev = 18;
+                _context.prev = 19;
+
+                if (!(_iteratorAbruptCompletion && _iterator.return != null)) {
+                  _context.next = 23;
+                  break;
+                }
+
+                _context.next = 23;
+                return regeneratorRuntime.awrap(_iterator.return());
+
+              case 23:
+                _context.prev = 23;
+
+                if (!_didIteratorError) {
+                  _context.next = 26;
+                  break;
+                }
+
+                throw _iteratorError;
+
+              case 26:
+                return _context.finish(23);
+
+              case 27:
+                return _context.finish(18);
+
+              case 28:
+              case "end":
+                return _context.stop();
+            }
+          }, null, null, [[2, 14, 18, 28], [19,, 23, 27]], Promise);
+        };
+
+        one = 1; // array destructuring: required for babel to crash
+
+      case 2:
+      case "end":
+        return _context2.stop();
+    }
+  }, null, null, null, Promise);
+}

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

@@ -61,6 +61,14 @@ export function isDenylisted(this: NodePath): boolean {
 // TODO: Remove in Babel 8
 export { isDenylisted as isBlacklisted };
 
+function restoreContext(path: NodePath, context: TraversalContext) {
+  if (path.context !== context) {
+    path.context = context;
+    path.state = context.state;
+    path.opts = context.opts;
+  }
+}
+
 export function visit(this: NodePath): boolean {
   if (!this.node) {
     return false;
@@ -74,15 +82,17 @@ export function visit(this: NodePath): boolean {
     return false;
   }
 
-  // Note: We need to check "this.shouldSkip" twice because
-  // the visitor can set it to true. Usually .shouldSkip is false
+  const currentContext = this.context;
+  // Note: We need to check "this.shouldSkip" first because
+  // another visitor can set it to true. Usually .shouldSkip is false
   // before calling the enter visitor, but it can be true in case of
   // a requeued node (e.g. by .replaceWith()) that is then marked
   // with .skip().
-  if (this.shouldSkip || this.call("enter") || this.shouldSkip) {
+  if (this.shouldSkip || this.call("enter")) {
     this.debug("Skip...");
     return this.shouldStop;
   }
+  restoreContext(this, currentContext);
 
   this.debug("Recursing into...");
   traverse.node(
@@ -94,6 +104,8 @@ export function visit(this: NodePath): boolean {
     this.skipKeys,
   );
 
+  restoreContext(this, currentContext);
+
   this.call("exit");
 
   return this.shouldStop;

packages/babel-traverse/test/traverse.js

@@ -216,4 +216,63 @@ describe("traverse", function () {
       expect(visited).toBe(true);
     });
   });
+  describe("path.visit()", () => {
+    it("should preserve traversal context after enter hook is executed", () => {
+      const ast = parse("{;}");
+      // The test initiates a sub-traverse from program. When the `enter` hook of BlockStatement
+      // is called, the unshiftContainer will change the traversal context of the BlockStatement
+      // to the one of Program which has an EmptyStatement visitor. If the traversal context
+      // is not restored after the `enter` hook is executed, the `EmptyStatement` visitor will
+      // be run twice: one in the sub-traverse and the other in the top level traverse.
+      let emptyStatementVisitedCounter = 0;
+      traverse(ast, {
+        noScope: true,
+        Program(path) {
+          path.traverse({
+            noScope: true,
+            BlockStatement: {
+              enter(path) {
+                path.parentPath.unshiftContainer("body", [t.numericLiteral(0)]);
+              },
+            },
+          });
+        },
+        EmptyStatement() {
+          ++emptyStatementVisitedCounter;
+        },
+      });
+      expect(emptyStatementVisitedCounter).toBe(1);
+    });
+    it("should preserve traversal context after visitor is executed", () => {
+      const ast = parse("{;}");
+      // The test initiates a sub-traverse from program. During the BlockStatement is traversed,
+      // the EmptyStatement visitor will be called and the unshiftContainer will change the
+      // traversal context of the BlockStatement to that of Program which has an EmptyStatement
+      // visitor. If the traversal context is not restored after `enter` hook is executed,
+      // the `BlockStatement:exit` visitor will be run twice: one in the sub-traverse and the other
+      // in the top level traverse.
+      let blockStatementVisitedCounter = 0;
+      traverse(ast, {
+        noScope: true,
+        Program(path) {
+          path.traverse({
+            noScope: true,
+            EmptyStatement: {
+              enter(path) {
+                path.parentPath.parentPath.unshiftContainer("body", [
+                  t.numericLiteral(0),
+                ]);
+              },
+            },
+          });
+        },
+        BlockStatement: {
+          exit() {
+            ++blockStatementVisitedCounter;
+          },
+        },
+      });
+      expect(blockStatementVisitedCounter).toBe(1);
+    });
+  });
 });