[Concurrency] Implement restrictions on calls to 'async' functions.

Implement missing restrictions on calls to 'async':
* Diagnose async calls/uses of await in illegal contexts (such as
default arguments)
* Diagnose async calls/uses of await in functions/closures that are not
asynchronous themselves
* Handle autoclosure arguments as their own separate contexts (so
'await' has to go on the argument), which differs from error handling
(where the 'try' can go outside) because we want to be more particular
about marking the specific suspension points.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -4038,8 +4038,24 @@ NOTE(note_disable_error_propagation,none,
       "did you mean to disable error propagation?", ())
 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'", ())
 WARNING(no_async_in_await,none,
         "no calls to 'async' functions occur within 'await' expression", ())
+ERROR(async_call_in_illegal_context,none,
+      "'async' call cannot occur in "
+      "%select{<<ERROR>>|a default argument|a property initializer|a global variable initializer|an enum case raw value|a catch pattern|a catch guard expression|a defer body}0",
+      (unsigned))
+ERROR(await_in_illegal_context,none,
+      "'await' operation cannot occur in "
+      "%select{<<ERROR>>|a default argument|a property initializer|a global variable initializer|an enum case raw value|a catch pattern|a catch guard expression|a defer body}0",
+      (unsigned))
+ERROR(async_in_nonasync_function,none,
+      "%select{'async'|'await'}0 in %select{a function|an autoclosure}1 that "
+      "does not support concurrency",
+      (bool, bool))
+NOTE(note_add_async_to_function,none,
+     "add 'async' to function %0 to make it asynchronous", (DeclName))
 
 WARNING(no_throw_in_try,none,
         "no calls to throwing functions occur within 'try' expression", ())

lib/Sema/TypeCheckEffects.cpp

