From 42f9143e114c5c07f40df83ed07ffeb3cbaf2101 Mon Sep 17 00:00:00 2001
From: Jack Works <jackworks@protonmail.com>
Date: Fri, 30 Sep 2022 07:13:25 +0800
Subject: [PATCH] feat: codefix for `for await of` (#50623)

---
 src/compiler/checker.ts                        | 18 +++++++++++++++++-
 src/compiler/types.ts                          |  1 +
 src/services/codefixes/addMissingAwait.ts      |  9 +++++++++
 .../codeFixAddMissingAwait_forAwaitOf.ts       | 17 +++++++++++++++++
 4 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 tests/cases/fourslash/codeFixAddMissingAwait_forAwaitOf.ts

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index 1677cc6a7cb4c..08d81c69e00e6 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -642,6 +642,11 @@ namespace ts {
             getOptionalType: () => optionalType,
             getPromiseType: () => getGlobalPromiseType(/*reportErrors*/ false),
             getPromiseLikeType: () => getGlobalPromiseLikeType(/*reportErrors*/ false),
+            getAsyncIterableType: () => {
+                const type = getGlobalAsyncIterableType(/*reportErrors*/ false);
+                if (type === emptyGenericType) return undefined;
+                return type;
+            },
             isSymbolAccessible,
             isArrayType,
             isTupleType,
@@ -39429,7 +39434,18 @@ namespace ts {
             const message = allowAsyncIterables
                 ? Diagnostics.Type_0_must_have_a_Symbol_asyncIterator_method_that_returns_an_async_iterator
                 : Diagnostics.Type_0_must_have_a_Symbol_iterator_method_that_returns_an_iterator;
-            return errorAndMaybeSuggestAwait(errorNode, !!getAwaitedTypeOfPromise(type), message, typeToString(type));
+            const suggestAwait =
+                // for (const x of Promise<...>) or [...Promise<...>]
+                !!getAwaitedTypeOfPromise(type)
+                // for (const x of AsyncIterable<...>)
+                || (
+                    !allowAsyncIterables &&
+                    isForOfStatement(errorNode.parent) &&
+                    errorNode.parent.expression === errorNode &&
+                    getGlobalAsyncIterableType(/** reportErrors */ false) !== emptyGenericType &&
+                    isTypeAssignableTo(type, getGlobalAsyncIterableType(/** reportErrors */ false)
+                ));
+            return errorAndMaybeSuggestAwait(errorNode, suggestAwait, message, typeToString(type));
         }
 
         /**
diff --git a/src/compiler/types.ts b/src/compiler/types.ts
index 66e55ab58b381..997c83b2f9db7 100644
--- a/src/compiler/types.ts
+++ b/src/compiler/types.ts
@@ -4760,6 +4760,7 @@ namespace ts {
         /* @internal */ createPromiseType(type: Type): Type;
         /* @internal */ getPromiseType(): Type;
         /* @internal */ getPromiseLikeType(): Type;
+        /* @internal */ getAsyncIterableType(): Type | undefined;
 
         /* @internal */ isTypeAssignableTo(source: Type, target: Type): boolean;
         /* @internal */ createAnonymousType(symbol: Symbol | undefined, members: SymbolTable, callSignatures: Signature[], constructSignatures: Signature[], indexInfos: IndexInfo[]): Type;
diff --git a/src/services/codefixes/addMissingAwait.ts b/src/services/codefixes/addMissingAwait.ts
index a03d355b0986b..9020b9791dfb3 100644
--- a/src/services/codefixes/addMissingAwait.ts
+++ b/src/services/codefixes/addMissingAwait.ts
@@ -226,6 +226,15 @@ namespace ts.codefix {
     }
 
     function makeChange(changeTracker: textChanges.ChangeTracker, errorCode: number, sourceFile: SourceFile, checker: TypeChecker, insertionSite: Expression, fixedDeclarations?: Set<number>) {
+        if (isForOfStatement(insertionSite.parent) && !insertionSite.parent.awaitModifier) {
+            const exprType = checker.getTypeAtLocation(insertionSite);
+            const asyncIter = checker.getAsyncIterableType();
+            if (asyncIter && checker.isTypeAssignableTo(exprType, asyncIter)) {
+                const forOf = insertionSite.parent;
+                changeTracker.replaceNode(sourceFile, forOf, factory.updateForOfStatement(forOf, factory.createToken(SyntaxKind.AwaitKeyword), forOf.initializer, forOf.expression, forOf.statement));
+                return;
+            }
+        }
         if (isBinaryExpression(insertionSite)) {
             for (const side of [insertionSite.left, insertionSite.right]) {
                 if (fixedDeclarations && isIdentifier(side)) {
diff --git a/tests/cases/fourslash/codeFixAddMissingAwait_forAwaitOf.ts b/tests/cases/fourslash/codeFixAddMissingAwait_forAwaitOf.ts
new file mode 100644
index 0000000000000..00b7e7ded8fbe
--- /dev/null
+++ b/tests/cases/fourslash/codeFixAddMissingAwait_forAwaitOf.ts
@@ -0,0 +1,17 @@
+/// <reference path="fourslash.ts" />
+// @lib: es2020
+// @target: es2020
+////async function* g() {}
+////async function fn() {
+////    for (const { } of g()) { }
+////}
+
+verify.codeFix({
+  description: ts.Diagnostics.Add_await.message,
+  index: 0,
+  newFileContent:
+`async function* g() {}
+async function fn() {
+    for await (const { } of g()) { }
+}`
+});