Convert async flag to source locations and add initial try support to for await in syntax

Author: phausler

include/swift/AST/Stmt.h

@@ -726,7 +726,8 @@ class RepeatWhileStmt : public LabeledStmt {
 /// \endcode
 class ForEachStmt : public LabeledStmt {
   SourceLoc ForLoc;
-  bool IsAsync;
+  SourceLoc TryLoc;
+  SourceLoc AwaitLoc;
   Pattern *Pat;
   SourceLoc InLoc;
   Expr *Sequence;
@@ -742,12 +743,12 @@ class ForEachStmt : public LabeledStmt {
   Expr *convertElementExpr = nullptr;
 
 public:
-  ForEachStmt(LabeledStmtInfo LabelInfo, SourceLoc ForLoc, bool IsAsync, Pattern *Pat,
-              SourceLoc InLoc, Expr *Sequence, SourceLoc WhereLoc,
+  ForEachStmt(LabeledStmtInfo LabelInfo, SourceLoc ForLoc, SourceLoc TryLoc, SourceLoc AwaitLoc, 
+              Pattern *Pat, SourceLoc InLoc, Expr *Sequence, SourceLoc WhereLoc,
               Expr *WhereExpr, BraceStmt *Body, Optional<bool> implicit = None)
     : LabeledStmt(StmtKind::ForEach, getDefaultImplicitFlag(implicit, ForLoc),
                   LabelInfo),
-      ForLoc(ForLoc), IsAsync(IsAsync), Pat(nullptr), InLoc(InLoc), Sequence(Sequence),
+      ForLoc(ForLoc), TryLoc(TryLoc), AwaitLoc(AwaitLoc), Pat(nullptr), InLoc(InLoc), Sequence(Sequence),
       WhereLoc(WhereLoc), WhereExpr(WhereExpr), Body(Body) {
     setPattern(Pat);
   }
@@ -780,7 +781,8 @@ class ForEachStmt : public LabeledStmt {
   /// getWhereLoc - Retrieve the location of the 'where' keyword.
   SourceLoc getWhereLoc() const { return WhereLoc; }
 
-  bool isAsync() const { return IsAsync; }
+  SourceLoc getAwaitLoc() const { return AwaitLoc; }
+  SourceLoc getTryLoc() const { return TryLoc; }
 
   /// getPattern - Retrieve the pattern describing the iteration variables.
   /// These variables will only be visible within the body of the loop.

lib/Parse/ParseStmt.cpp

@@ -2124,12 +2124,18 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
   // lookahead to resolve what is going on.
   bool IsCStyleFor = isStmtForCStyle(*this);
   auto StartOfControl = Tok.getLoc();
-  bool IsAsync = false;
+  SourceLoc AwaitLoc;
+  SourceLoc TryLoc;
 
   if (shouldParseExperimentalConcurrency() && 
     Tok.isContextualKeyword("await")) {
-    consumeToken();
-    IsAsync = true;
+    AwaitLoc = consumeToken();
+  } if (shouldParseExperimentalConcurrency() && 
+    Tok.is(tok::kw_try)) {
+    TryLoc = consumeToken();
+    if (Tok.isContextualKeyword("await")) {
+      AwaitLoc = consumeToken();
+    }
   }
 
   // Parse the pattern.  This is either 'case <refutable pattern>' or just a
@@ -2225,7 +2231,7 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
 
   return makeParserResult(
       Status,
-      new (Context) ForEachStmt(LabelInfo, ForLoc, IsAsync, pattern.get(), InLoc,
+      new (Context) ForEachStmt(LabelInfo, ForLoc, TryLoc, AwaitLoc, pattern.get(), InLoc,
                                 Container.get(), WhereLoc, Where.getPtrOrNull(),
                                 Body.get()));
 }

lib/SILGen/SILGenStmt.cpp

@@ -1198,7 +1198,7 @@ void StmtEmitter::visitAsyncForEachStmt(ForEachStmt *S) {
 }
 
 void StmtEmitter::visitForEachStmt(ForEachStmt *S) {
-  if (S->isAsync()) {
+  if (S->getAwaitLoc().isValid()) {
     visitAsyncForEachStmt(S);
     return;
   }

lib/Sema/BuilderTransform.cpp

@@ -774,7 +774,7 @@ class BuilderClosureVisitor
     // take care of this.
     auto sequenceProto = TypeChecker::getProtocol(
         dc->getASTContext(), forEachStmt->getForLoc(),
-        forEachStmt->isAsync() ? 
+        forEachStmt->getAwaitLoc().isValid() ? 
           KnownProtocolKind::AsyncSequence : KnownProtocolKind::Sequence);
     if (!sequenceProto) {
       if (!unhandledNode)

lib/Sema/CSGen.cpp

@@ -3655,7 +3655,7 @@ generateForEachStmtConstraints(
     ConstraintSystem &cs, SolutionApplicationTarget target, Expr *sequence) {
   auto forEachStmtInfo = target.getForEachStmtInfo();
   ForEachStmt *stmt = forEachStmtInfo.stmt;
-  bool isAsync = stmt->isAsync();
+  bool isAsync = stmt->getAwaitLoc().isValid();
 
   auto locator = cs.getConstraintLocator(sequence);
   auto contextualLocator =

lib/Sema/TypeCheckConstraints.cpp

@@ -563,7 +563,7 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
 
   auto sequenceProto = TypeChecker::getProtocol(
       dc->getASTContext(), stmt->getForLoc(), 
-      stmt->isAsync() ? 
+      stmt->getAwaitLoc().isValid() ? 
         KnownProtocolKind::AsyncSequence : KnownProtocolKind::Sequence);
   if (!sequenceProto)
     return failed();