[Concurrency] @asyncHandler functions are an async context.

Allow async calls and await expressions within @asyncHandler functions.
If we see an async call or await expression in a function that is not
an async context, also suggest adding @asyncHandler if the function
meets the semantic constraints.

Author: DougGregor

include/swift/AST/Decl.h

@@ -5888,6 +5888,13 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// type of the function will be `async` as well.
   bool hasAsync() const { return Bits.AbstractFunctionDecl.Async; }
 
+  /// Returns true if the function is a suitable 'async' context.
+  ///
+  /// Functions that are an 'async' context can make calls to 'async' functions.
+  bool isAsyncContext() const {
+    return hasAsync() || getAttrs().hasAttribute<AsyncHandlerAttr>();
+  }
+
   /// Returns true if the function body throws.
   bool hasThrows() const { return Bits.AbstractFunctionDecl.Throws; }
 

include/swift/AST/DiagnosticsSema.def

@@ -4047,7 +4047,7 @@ WARNING(no_throw_in_do_with_catch,none,
 ERROR(async_call_without_await,none,
       "call is 'async' but is not marked with 'await'", ())
 ERROR(async_call_without_await_in_autoclosure,none,
-      "call is 'async' in an autoclosure argument is not marked with 'await'", ())
+      "call is 'async' in an autoclosure argument that is not marked with 'await'", ())
 WARNING(no_async_in_await,none,
         "no calls to 'async' functions occur within 'await' expression", ())
 ERROR(async_call_in_illegal_context,none,
@@ -4064,6 +4064,8 @@ ERROR(async_in_nonasync_function,none,
       (bool, bool))
 NOTE(note_add_async_to_function,none,
      "add 'async' to function %0 to make it asynchronous", (DeclName))
+NOTE(note_add_asynchandler_to_function,none,
+     "add '@asyncHandler' to function %0 to create an implicit asynchronous context", (DeclName))
 ERROR(not_objc_function_async,none,
       "'async' function cannot be represented in Objective-C", ())
 NOTE(not_objc_function_type_async,none,

lib/Sema/TypeCheckAttr.cpp

@@ -41,6 +41,66 @@
 
 using namespace swift;
 
+/// Check whether the @asyncHandler attribute can be applied to the given
+/// function declaration.
+///
+/// \param diagnose Whether to emit a diagnostic when a problem is encountered.
+///
+/// \returns \c true if there was a problem with adding the attribute, \c false
+/// otherwise.
+static bool checkAsyncHandler(FuncDecl *func, bool diagnose) {
+  if (!func->getResultInterfaceType()->isVoid()) {
+    if (diagnose) {
+      func->diagnose(diag::asynchandler_returns_value)
+          .highlight(func->getBodyResultTypeLoc().getSourceRange());
+    }
+
+    return true;
+  }
+
+  if (func->hasThrows()) {
+    if (diagnose) {
+      func->diagnose(diag::asynchandler_throws)
+          .fixItRemove(func->getThrowsLoc());
+    }
+
+    return true;
+  }
+
+  if (func->hasAsync()) {
+    if (diagnose) {
+      func->diagnose(diag::asynchandler_async)
+          .fixItRemove(func->getAsyncLoc());
+    }
+
+    return true;
+  }
+
+  for (auto param : *func->getParameters()) {
+    if (param->isInOut()) {
+      if (diagnose) {
+        param->diagnose(diag::asynchandler_inout_parameter)
+            .fixItRemove(param->getSpecifierLoc());
+      }
+
+      return true;
+    }
+  }
+
+  if (func->isMutating()) {
+    if (diagnose) {
+      auto diag = func->diagnose(diag::asynchandler_mutating);
+      if (auto mutatingAttr = func->getAttrs().getAttribute<MutatingAttr>()) {
+        diag.fixItRemove(mutatingAttr->getRange());
+      }
+    }
+
+    return true;
+  }
+
+  return false;
+}
+
 namespace {
   /// This emits a diagnostic with a fixit to remove the attribute.
   template<typename ...ArgTypes>
@@ -270,42 +330,7 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
       return;
     }
 
-    if (!func->getResultInterfaceType()->isVoid()) {
-      func->diagnose(diag::asynchandler_returns_value)
-          .highlight(func->getBodyResultTypeLoc().getSourceRange());
-      attr->setInvalid();
-      return;
-    }
-
-    if (func->hasThrows()) {
-      func->diagnose(diag::asynchandler_throws)
-          .fixItRemove(func->getThrowsLoc());
-      attr->setInvalid();
-      return;
-    }
-
-    if (func->hasAsync()) {
-      func->diagnose(diag::asynchandler_async)
-          .fixItRemove(func->getAsyncLoc());
-      attr->setInvalid();
-      return;
-    }
-
-    for (auto param : *func->getParameters()) {
-      if (param->isInOut()) {
-        param->diagnose(diag::asynchandler_inout_parameter)
-            .fixItRemove(param->getSpecifierLoc());
-        attr->setInvalid();
-        return;
-      }
-    }
-
-    if (func->isMutating()) {
-      auto diag = Ctx.Diags.diagnose(func, diag::asynchandler_mutating);
-      diag.highlight(attr->getRangeWithAt());
-      if (auto mutatingAttr = func->getAttrs().getAttribute<MutatingAttr>()) {
-        diag.fixItRemove(mutatingAttr->getRange());
-      }
+    if (checkAsyncHandler(func, /*diagnose=*/true)) {
       attr->setInvalid();
       return;
     }
@@ -5201,3 +5226,13 @@ void AttributeChecker::visitTransposeAttr(TransposeAttr *attr) {
   // Set the resolved linearity parameter indices in the attribute.
   attr->setParameterIndices(linearParamIndices);
 }
+
+void swift::addAsyncNotes(FuncDecl *func) {
+  func->diagnose(diag::note_add_async_to_function, func->getName());
+
+  if (!checkAsyncHandler(func, /*diagnose=*/false)) {
+    func->diagnose(
+            diag::note_add_asynchandler_to_function, func->getName())
+        .fixItInsert(func->getAttributeInsertionLoc(false), "@asyncHandler ");
+  }
+}

lib/Sema/TypeCheckDecl.cpp

@@ -2155,11 +2155,10 @@ static Type validateParameterType(ParamDecl *decl) {
     if (auto fnType = Ty->getAs<FunctionType>()) {
       if (fnType->isAsync() &&
           !(isa<AbstractFunctionDecl>(dc) &&
-            cast<AbstractFunctionDecl>(dc)->hasAsync())) {
+            cast<AbstractFunctionDecl>(dc)->isAsyncContext())) {
         decl->diagnose(diag::async_autoclosure_nonasync_function);
-        if (auto func = dyn_cast<FuncDecl>(dc)) {
-          func->diagnose(diag::note_add_async_to_function, func->getName());
-        }
+        if (auto func = dyn_cast<FuncDecl>(dc))
+          addAsyncNotes(func);
       }
     }
   }

lib/Sema/TypeCheckEffects.cpp

@@ -932,7 +932,7 @@ class Context {
       }
     }
 
-    return Context(D->hasThrows(), D->hasAsync(), AnyFunctionRef(D));
+    return Context(D->hasThrows(), D->isAsyncContext(), AnyFunctionRef(D));
   }
 
   static Context forDeferBody() {
@@ -1254,7 +1254,7 @@ class Context {
     if (!func)
       return;
 
-    func->diagnose(diag::note_add_async_to_function, func->getName());
+    addAsyncNotes(func);
   }
 
   void diagnoseUnhandledAsyncSite(DiagnosticEngine &Diags, ASTNode node) {

lib/Sema/TypeChecker.h

@@ -1391,6 +1391,10 @@ void checkUnknownAttrRestrictions(
 /// it to later stages.
 void bindSwitchCasePatternVars(DeclContext *dc, CaseStmt *stmt);
 
+/// Add notes suggesting the addition of 'async' or '@asyncHandler', as
+/// appropriate, to a diagnostic for a function that isn't an async context.
+void addAsyncNotes(FuncDecl *func);
+
 } // end namespace swift
 
 #endif

test/attr/asynchandler.swift

@@ -1,6 +1,15 @@
 // RUN: %target-swift-frontend -typecheck -verify %s -enable-experimental-concurrency
 
-@asyncHandler func asyncHandler1() { }
+func globalAsyncFunction() async -> Int { 0 }
+
+@asyncHandler func asyncHandler1() {
+  // okay, it's an async context
+  let _ = await globalAsyncFunction()
+}
+
+@asyncHandler func asyncHandler2(fn: @autoclosure () async -> Int ) {
+  // okay, it's an async context
+}
 
 @asyncHandler
 func asyncHandlerBad1() -> Int { 0 }

test/expr/unary/async_await.swift

@@ -21,6 +21,7 @@ func test2(
 }
 
 func test3() { // expected-note{{add 'async' to function 'test3()' to make it asynchronous}}
+  // expected-note@-1{{add '@asyncHandler' to function 'test3()' to create an implicit asynchronous context}}{{1-1=@asyncHandler }}
   _ = await getInt() // expected-error{{'async' in a function that does not support concurrency}}
 }
 
@@ -36,7 +37,7 @@ struct SomeStruct {
 func acceptAutoclosureNonAsync(_: @autoclosure () -> Int) async { }
 func acceptAutoclosureAsync(_: @autoclosure () async -> Int) async { }
 
-func acceptAutoclosureNonAsyncBad(_: @autoclosure () async -> Int) { }
+func acceptAutoclosureNonAsyncBad(_: @autoclosure () async -> Int) -> Int { 0 }
 // expected-error@-1{{'async' autoclosure parameter in a non-'async' function}}
 // expected-note@-2{{add 'async' to function 'acceptAutoclosureNonAsyncBad' to make it asynchronous}}
 
@@ -46,13 +47,13 @@ struct HasAsyncBad {
 }
 
 func testAutoclosure() async {
-  await acceptAutoclosureAsync(getInt()) // expected-error{{call is 'async' in an autoclosure argument is not marked with 'await'}}
+  await acceptAutoclosureAsync(getInt()) // expected-error{{call is 'async' in an autoclosure argument that is not marked with 'await'}}
   await acceptAutoclosureNonAsync(getInt()) // expected-error{{'async' in an autoclosure that does not support concurrency}}
 
   await acceptAutoclosureAsync(await getInt())
   await acceptAutoclosureNonAsync(await getInt()) // expected-error{{'async' in an autoclosure that does not support concurrency}}
 
-  await acceptAutoclosureAsync(getInt()) // expected-error{{call is 'async' in an autoclosure argument is not marked with 'await'}}
+  await acceptAutoclosureAsync(getInt()) // expected-error{{call is 'async' in an autoclosure argument that is not marked with 'await'}}
   await acceptAutoclosureNonAsync(getInt()) // expected-error{{'async' in an autoclosure that does not support concurrency}}
 }