Await async iterator return() when a for-await loop breaks

quickjs.c

@@ -22076,6 +22076,7 @@ typedef struct BlockEnv {
     int label_finally; /* -1 if none */
     int scope_level;
     uint8_t has_iterator : 1;
+    uint8_t is_async_iterator : 1;
     uint8_t is_regular_stmt : 1; // i.e. not a loop statement
     uint8_t has_using : 1; /* scope has using declarations needing disposal */
     int using_scope_level; /* scope level for OP_dispose_scope (-1 if none) */
@@ -28422,6 +28423,7 @@ static void push_break_entry(JSFunctionDef *fd, BlockEnv *be,
     be->label_finally = -1;
     be->scope_level = fd->scope_level;
     be->has_iterator = false;
+    be->is_async_iterator = false;
     be->is_regular_stmt = false;
     be->has_using = false;
     be->using_scope_level = -1;
@@ -28434,6 +28436,45 @@ static void pop_break_entry(JSFunctionDef *fd)
     fd->top_break = be->prev;
 }
 
+static void emit_iterator_close(JSParseState *s, bool is_async,
+                                bool preserve_completion)
+{
+    int label_no_return, label_done;
+
+    if (!is_async) {
+        emit_op(s, OP_iterator_close);
+        return;
+    }
+
+    if (preserve_completion) {
+        /* iter_obj next completion -> completion iter_obj */
+        emit_op(s, OP_nip);
+        emit_op(s, OP_swap);
+    } else {
+        /* iter_obj next catch_offset -> iter_obj */
+        emit_op(s, OP_drop);
+        emit_op(s, OP_drop);
+    }
+    emit_op(s, OP_dup);
+    emit_op(s, OP_is_undefined);
+    label_done = emit_goto(s, OP_if_true, -1);
+    emit_op(s, OP_get_field2);
+    emit_atom(s, JS_ATOM_return);
+    /* completion? iter_obj return_func */
+    emit_op(s, OP_dup);
+    emit_op(s, OP_is_undefined_or_null);
+    label_no_return = emit_goto(s, OP_if_true, -1);
+    emit_op(s, OP_call_method);
+    emit_u16(s, 0);
+    emit_op(s, OP_await);
+    emit_op(s, OP_check_object);
+    emit_goto(s, OP_goto, label_done);
+    emit_label(s, label_no_return);
+    emit_op(s, OP_drop);
+    emit_label(s, label_done);
+    emit_op(s, OP_drop);
+}
+
 static __exception int emit_break(JSParseState *s, JSAtom name, int is_cont)
 {
     BlockEnv *top;
@@ -28460,7 +28501,7 @@ static __exception int emit_break(JSParseState *s, JSAtom name, int is_cont)
         }
         i = 0;
         if (top->has_iterator) {
-            emit_op(s, OP_iterator_close);
+            emit_iterator_close(s, top->is_async_iterator, false);
             i += 3;
         }
         for(; i < top->drop_count; i++)
@@ -28521,25 +28562,8 @@ static void emit_return(JSParseState *s, bool hasval)
             emit_op(s, OP_nip_catch);
             /* stack: iter_obj next ret_val */
             if (top->has_iterator) {
-                if (s->cur_func->func_kind == JS_FUNC_ASYNC_GENERATOR) {
-                    int label_next, label_next2;
-                    emit_op(s, OP_nip); /* next */
-                    emit_op(s, OP_swap);
-                    emit_op(s, OP_get_field2);
-                    emit_atom(s, JS_ATOM_return);
-                    /* stack: iter_obj return_func */
-                    emit_op(s, OP_dup);
-                    emit_op(s, OP_is_undefined_or_null);
-                    label_next = emit_goto(s, OP_if_true, -1);
-                    emit_op(s, OP_call_method);
-                    emit_u16(s, 0);
-                    emit_op(s, OP_check_object);
-                    emit_op(s, OP_await);
-                    label_next2 = emit_goto(s, OP_goto, -1);
-                    emit_label(s, label_next);
-                    emit_op(s, OP_drop);
-                    emit_label(s, label_next2);
-                    emit_op(s, OP_drop);
+                if (top->is_async_iterator) {
+                    emit_iterator_close(s, true, true);
                 } else {
                     emit_op(s, OP_rot3r);
                     emit_op(s, OP_undefined); /* dummy catch offset */
@@ -29088,6 +29112,7 @@ static __exception int js_parse_for_in_of(JSParseState *s, int label_name,
            that a yield in the expression does not try to close a
            not-yet-created iterator */
         break_entry.has_iterator = true;
+        break_entry.is_async_iterator = is_async;
         break_entry.drop_count += 2;
         if (is_async)
             emit_op(s, OP_for_await_of_start);
@@ -29196,7 +29221,7 @@ static __exception int js_parse_for_in_of(JSParseState *s, int label_name,
     if (is_for_of) {
         /* close and drop enum_rec */
         emit_source_loc_at(s, source_line_num, source_col_num);
-        emit_op(s, OP_iterator_close);
+        emit_iterator_close(s, is_async, false);
     } else {
         emit_op(s, OP_drop);
     }

tests/for-await-normal-close.js

@@ -33,3 +33,18 @@ function makeAsyncIter(log) {
     } catch (e) {}
     assert(log.returned, 1);
 }
+
+/* leaving the loop via break must await the result of return() */
+{
+    const order = [];
+    const it = { [Symbol.asyncIterator]() { let i = 0; return {
+        next() { return Promise.resolve({ value: i++, done: false }); },
+        return() {
+            order.push("return");
+            return Promise.resolve().then(() => { order.push("resolved"); return { done: true }; });
+        },
+    }; } };
+    for await (const x of it) break;
+    order.push("after");
+    assert(order.join(","), "return,resolved,after");
+}