for await runs its body one iteration at a time (#498)

tslang/lib/TypeScript/MLIRGenStatements.cpp

@@ -815,17 +815,6 @@ namespace mlirgen
         SmallVector<mlir::Type, 0> types;
         SmallVector<mlir::Value, 0> operands;
 
-        mlir::Value asyncGroupResult;
-        if (hasAwait)
-        {
-            auto groupType = mlir::async::GroupType::get(builder.getContext());
-            auto blockSize = builder.create<mlir_ts::ConstantOp>(location, builder.getIndexAttr(0));
-            auto asyncGroupOp = builder.create<mlir::async::CreateGroupOp>(location, groupType, blockSize);
-            asyncGroupResult = asyncGroupOp.getResult();
-            // operands.push_back(asyncGroupOp);
-            // types.push_back(groupType);
-        }
-
         auto forOp = builder.create<mlir_ts::ForOp>(location, types, operands);
         if (!label.empty())
         {
@@ -908,10 +897,10 @@ namespace mlirgen
                 return mlir::failure();
             }
 
-            // add to group
-            auto rankType = mlir::IndexType::get(builder.getContext());
-            // TODO: should i replace with value from arg0?
-            builder.create<mlir::async::AddToGroupOp>(location, rankType, asyncExecOp.getToken(), asyncGroupResult);
+            // One iteration at a time, as `for await` runs its body: the next value is asked for only
+            // once this body is done. Each body was added to a group awaited after the loop, so the
+            // bodies ran as concurrent tasks - interleaved, losing each other's writes (#498).
+            builder.create<mlir::async::AwaitOp>(location, asyncExecOp.getToken());
         }
         else
         {
@@ -929,25 +918,6 @@ namespace mlirgen
 
         builder.setInsertionPointAfter(forOp);
 
-        if (hasAwait)
-        {
-            // Not helping
-            /*
-            // async await all, see convert-to-llvm.mlir
-            auto asyncExecAwaitAllOp =
-                builder.create<mlir::async::ExecuteOp>(location, mlir::TypeRange{}, mlir::ValueRange{},
-            mlir::ValueRange{},
-                                                       [&](mlir::OpBuilder &builder, mlir::Location location,
-            mlir::ValueRange values) { builder.create<mlir::async::AwaitAllOp>(location, asyncGroupResult);
-                                                           builder.create<mlir::async::YieldOp>(location,
-            mlir::ValueRange{});
-                                                       });
-            */
-
-            // Wait for the completion of all subtasks.
-            builder.create<mlir::async::AwaitAllOp>(location, asyncGroupResult);
-        }
-
         return mlir::success();
     }
 

tslang/test/tester/CMakeLists.txt

@@ -352,6 +352,7 @@ tslang_add_test(NAME test-compile-00-tuple-runtime-arrays COMMAND test-runner "$
 tslang_add_test(NAME test-compile-00-record-to-class COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00record_to_class.ts")
 tslang_add_test(NAME test-compile-00-array-insert-functions COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_insert_functions.ts")
 tslang_add_test(NAME test-compile-00-await-void COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00await_void.ts")
+tslang_add_test(NAME test-compile-00-for-await-sequential COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await_sequential.ts")
 tslang_add_test(NAME test-compile-00-const-record-owned-fields COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts")
 tslang_add_test(NAME test-compile-00-array-length-valid COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_length_valid.ts")
 tslang_add_test(NAME test-compile-00-string-view-copy COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_view_copy.ts")
@@ -861,6 +862,7 @@ tslang_add_test(NAME test-jit-00-tuple-runtime-arrays COMMAND test-runner -jit "
 tslang_add_test(NAME test-jit-00-record-to-class COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00record_to_class.ts")
 tslang_add_test(NAME test-jit-00-array-insert-functions COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_insert_functions.ts")
 tslang_add_test(NAME test-jit-00-await-void COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00await_void.ts")
+tslang_add_test(NAME test-jit-00-for-await-sequential COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00for_await_sequential.ts")
 tslang_add_test(NAME test-jit-00-const-record-owned-fields COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00const_record_owned_fields.ts")
 tslang_add_test(NAME test-jit-00-array-length-valid COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00array_length_valid.ts")
 tslang_add_test(NAME test-jit-00-string-view-copy COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00string_view_copy.ts")
@@ -1699,6 +1701,7 @@ set(TSLANG_CORPUS
     00record_to_class.ts
     00array_insert_functions.ts
     00await_void.ts
+    00for_await_sequential.ts
     00const_record_owned_fields.ts
     00array_length_valid.ts
     00string_view_copy.ts

tslang/test/tester/tests/00for_await_sequential.ts

@@ -0,0 +1,30 @@
+// The body of a `for await` runs one iteration at a time, in order (#498): each body was a task
+// of its own, all awaited together after the loop, so iterations interleaved and lost each
+// other's writes.
+let log = "";
+let keep: number[] = [];
+async function* numbers() {
+    for (let i = 1; i <= 8; i++) {
+        yield i;
+    }
+}
+async function note(s: string) {
+    log += s;
+}
+async function main() {
+    for await (const x of numbers()) {
+        log += "<" + x;
+        keep.push(x);
+        log += ">";
+    }
+    assert(log == "<1><2><3><4><5><6><7><8>", "each body whole, in order");
+    assert(keep.length == 8, "every push kept");
+
+    log = "";
+    for await (const x of numbers()) {
+        await note("[" + x);
+        await note("]");
+    }
+    assert(log == "[1][2][3][4][5][6][7][8]", "awaits in the body, in order");
+    print("done.");
+}