@@ -857,6 +857,7 @@ class Context {
   Kind TheKind;
   Optional<AnyFunctionRef> Function;
   bool HandlesErrors = false;
+  bool HandlesAsync = false;
 
   /// Whether error-handling queries should ignore the function context, e.g.,
   /// for autoclosure and rethrows checks.
@@ -870,9 +871,10 @@ class Context {
     assert(TheKind != Kind::PotentiallyHandled);
   }
 
-  explicit Context(bool handlesErrors, Optional<AnyFunctionRef> function)
+  explicit Context(bool handlesErrors, bool handlesAsync,
+                   Optional<AnyFunctionRef> function)
     : TheKind(Kind::PotentiallyHandled), Function(function),
-      HandlesErrors(handlesErrors) { }
+      HandlesErrors(handlesErrors), HandlesAsync(handlesAsync) { }
 
 public:
   /// Whether this is a function that rethrows.
@@ -910,7 +912,7 @@ class Context {
 
   static Context forTopLevelCode(TopLevelCodeDecl *D) {
     // Top-level code implicitly handles errors and 'async' calls.
-    return Context(/*handlesErrors=*/true, None);
+    return Context(/*handlesErrors=*/true, /*handlesAsync=*/true, None);
   }
 
   static Context forFunction(AbstractFunctionDecl *D) {
@@ -930,8 +932,7 @@ class Context {
       }
     }
 
-    bool handlesErrors = D->hasThrows();
-    return Context(handlesErrors, AnyFunctionRef(D));
+    return Context(D->hasThrows(), D->hasAsync(), AnyFunctionRef(D));
   }
 
   static Context forDeferBody() {
@@ -956,12 +957,15 @@ class Context {
   static Context forClosure(AbstractClosureExpr *E) {
     // Determine whether the closure has throwing function type.
     bool closureTypeThrows = true;
+    bool closureTypeIsAsync = true;
     if (auto closureType = E->getType()) {
-      if (auto fnType = closureType->getAs<AnyFunctionType>())
+      if (auto fnType = closureType->getAs<AnyFunctionType>()) {
         closureTypeThrows = fnType->isThrowing();
+        closureTypeIsAsync = fnType->isAsync();
+      }
     }
 
-    return Context(closureTypeThrows, AnyFunctionRef(E));
+    return Context(closureTypeThrows, closureTypeIsAsync, AnyFunctionRef(E));
   }
 
   static Context forCatchPattern(CaseStmt *S) {
@@ -1013,6 +1017,10 @@ class Context {
     llvm_unreachable("bad error kind");
   }
 
+  bool handlesAsync() const {
+    return HandlesAsync;
+  }
+
   DeclContext *getRethrowsDC() const {
     if (!isRethrows())
       return nullptr;
@@ -1182,7 +1190,6 @@ class Context {
     case Kind::DeferBody:
       diagnoseThrowInIllegalContext(Diags, E, getKind());
       return;
-
     }
     llvm_unreachable("bad context kind");
   }
@@ -1211,6 +1218,64 @@ class Context {
     }
     llvm_unreachable("bad context kind");
   }
+
+  void diagnoseUncoveredAsyncSite(ASTContext &ctx, ASTNode node) {
+    SourceRange highlight;
+
+    // Generate more specific messages in some cases.
+    if (auto apply = dyn_cast_or_null<ApplyExpr>(node.dyn_cast<Expr*>()))
+      highlight = apply->getSourceRange();
+
+    auto diag = diag::async_call_without_await;
+    if (isAutoClosure())
+      diag = diag::async_call_without_await_in_autoclosure;
+    ctx.Diags.diagnose(node.getStartLoc(), diag)
+        .highlight(highlight);
+  }
+
+  void diagnoseAsyncInIllegalContext(DiagnosticEngine &Diags, ASTNode node) {
+    if (auto *e = node.dyn_cast<Expr*>()) {
+      if (isa<ApplyExpr>(e)) {
+        Diags.diagnose(e->getLoc(), diag::async_call_in_illegal_context,
+                       static_cast<unsigned>(getKind()));
+        return;
+      }
+    }
+
+    Diags.diagnose(node.getStartLoc(), diag::await_in_illegal_context,
+                   static_cast<unsigned>(getKind()));
+  }
+
+  void maybeAddAsyncNote(DiagnosticEngine &Diags) {
+    if (!Function)
+      return;
+
+    auto func = dyn_cast_or_null<FuncDecl>(Function->getAbstractFunctionDecl());
+    if (!func)
+      return;
+
+    func->diagnose(diag::note_add_async_to_function, func->getName());
+  }
+
+  void diagnoseUnhandledAsyncSite(DiagnosticEngine &Diags, ASTNode node) {
+    switch (getKind()) {
+    case Kind::PotentiallyHandled:
+      Diags.diagnose(node.getStartLoc(), diag::async_in_nonasync_function,
+                     node.isExpr(ExprKind::Await), isAutoClosure());
+      maybeAddAsyncNote(Diags);
+      return;
+
+    case Kind::EnumElementInitializer:
+    case Kind::GlobalVarInitializer:
+    case Kind::IVarInitializer:
+    case Kind::DefaultArgument:
+    case Kind::CatchPattern:
+    case Kind::CatchGuard:
+    case Kind::DeferBody:
+      diagnoseAsyncInIllegalContext(Diags, node);
+      return;
+    }
+  }
 };
 
 /// A class to walk over a local context and validate the correctness
@@ -1322,6 +1387,12 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
       Self.MaxThrowingKind = ThrowingKind::None;
     }
 
+    void resetCoverageForAutoclosureBody() {
+      Self.Flags.clear(ContextFlags::IsAsyncCovered);
+      Self.Flags.clear(ContextFlags::HasAnyAsyncSite);
+      Self.Flags.clear(ContextFlags::HasAnyAwait);
+    }
+
     void resetCoverageForDoCatch() {
       Self.Flags.reset();
       Self.MaxThrowingKind = ThrowingKind::None;
@@ -1409,6 +1480,7 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
   ShouldRecurse_t checkAutoClosure(AutoClosureExpr *E) {
     ContextScope scope(*this, Context::forClosure(E));
     scope.enterSubFunction();
+    scope.resetCoverageForAutoclosureBody();
     E->getBody()->walk(*this);
     scope.preserveCoverageFromAutoclosureBody();
     return ShouldNotRecurse;
@@ -1572,17 +1644,14 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     if (classification.isAsync()) {
       // Remember that we've seen an async call.
       Flags.set(ContextFlags::HasAnyAsyncSite);
-      
+
+      // Diagnose async calls in a context that doesn't handle async.
+      if (!CurContext.handlesAsync()) {
+        CurContext.diagnoseUnhandledAsyncSite(Ctx.Diags, E);
+      }
       // Diagnose async calls that are outside of an await context.
-      if (!Flags.has(ContextFlags::IsAsyncCovered)) {
-        SourceRange highlight;
-        
-        // Generate more specific messages in some cases.
-        if (auto e = dyn_cast_or_null<ApplyExpr>(E.dyn_cast<Expr*>()))
-          highlight = e->getSourceRange();
-
-        Ctx.Diags.diagnose(E.getStartLoc(), diag::async_call_without_await)
-              .highlight(highlight);
+      else if (!Flags.has(ContextFlags::IsAsyncCovered)) {
+        CurContext.diagnoseUncoveredAsyncSite(Ctx, E);
       }
     }
 
@@ -1626,10 +1695,16 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     scope.enterAwait();
 
     E->getSubExpr()->walk(*this);
-    
-    // Warn about 'await' expressions that weren't actually needed.
-    if (!Flags.has(ContextFlags::HasAnyAsyncSite))
-      Ctx.Diags.diagnose(E->getAwaitLoc(), diag::no_async_in_await);
+
+    // Warn about 'await' expressions that weren't actually needed, unless of
+    // course we're in a context that could never handle an 'async'. Then, we
+    // produce an error.
+    if (!Flags.has(ContextFlags::HasAnyAsyncSite)) {
+      if (CurContext.handlesAsync())
+        Ctx.Diags.diagnose(E->getAwaitLoc(), diag::no_async_in_await);
+      else
+        CurContext.diagnoseUnhandledAsyncSite(Ctx.Diags, E);
+    }
 
     // Inform the parent of the walk that an 'await' exists here.
     scope.preserveCoverageFromAwaitOperand();

test/expr/unary/async_await.swift

@@ -8,3 +8,43 @@ func test1(asyncfp : () async -> Int, fp : () -> Int) async {
   _ = asyncfp() // expected-error {{call is 'async' but is not marked with 'await'}}
 }
 
+func getInt() async -> Int { return 5 }
+
+// Locations where "await" is prohibited.
+func test2(
+  defaulted: Int = __await getInt() // expected-error{{'async' call cannot occur in a default argument}}
+) async {
+  defer {
+    _ = __await getInt() // expected-error{{'async' call cannot occur in a defer body}}
+  }
+  print("foo")
+}
+
+func test3() { // expected-note{{add 'async' to function 'test3()' to make it asynchronous}}
+  _ = __await getInt() // expected-error{{'async' in a function that does not support concurrency}}
+}
+
+enum SomeEnum: Int {
+case foo = __await 5 // expected-error{{raw value for enum case must be a literal}}
+}
+
+struct SomeStruct {
+  var x = __await getInt() // expected-error{{'async' call cannot occur in a property initializer}}
+  static var y = __await getInt() // expected-error{{'async' call cannot occur in a global variable initializer}}
+}
+
+func acceptAutoclosureNonAsync(_: @autoclosure () -> Int) { }
+func acceptAutoclosureAsync(_: @autoclosure () async -> Int) { }
+
+func testAutoclosure() async {
+  acceptAutoclosureAsync(getInt()) // expected-error{{call is 'async' in an autoclosure argument is not marked with 'await'}}
+  acceptAutoclosureNonAsync(getInt()) // expected-error{{'async' in an autoclosure that does not support concurrency}}
+
+  acceptAutoclosureAsync(__await getInt())
+  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'}}
+  // expected-warning@-1{{no calls to 'async' functions occur within 'await' expression}}
+  __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}}
+}