[Sema] Fix try/await handling for if/switch expressions

Previously we allowed an attached `try` or `await`
on an `if`/`switch` expression to cover the branches,
which does not match what was proposed, and is
especially harmful for multi-statement cases. Fix
the effect handling logic such that we reset
effect coverage for `if`/`switch` expressions
similar to closures, but continue to maintain the
state needed for rethrows checking.

rdar://116066748

Author: hamishknight

lib/Sema/TypeCheckEffects.cpp

@@ -448,6 +448,8 @@ class EffectsHandlingWalker : public ASTWalker {
       recurse = asImpl().checkInterpolatedStringLiteral(interpolated);
     } else if (auto macroExpansionExpr = dyn_cast<MacroExpansionExpr>(E)) {
       recurse = ShouldRecurse;
+    } else if (auto *SVE = dyn_cast<SingleValueStmtExpr>(E)) {
+      recurse = asImpl().checkSingleValueStmtExpr(SVE);
     }
     // Error handling validation (via checkTopLevelEffects) happens after
     // type checking. If an unchecked expression is still around, the code was
@@ -1076,6 +1078,10 @@ class ApplyClassifier {
       return ShouldRecurse;
     }
 
+    ShouldRecurse_t checkSingleValueStmtExpr(SingleValueStmtExpr *SVE) {
+      return ShouldRecurse;
+    }
+
     ConditionalEffectKind checkExhaustiveDoBody(DoCatchStmt *S) {
       // All errors thrown by the do body are caught, but any errors thrown
       // by the catch bodies are bounded by the throwing kind of the do body.
@@ -1196,6 +1202,10 @@ class ApplyClassifier {
       return ShouldRecurse;
     }
 
+    ShouldRecurse_t checkSingleValueStmtExpr(SingleValueStmtExpr *SVE) {
+      return ShouldRecurse;
+    }
+
     void visitExprPre(Expr *expr) { return; }
   };
 
@@ -2211,6 +2221,15 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
       // 'async'.
     }
 
+    void preserveCoverageFromSingleValueStmtExpr() {
+      // We need to preserve whether we saw any throwing sites, to avoid warning
+      // on 'do { let x = if .random() { try ... } else { ... } } catch { ... }'
+      OldFlags.mergeFrom(ContextFlags::HasAnyThrowSite, Self.Flags);
+
+      // We need to preserve the throwing kind to correctly handle rethrows.
+      OldMaxThrowingKind = std::max(OldMaxThrowingKind, Self.MaxThrowingKind);
+    }
+
     void preserveCoverageFromNonExhaustiveCatch() {
       OldFlags.mergeFrom(ContextFlags::HasAnyThrowSite, Self.Flags);
       OldMaxThrowingKind = std::max(OldMaxThrowingKind, Self.MaxThrowingKind);
@@ -2354,6 +2373,17 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     return ShouldNotRecurse;
   }
 
+  ShouldRecurse_t
+  checkSingleValueStmtExpr(SingleValueStmtExpr *SVE) {
+    // 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();
+    SVE->getStmt()->walk(*this);
+    scope.preserveCoverageFromSingleValueStmtExpr();
+    return ShouldNotRecurse;
+  }
+
   ConditionalEffectKind checkExhaustiveDoBody(DoCatchStmt *S) {
     // This is a context where errors are handled.
     ContextScope scope(*this, CurContext.withHandlesErrors());