Ensure AsyncContext is correctly exited when a for-await loop i…

…s abruptly exited (via `break` or `continue`).

The simple solution is to just move the context swap at the end of the loop body into a `finally` block, so that it runs even if the loop exits via a different code path.  Note that if the abrupt exit is a `return`, then swap runs twice, but this is okay since the exit operation is idempotent.

PiperOrigin-RevId: 715856725

src/com/google/javascript/jscomp/InstrumentAsyncContext.java

@@ -178,8 +178,11 @@ void instrumentForAwaitOf(NodeTraversal t, Node n, Node parent) {
     //     becomes
     // for await (const x of swap(expr())) { // exit
     //   swap(0, 1); // reenter
-    //   use(x);
-    //   swap(); // exit
+    //   try {
+    //     use(x);
+    //   } finally {
+    //     swap(); // exit
+    //   }
     // }
     // swap(0, 1); // reenter
 
@@ -191,10 +194,13 @@ void instrumentForAwaitOf(NodeTraversal t, Node n, Node parent) {
 
     Node body = n.getLastChild();
     checkState(body.isBlock()); // NOTE: IRFactory normalizes non-block `for` bodies
-
-    // Add reenter/exit calls to start/end of block
-    body.addChildToFront(IR.exprResult(createReenter().srcrefTreeIfMissing(arg)));
-    body.addChildToBack(IR.exprResult(createExit().srcrefTreeIfMissing(arg)));
+    body.replaceWith(placeholder);
+    placeholder.replaceWith(
+        IR.block(
+            IR.exprResult(createReenter().srcrefTreeIfMissing(arg)),
+            IR.tryFinally(
+                body, //
+                IR.block(IR.exprResult(createExit().srcrefTreeIfMissing(arg))))));
 
     // Add reenter call after for-await-of
     if (!parent.isBlock()) {

test/com/google/javascript/jscomp/InstrumentAsyncContextTest.java

@@ -124,7 +124,9 @@ public void testAwait() {
   }
 
   // BEFORE: for await (const x of EXPR) { BODY }
-  //  AFTER: for await (const x of swap(EXPR)) { swap(0, 1); BODY; swap() } swap(0, 1)
+  //  AFTER: for await (const x of swap(EXPR)) {
+  //           swap(0, 1); try { BODY } finally { swap() }
+  //         } swap(0, 1)
   //
   // Explanation:
   //   * each iteration of the for-await loop is effectively an await; the top of the loop body is
@@ -153,8 +155,11 @@ public void testForAwait() {
             "  try {",
             "    for await (const x of $jscomp$swapContext(gen())) {",
             "      $jscomp$swapContext(0, 1);",
-            "      use(x);",
-            "      $jscomp$swapContext();",
+            "      try {",
+            "        use(x);",
+            "      } finally {",
+            "        $jscomp$swapContext();",
+            "      }",
             "    }",
             "    $jscomp$swapContext(0, 1);",
             "  } finally {",
@@ -187,8 +192,11 @@ public void testForAwaitNotInBlock() {
             "    if (1) {",
             "      for await (const x of $jscomp$swapContext(gen())) {",
             "        $jscomp$swapContext(0, 1);",
-            "        use(x);",
-            "        $jscomp$swapContext();",
+            "        try {",
+            "          use(x);",
+            "        } finally {",
+            "          $jscomp$swapContext();",
+            "        }",
             "      }",
             "      $jscomp$swapContext(0, 1);",
             "    }",
@@ -220,9 +228,12 @@ public void testForAwaitMixedWithAwait() {
             "                   $jscomp$swapContext(",
             "                       $jscomp$swapContext(await $jscomp$swapContext(gen()), 1))) {",
             "      $jscomp$swapContext(0, 1);",
-            "      $jscomp$swapContext(",
-            "          await $jscomp$swapContext(use(x)), 1);",
-            "      $jscomp$swapContext();",
+            "      try {",
+            "        $jscomp$swapContext(",
+            "            await $jscomp$swapContext(use(x)), 1);",
+            "      } finally {",
+            "        $jscomp$swapContext();",
+            "      }",
             "    }",
             "    $jscomp$swapContext(0, 1);",
             "  } finally {",