[Sema] Correct 'await' insertion fixIt

lib/Sema/TypeCheckEffects.cpp

@@ -2801,7 +2801,7 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
 
   static bool isEffectAnchor(Expr *e) {
     return isa<AbstractClosureExpr>(e) || isa<DiscardAssignmentExpr>(e) ||
-           isa<AssignExpr>(e);
+           isa<AssignExpr>(e) || (isa<DeclRefExpr>(e) && e->isImplicit());
   }
 
   static bool isAnchorTooEarly(Expr *e) {
@@ -3581,27 +3581,43 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     Ctx.Diags.diagnose(E->getAwaitLoc(), diag::no_async_in_await);
   }
 
-  void diagnoseUncoveredAsyncSite(const Expr *anchor) const {
-    auto asyncPointIter = uncoveredAsync.find(anchor);
-    if (asyncPointIter == uncoveredAsync.end())
-      return;
-    const std::vector<DiagnosticInfo> &errors = asyncPointIter->getSecond();
+  std::pair<SourceLoc, std::string>
+  getFixItForUncoveredAsyncSite(const Expr *anchor) const {
     SourceLoc awaitInsertLoc = anchor->getStartLoc();
-    if (const AnyTryExpr *tryExpr = dyn_cast<AnyTryExpr>(anchor))
+    std::string insertText = "await ";
+    if (auto *tryExpr = dyn_cast<AnyTryExpr>(anchor))
       awaitInsertLoc = tryExpr->getSubExpr()->getStartLoc();
-    else if (const AutoClosureExpr *autoClosure = dyn_cast<AutoClosureExpr>(anchor)) {
-      if (const AnyTryExpr *tryExpr = dyn_cast<AnyTryExpr>(autoClosure->getSingleExpressionBody()))
+    else if (auto *autoClosure = dyn_cast<AutoClosureExpr>(anchor)) {
+      if (auto *tryExpr =
+              dyn_cast<AnyTryExpr>(autoClosure->getSingleExpressionBody()))
         awaitInsertLoc = tryExpr->getSubExpr()->getStartLoc();
+      // Supply a tailored fixIt including the identifier if we are
+      // looking at a shorthand optional binding.
+    } else if (anchor->isImplicit()) {
+      if (auto declRef = dyn_cast<DeclRefExpr>(anchor))
+        if (auto var = dyn_cast_or_null<VarDecl>(declRef->getDecl())) {
+          insertText = " = await " + var->getNameStr().str();
+          awaitInsertLoc = Lexer::getLocForEndOfToken(Ctx.Diags.SourceMgr,
+                                                       anchor->getStartLoc());
+        }
     }
+    return std::make_pair(awaitInsertLoc, insertText);
+  }
 
+  void diagnoseUncoveredAsyncSite(const Expr *anchor) const {
+    auto asyncPointIter = uncoveredAsync.find(anchor);
+    if (asyncPointIter == uncoveredAsync.end())
+      return;
+    const auto &errors = asyncPointIter->getSecond();
+    const auto &[loc, insertText] = getFixItForUncoveredAsyncSite(anchor);
     bool downgradeToWarning = llvm::all_of(errors,
         [&](DiagnosticInfo diag) -> bool {
           return diag.downgradeToWarning;
         });
 
     Ctx.Diags.diagnose(anchor->getStartLoc(), diag::async_expr_without_await)
       .warnUntilSwiftVersionIf(downgradeToWarning, 6)
-      .fixItInsert(awaitInsertLoc, "await ")
+      .fixItInsert(loc, insertText)
       .highlight(anchor->getSourceRange());
 
     for (const DiagnosticInfo &diag: errors) {

test/expr/unary/async_await.swift

@@ -217,3 +217,25 @@ func testAsyncLetOutOfAsync() {
   _ = await x  // expected-error{{'async let' in a function that does not support concurrency}}
   _ = x // expected-error{{'async let' in a function that does not support concurrency}}
 }
+
+class A {}
+class B: A {}
+func f(_ x: String) async -> String? { x }
+func testAsyncExprWithoutAwait() async {
+  async let result: B? = nil
+  if let result: A = result {} // expected-error {{expression is 'async' but is not marked with 'await'}} {{22-22=await }}
+  // expected-warning@-1 {{immutable value 'result' was never used; consider replacing with '_' or removing it}}
+  // expected-note@-2 {{reference to async let 'result' is 'async'}}
+  if let result: A {} // expected-error {{expression is 'async' but is not marked with 'await'}} {{19-19= = await result}}
+  // expected-warning@-1 {{immutable value 'result' was never used; consider replacing with '_' or removing it}}
+  // expected-note@-2 {{reference to async let 'result' is 'async'}}
+  if let result = result {} // expected-error {{expression is 'async' but is not marked with 'await'}} {{19-19=await }}
+  // expected-warning@-1 {{value 'result' was defined but never used; consider replacing with boolean test}}
+  // expected-note@-2 {{reference to async let 'result' is 'async'}}
+  if let result {} // expected-error {{expression is 'async' but is not marked with 'await'}} {{16-16= = await result}}
+  // expected-warning@-1 {{value 'result' was defined but never used; consider replacing with boolean test}}
+  // expected-note@-2 {{reference to async let 'result' is 'async'}}
+  let a = f("a") // expected-error {{expression is 'async' but is not marked with 'await'}} {{11-11=await }}
+  // expected-warning@-1 {{initialization of immutable value 'a' was never used; consider replacing with assignment to '_' or removing it}}
+  // expected-note@-2 {{call is 'async'}}
+}