Merge pull request #33408 from DougGregor/async-closures

[Concurrency] Add support for 'async' closures.

Author: DougGregor

include/swift/AST/Expr.h

@@ -3815,6 +3815,9 @@ class ClosureExpr : public AbstractClosureExpr {
   /// this information directly on the ClosureExpr.
   VarDecl * CapturedSelfDecl;
 
+  /// The location of the "async", if present.
+  SourceLoc AsyncLoc;
+
   /// The location of the "throws", if present.
   SourceLoc ThrowsLoc;
 
@@ -3833,14 +3836,15 @@ class ClosureExpr : public AbstractClosureExpr {
   llvm::PointerIntPair<BraceStmt *, 1, bool> Body;
 public:
   ClosureExpr(SourceRange bracketRange, VarDecl *capturedSelfDecl,
-              ParameterList *params, SourceLoc throwsLoc, SourceLoc arrowLoc,
-              SourceLoc inLoc, TypeExpr *explicitResultType,
+              ParameterList *params, SourceLoc asyncLoc, SourceLoc throwsLoc,
+              SourceLoc arrowLoc, SourceLoc inLoc, TypeExpr *explicitResultType,
               unsigned discriminator, DeclContext *parent)
     : AbstractClosureExpr(ExprKind::Closure, Type(), /*Implicit=*/false,
                           discriminator, parent),
       BracketRange(bracketRange),
       CapturedSelfDecl(capturedSelfDecl),
-      ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
+      AsyncLoc(asyncLoc), ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc),
+      InLoc(inLoc),
       ExplicitResultTypeAndBodyState(explicitResultType, BodyState::Parsed),
       Body(nullptr) {
     setParameterList(params);
@@ -3888,7 +3892,12 @@ class ClosureExpr : public AbstractClosureExpr {
   SourceLoc getInLoc() const {
     return InLoc;
   }
-  
+
+  /// Retrieve the location of the 'async' for a closure that has it.
+  SourceLoc getAsyncLoc() const {
+    return AsyncLoc;
+  }
+
   /// Retrieve the location of the 'throws' for a closure that has it.
   SourceLoc getThrowsLoc() const {
     return ThrowsLoc;

include/swift/AST/ExtInfo.h

@@ -436,6 +436,14 @@ class ASTExtInfo {
     return builder.withThrows(throws).build();
   }
 
+  /// Helper method for changing only the async field.
+  ///
+  /// Prefer using \c ASTExtInfoBuilder::withAsync for chaining.
+  LLVM_NODISCARD
+  ASTExtInfo withAsync(bool async = true) const {
+    return builder.withAsync(async).build();
+  }
+
   bool isEqualTo(ASTExtInfo other, bool useClangTypes) const {
     return builder.isEqualTo(other.builder, useClangTypes);
   }

include/swift/Parse/Parser.h

@@ -1572,6 +1572,7 @@ class Parser {
           SmallVectorImpl<CaptureListEntry> &captureList,
           VarDecl *&capturedSelfParamDecl,
           ParameterList *&params,
+          SourceLoc &asyncLoc,
           SourceLoc &throwsLoc,
           SourceLoc &arrowLoc,
           TypeExpr *&explicitResultType,

lib/Parse/ParseExpr.cpp

@@ -2432,7 +2432,8 @@ bool Parser::
 parseClosureSignatureIfPresent(SourceRange &bracketRange,
                                SmallVectorImpl<CaptureListEntry> &captureList,
                                VarDecl *&capturedSelfDecl,
-                               ParameterList *&params, SourceLoc &throwsLoc,
+                               ParameterList *&params,
+                               SourceLoc &asyncLoc, SourceLoc &throwsLoc,
                                SourceLoc &arrowLoc,
                                TypeExpr *&explicitResultType, SourceLoc &inLoc){
   // Clear out result parameters.
@@ -2444,6 +2445,24 @@ parseClosureSignatureIfPresent(SourceRange &bracketRange,
   explicitResultType = nullptr;
   inLoc = SourceLoc();
 
+  // Consume 'async', 'throws', and 'rethrows', but in any order.
+  auto consumeAsyncThrows = [&] {
+    bool hadAsync = false;
+    if (Context.LangOpts.EnableExperimentalConcurrency &&
+        Tok.isContextualKeyword("async")) {
+      consumeToken();
+      hadAsync = true;
+    }
+
+    if (!consumeIf(tok::kw_throws) && !consumeIf(tok::kw_rethrows))
+      return;
+
+    if (Context.LangOpts.EnableExperimentalConcurrency && !hadAsync &&
+        Tok.isContextualKeyword("async")) {
+      consumeToken();
+    }
+  };
+
   // If we have a leading token that may be part of the closure signature, do a
   // speculative parse to validate it and look for 'in'.
   if (Tok.isAny(tok::l_paren, tok::l_square, tok::identifier, tok::kw__)) {
@@ -2465,7 +2484,8 @@ parseClosureSignatureIfPresent(SourceRange &bracketRange,
 
       // Consume the ')', if it's there.
       if (consumeIf(tok::r_paren)) {
-        consumeIf(tok::kw_throws) || consumeIf(tok::kw_rethrows);
+        consumeAsyncThrows();
+
         // Parse the func-signature-result, if present.
         if (consumeIf(tok::arrow)) {
           if (!canParseType())
@@ -2485,8 +2505,8 @@ parseClosureSignatureIfPresent(SourceRange &bracketRange,
 
         return false;
       }
-      
-      consumeIf(tok::kw_throws) || consumeIf(tok::kw_rethrows);
+
+      consumeAsyncThrows();
 
       // Parse the func-signature-result, if present.
       if (consumeIf(tok::arrow)) {
@@ -2682,11 +2702,10 @@ parseClosureSignatureIfPresent(SourceRange &bracketRange,
 
       params = ParameterList::create(Context, elements);
     }
-    
-    if (Tok.is(tok::kw_throws)) {
-      throwsLoc = consumeToken();
-    } else if (Tok.is(tok::kw_rethrows)) {
-      throwsLoc = consumeToken();
+
+    bool rethrows = false;
+    parseAsyncThrows(SourceLoc(), asyncLoc, throwsLoc, &rethrows);
+    if (rethrows) {
       diagnose(throwsLoc, diag::rethrowing_function_type)
         .fixItReplace(throwsLoc, "throws");
     }
@@ -2803,13 +2822,14 @@ ParserResult<Expr> Parser::parseExprClosure() {
   SmallVector<CaptureListEntry, 2> captureList;
   VarDecl *capturedSelfDecl;
   ParameterList *params = nullptr;
+  SourceLoc asyncLoc;
   SourceLoc throwsLoc;
   SourceLoc arrowLoc;
   TypeExpr *explicitResultType;
   SourceLoc inLoc;
-  parseClosureSignatureIfPresent(bracketRange, captureList,
-                                 capturedSelfDecl, params, throwsLoc,
-                                 arrowLoc, explicitResultType, inLoc);
+  parseClosureSignatureIfPresent(
+      bracketRange, captureList, capturedSelfDecl, params, asyncLoc, throwsLoc,
+      arrowLoc, explicitResultType, inLoc);
 
   // If the closure was created in the context of an array type signature's
   // size expression, there will not be a local context. A parse error will
@@ -2824,10 +2844,9 @@ ParserResult<Expr> Parser::parseExprClosure() {
   unsigned discriminator = CurLocalContext->claimNextClosureDiscriminator();
 
   // Create the closure expression and enter its context.
-  auto *closure = new (Context) ClosureExpr(bracketRange, capturedSelfDecl,
-                                            params, throwsLoc, arrowLoc, inLoc,
-                                            explicitResultType, discriminator,
-                                            CurDeclContext);
+  auto *closure = new (Context) ClosureExpr(
+      bracketRange, capturedSelfDecl, params, asyncLoc, throwsLoc, arrowLoc,
+      inLoc, explicitResultType, discriminator, CurDeclContext);
   // The arguments to the func are defined in their own scope.
   Scope S(this, ScopeKind::ClosureParams);
   ParseFunctionBody cc(*this, closure);

lib/Sema/CSGen.cpp

@@ -2114,9 +2114,7 @@ namespace {
         }
       }
 
-      auto extInfo = FunctionType::ExtInfo();
-      if (closureCanThrow(closure))
-        extInfo = extInfo.withThrows();
+      auto extInfo = closureEffects(closure);
 
       // Closure expressions always have function type. In cases where a
       // parameter or return type is omitted, a fresh type variable is used to
@@ -2596,9 +2594,12 @@ namespace {
       return CS.getType(expr->getClosureBody());
     }
 
-    /// Walk a closure AST to determine if it can throw.
-    bool closureCanThrow(ClosureExpr *expr) {
-      // A walker that looks for 'try' or 'throw' expressions
+    /// Walk a closure AST to determine its effects.
+    ///
+    /// \returns a function's extended info describing the effects, as
+    /// determined syntactically.
+    FunctionType::ExtInfo closureEffects(ClosureExpr *expr) {
+      // A walker that looks for 'try' and 'throw' expressions
       // that aren't nested within closures, nested declarations,
       // or exhaustive catches.
       class FindInnerThrows : public ASTWalker {
@@ -2743,18 +2744,62 @@ namespace {
 
         bool foundThrow() { return FoundThrow; }
       };
-      
-      if (expr->getThrowsLoc().isValid())
-        return true;
-      
+
+      // A walker that looks for 'async' and 'await' expressions
+      // that aren't nested within closures or nested declarations.
+      class FindInnerAsync : public ASTWalker {
+        bool FoundAsync = false;
+
+        std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
+          // If we've found an 'await', record it and terminate the traversal.
+          if (isa<AwaitExpr>(expr)) {
+            FoundAsync = true;
+            return { false, nullptr };
+          }
+
+          // Do not recurse into other closures.
+          if (isa<ClosureExpr>(expr))
+            return { false, expr };
+
+          return { true, expr };
+        }
+
+        bool walkToDeclPre(Decl *decl) override {
+          // Do not walk into function or type declarations.
+          if (!isa<PatternBindingDecl>(decl))
+            return false;
+
+          return true;
+        }
+
+      public:
+        bool foundAsync() { return FoundAsync; }
+      };
+
+      // If either 'throws' or 'async' was explicitly specified, use that
+      // set of effects.
+      bool throws = expr->getThrowsLoc().isValid();
+      bool async = expr->getAsyncLoc().isValid();
+      if (throws || async) {
+        return ASTExtInfoBuilder()
+          .withThrows(throws)
+          .withAsync(async)
+          .build();
+      }
+
+      // Scan the body to determine the effects.
       auto body = expr->getBody();
-      
       if (!body)
-        return false;
-      
-      auto tryFinder = FindInnerThrows(CS, expr);
-      body->walk(tryFinder);
-      return tryFinder.foundThrow();
+        return FunctionType::ExtInfo();
+
+      auto throwFinder = FindInnerThrows(CS, expr);
+      body->walk(throwFinder);
+      auto asyncFinder = FindInnerAsync();
+      body->walk(asyncFinder);
+      return ASTExtInfoBuilder()
+        .withThrows(throwFinder.foundThrow())
+        .withAsync(asyncFinder.foundAsync())
+        .build();
     }
 
     Type visitClosureExpr(ClosureExpr *closure) {

lib/Sema/DebuggerTestingTransform.cpp

@@ -227,8 +227,8 @@ class DebuggerTestingTransform : public ASTWalker {
     auto *Params = ParameterList::createEmpty(Ctx);
     auto *Closure = new (Ctx)
         ClosureExpr(SourceRange(), nullptr, Params, SourceLoc(), SourceLoc(),
-                    SourceLoc(), nullptr, DF.getNextDiscriminator(),
-                    getCurrentDeclContext());
+                    SourceLoc(), SourceLoc(), nullptr,
+                    DF.getNextDiscriminator(), getCurrentDeclContext());
     Closure->setImplicit(true);
 
     // TODO: Save and return the value of $OriginalExpr.

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -393,9 +393,9 @@ deriveBodyDifferentiable_zeroTangentVectorInitializer(
 
     auto *closureParams = ParameterList::createEmpty(C);
     auto *closure = new (C) ClosureExpr(
-        SourceRange(), /*capturedSelfDecl*/ nullptr, closureParams, SourceLoc(),
-        SourceLoc(), SourceLoc(), TypeExpr::createImplicit(resultTy, C),
-        discriminator, funcDecl);
+        SourceRange(), /*capturedSelfDecl*/ nullptr, closureParams,
+        SourceLoc(), SourceLoc(), SourceLoc(), SourceLoc(),
+        TypeExpr::createImplicit(resultTy, C), discriminator, funcDecl);
     closure->setImplicit();
     auto *closureReturn = new (C) ReturnStmt(SourceLoc(), zeroExpr, true);
     auto *closureBody =
@@ -504,8 +504,8 @@ deriveBodyDifferentiable_zeroTangentVectorInitializer(
   auto *closureParams = ParameterList::createEmpty(C);
   auto *closure = new (C) ClosureExpr(
       SourceRange(), /*capturedSelfDecl*/ nullptr, closureParams, SourceLoc(),
-      SourceLoc(), SourceLoc(), TypeExpr::createImplicit(resultTy, C),
-      discriminator, funcDecl);
+      SourceLoc(), SourceLoc(), SourceLoc(),
+      TypeExpr::createImplicit(resultTy, C), discriminator, funcDecl);
   closure->setImplicit();
   auto *closureReturn = new (C) ReturnStmt(SourceLoc(), callExpr, true);
   auto *closureBody =

test/Parse/async-syntax.swift

@@ -14,3 +14,9 @@ func testTypeExprs() {
 func testAwaitOperator() async {
   let _ = __await asyncGlobal1()
 }
+
+func testAsyncClosure() {
+  let _ = { () async in 5 }
+  let _ = { () throws in 5 }
+  let _ = { () async throws in 5 }
+}

test/expr/unary/async_await.swift

@@ -48,3 +48,19 @@ func testAutoclosure() async {
   __await acceptAutoclosureNonAsync(getInt()) // expected-error{{'async' in an autoclosure that does not support concurrency}}
   // expected-warning@-1{{no calls to 'async' functions occur within 'await' expression}}
 }
+
+// Test inference of 'async' from the body of a closure.
+func testClosure() {
+  let closure = {
+     __await getInt()
+  }
+
+  let _: () -> Int = closure // expected-error{{cannot convert value of type '() async -> Int' to specified type '() -> Int'}}
+
+  let closure2 = { () async -> Int in
+    print("here")
+    return __await getInt()
+  }
+
+  let _: () -> Int = closure2 // expected-error{{cannot convert value of type '() async -> Int' to specified type '() -> Int'}}
+}

utils/gyb_syntax_support/ExprNodes.py

@@ -390,6 +390,9 @@
                        Child('SimpleInput', kind='ClosureParamList'),
                        Child('Input', kind='ParameterClause'),
                    ]),
+             Child('AsyncKeyword', kind='IdentifierToken',
+                   classification='Keyword',
+                   text_choices=['async'], is_optional=True),
              Child('ThrowsTok', kind='ThrowsToken', is_optional=True),
              Child('Output', kind='ReturnClause', is_optional=True),
              Child('InTok', kind='InToken'),