Merge pull request #33399 from DougGregor/concurrency-async-call-restrictions

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

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}}
+}