fix: avoid popContext on unvisited node paths (#16305)

JLHwung · web-flow · commit bd5abd5ea14a · 2024-02-28T11:31:25.000+01:00
diff --git a/packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/regression-16293/input.js b/packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/regression-16293/input.js
@@ -0,0 +1,5 @@
+async () => do {
+  await 0
+  while (0) {}
+  0
+}
diff --git a/packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/regression-16293/output.js b/packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expressions/regression-16293/output.js
@@ -0,0 +1,5 @@
+async () => await async function () {
+  await 0;
+  while (0) {}
+  return 0;
+}();
diff --git a/packages/babel-traverse/src/context.ts b/packages/babel-traverse/src/context.ts
@@ -118,9 +118,12 @@ export default class TraversalContext<S = unknown> {
 
     const visited = new WeakSet();
     let stop = false;
+    let visitIndex = 0;
 
     // visit the queue
-    for (const path of queue) {
+    for (; visitIndex < queue.length; ) {
+      const path = queue[visitIndex];
+      visitIndex++;
       path.resync();
 
       if (
@@ -154,9 +157,9 @@ export default class TraversalContext<S = unknown> {
       }
     }
 
-    // clear queue
-    for (const path of queue) {
-      path.popContext();
+    // pop contexts
+    for (let i = 0; i < visitIndex; i++) {
+      queue[i].popContext();
     }
 
     // clear queue
diff --git a/packages/babel-traverse/test/traverse.js b/packages/babel-traverse/test/traverse.js
@@ -314,6 +314,44 @@ describe("traverse", function () {
         ["EXIT", "./Bar"],
       ]);
     });
+    it("should preserve the context for those nodes that are not visited in sub-traversal", () => {
+      const code = `{ var first; function second() {} }`;
+      const ast = parse(code);
+      let contextLevel;
+      traverse(
+        ast,
+        {
+          enter(path) {
+            if (path.isFunctionDeclaration()) {
+              path.parentPath.traverse(
+                {
+                  enter(path) {
+                    if (path.isFunctionDeclaration()) {
+                      path.parentPath.traverse(
+                        {
+                          enter(path) {
+                            if (path.isVariableDeclaration()) path.stop();
+                          },
+                        },
+                        { level: 3 },
+                      );
+                      // the function declaration path should have state.level as 2
+                      // as it is defined within level 2 traversal and the node is
+                      // not visited in the next sub-traversal
+                      contextLevel = path.state.level;
+                    }
+                  },
+                },
+                { level: 2 },
+              );
+            }
+          },
+        },
+        undefined,
+        { level: 1 },
+      );
+      expect(contextLevel).toBe(2);
+    });
   });
   describe("path.stop()", () => {
     it("should stop the traversal when a grand child is stopped", () => {

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +async () => do {
 +  await 0
 +  while (0) {}
 +  0
 +}