[TypeCheckEffects] Do not diagnose missing try/await in async let

initializers that involve `if` or `switch` expressions.

Author: hborla

lib/Sema/TypeCheckEffects.cpp

@@ -2562,6 +2562,9 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
 
       /// Do we have any 'await's in this context?
       HasAnyAwait = 0x80,
+
+      /// Are we in an 'async let' initializer context?
+      InAsyncLet = 0x100,
     };
   private:
     unsigned Bits;
@@ -2707,8 +2710,7 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     }
 
     void enterAsyncLet() {
-      Self.Flags.set(ContextFlags::IsTryCovered);
-      Self.Flags.set(ContextFlags::IsAsyncCovered);
+      Self.Flags.set(ContextFlags::InAsyncLet);
     }
 
     void refineLocalContext(Context newContext) {
@@ -2747,6 +2749,11 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
       // 'async'.
     }
 
+    void setCoverageForSingleValueStmtExpr() {
+      resetCoverage();
+      Self.Flags.mergeFrom(ContextFlags::InAsyncLet, OldFlags);
+    }
+
     void preserveCoverageFromSingleValueStmtExpr() {
       // We need to preserve whether we saw any throwing sites, to avoid warning
       // on 'do { let x = if .random() { try ... } else { ... } } catch { ... }'
@@ -2918,7 +2925,7 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     // For an if/switch expression, we reset coverage such that a 'try'/'await'
     // does not cover the branches.
     ContextScope scope(*this, /*newContext*/ llvm::None);
-    scope.resetCoverage();
+    scope.setCoverageForSingleValueStmtExpr();
     SVE->getStmt()->walk(*this);
     scope.preserveCoverageFromSingleValueStmtExpr();
     return ShouldNotRecurse;
@@ -3173,7 +3180,8 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
                                               classification.getAsyncReason());
       }
       // Diagnose async calls that are outside of an await context.
-      else if (!Flags.has(ContextFlags::IsAsyncCovered)) {
+      else if (!(Flags.has(ContextFlags::IsAsyncCovered) ||
+                 Flags.has(ContextFlags::InAsyncLet))) {
         Expr *expr = E.dyn_cast<Expr*>();
         Expr *anchor = walkToAnchor(expr, parentMap,
                                     CurContext.isWithinInterpolatedString());
@@ -3209,7 +3217,8 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
             break;
 
       bool isTryCovered =
-        (!requiresTry || Flags.has(ContextFlags::IsTryCovered));
+        (!requiresTry || Flags.has(ContextFlags::IsTryCovered) ||
+         Flags.has(ContextFlags::InAsyncLet));
       if (!CurContext.handlesThrows(throwsKind)) {
         CurContext.diagnoseUnhandledThrowSite(Ctx.Diags, E, isTryCovered,
                                               classification.getThrowReason());

test/expr/unary/if_expr.swift

@@ -1492,6 +1492,16 @@ func tryAwaitIf15(_ fn: () async throws -> Int) async rethrows -> Int {
   }
 }
 
+func asyncLetIf(cond: Bool, _ fn: () async throws -> Int) async throws -> Int {
+  async let x = if cond {
+    fn()
+  } else {
+    0
+  }
+
+  return try await x
+}
+
 struct AnyEraserP: EraserP {
   init<T: EraserP>(erasing: T) {}
 }