[Diagnostics] Apply "unhandled throw" diagnostic for for-in loop in closures

Extract diagnostic into a method and use it while type-checking
`for-in` in top-level code and in closures.

Author: xedin

lib/Sema/CSClosure.cpp

@@ -16,6 +16,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "MiscDiagnostics.h"
 #include "TypeChecker.h"
 #include "swift/Sema/ConstraintSystem.h"
 
@@ -1163,12 +1164,17 @@ class ClosureConstraintApplication
 
     auto forEachTarget =
         rewriteTarget(*cs.getSolutionApplicationTarget(forEachStmt));
+
     if (!forEachTarget)
       hadError = true;
 
     auto body = visit(forEachStmt->getBody()).get<Stmt *>();
     forEachStmt->setBody(cast<BraceStmt>(body));
 
+    // Check to see if the sequence expr is throwing (in async context),
+    // if so require the stmt to have a `try`.
+    hadError |= diagnoseUnhandledThrowsInAsyncContext(closure, forEachStmt);
+
     return forEachStmt;
   }
 

lib/Sema/MiscDiagnostics.cpp

@@ -5140,3 +5140,33 @@ Optional<Identifier> TypeChecker::omitNeedlessWords(VarDecl *var) {
 
   return None;
 }
+
+bool swift::diagnoseUnhandledThrowsInAsyncContext(DeclContext *dc,
+                                                  ForEachStmt *forEach) {
+  if (!forEach->getAwaitLoc().isValid())
+    return false;
+
+  auto &ctx = dc->getASTContext();
+
+  auto sequenceProto = TypeChecker::getProtocol(
+      ctx, forEach->getForLoc(), KnownProtocolKind::AsyncSequence);
+
+  if (!sequenceProto)
+    return false;
+
+  // fetch the sequence out of the statement
+  // else wise the value is potentially unresolved
+  auto Ty = forEach->getSequence()->getType();
+  auto module = dc->getParentModule();
+  auto conformanceRef = module->lookupConformance(Ty, sequenceProto);
+
+  if (conformanceRef.hasEffect(EffectKind::Throws) &&
+      forEach->getTryLoc().isInvalid()) {
+    ctx.Diags
+        .diagnose(forEach->getAwaitLoc(), diag::throwing_call_unhandled, "call")
+        .fixItInsert(forEach->getAwaitLoc(), "try");
+    return true;
+  }
+
+  return false;
+}

lib/Sema/MiscDiagnostics.h

@@ -32,6 +32,7 @@ namespace swift {
   class Stmt;
   class TopLevelCodeDecl;
   class ValueDecl;
+  class ForEachStmt;
 
 /// Emit diagnostics for syntactic restrictions on a given expression.
 void performSyntacticExprDiagnostics(
@@ -115,6 +116,12 @@ void fixItEncloseTrailingClosure(ASTContext &ctx,
 /// we emit a diagnostic suggesting the async call.
 void checkFunctionAsyncUsage(AbstractFunctionDecl *decl);
 void checkPatternBindingDeclAsyncUsage(PatternBindingDecl *decl);
+
+/// Detect and diagnose a missing `try` in `for-in` loop sequence
+/// expression in async context (denoted with `await` keyword).
+bool diagnoseUnhandledThrowsInAsyncContext(DeclContext *dc,
+                                           ForEachStmt *forEach);
+
 } // namespace swift
 
 #endif // SWIFT_SEMA_MISC_DIAGNOSTICS_H

lib/Sema/TypeCheckConstraints.cpp

@@ -610,24 +610,10 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
   if (!typeCheckExpression(target))
     return failed();
 
-  // check to see if the sequence expr is throwing (and async), if so require 
-  // the stmt to have a try loc
-  if (stmt->getAwaitLoc().isValid()) {
-    // fetch the sequence out of the statement
-    // else wise the value is potentially unresolved
-    auto Ty = stmt->getSequence()->getType();
-    auto module = dc->getParentModule();
-    auto conformanceRef = module->lookupConformance(Ty, sequenceProto);
-    
-    if (conformanceRef.hasEffect(EffectKind::Throws) &&
-        stmt->getTryLoc().isInvalid()) {
-      auto &diags = dc->getASTContext().Diags;
-      diags.diagnose(stmt->getAwaitLoc(), diag::throwing_call_unhandled, "call")
-        .fixItInsert(stmt->getAwaitLoc(), "try");
-
-      return failed();
-    }
-  }
+  // Check to see if the sequence expr is throwing (in async context),
+  // if so require the stmt to have a `try`.
+  if (diagnoseUnhandledThrowsInAsyncContext(dc, stmt))
+    return failed();
 
   return false;
 }