3 files changed, 104 insertions, 59 deletions

diff --git a/clang/include/clang/Basic/DiagnosticSemaKinds.td b/clang/include/clang/Basic/DiagnosticSemaKinds.td
index f0c8605ffba..29a3e089913 100644
--- a/clang/include/clang/Basic/DiagnosticSemaKinds.td
+++ b/clang/include/clang/Basic/DiagnosticSemaKinds.td
@@ -8628,14 +8628,11 @@ def err_coroutine_unevaluated_context : Error<
   "'%0' cannot be used in an unevaluated context">;
 def err_coroutine_outside_function : Error<
   "'%0' cannot be used outside a function">;
-def err_coroutine_ctor_dtor : Error<
-  "'%1' cannot be used in a %select{constructor|destructor}0">;
-def err_coroutine_constexpr : Error<
-  "'%0' cannot be used in a constexpr function">;
-def err_coroutine_main : Error<
-  "'main' cannot be a coroutine">;
-def err_coroutine_varargs : Error<
-  "'%0' cannot be used in a varargs function">;
+def err_coroutine_invalid_func_context : Error<
+  "'%1' cannot be used in %select{a constructor|a destructor"
+  "|a copy assignment operator|a move assignment operator|the 'main' function"
+  "|a constexpr function|a function with a deduced return type"
+  "|a varargs function}0">;
 def ext_coroutine_without_co_await_co_yield : ExtWarn<
   "'co_return' used in a function "
   "that uses neither 'co_await' nor 'co_yield'">,
diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index 873a830d975..3109358df46 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -99,67 +99,98 @@ static QualType lookupPromiseType(Sema &S, const FunctionProtoType *FnType,
   return PromiseType;
 }
 
-/// Check that this is a context in which a coroutine suspension can appear.
-static FunctionScopeInfo *
-checkCoroutineContext(Sema &S, SourceLocation Loc, StringRef Keyword) {
+static bool isValidCoroutineContext(Sema &S, SourceLocation Loc,
+                                    StringRef Keyword) {
   // 'co_await' and 'co_yield' are not permitted in unevaluated operands.
   if (S.isUnevaluatedContext()) {
     S.Diag(Loc, diag::err_coroutine_unevaluated_context) << Keyword;
-    return nullptr;
+    return false;
   }
 
   // Any other usage must be within a function.
-  // FIXME: Reject a coroutine with a deduced return type.
   auto *FD = dyn_cast<FunctionDecl>(S.CurContext);
   if (!FD) {
     S.Diag(Loc, isa<ObjCMethodDecl>(S.CurContext)
                     ? diag::err_coroutine_objc_method
                     : diag::err_coroutine_outside_function) << Keyword;
-  } else if (isa<CXXConstructorDecl>(FD) || isa<CXXDestructorDecl>(FD)) {
-    // Coroutines TS [special]/6:
-    //   A special member function shall not be a coroutine.
-    //
-    // FIXME: We assume that this really means that a coroutine cannot
-    //        be a constructor or destructor.
-    S.Diag(Loc, diag::err_coroutine_ctor_dtor) 
-      << isa<CXXDestructorDecl>(FD) << Keyword;
-  } else if (FD->isConstexpr()) {
-    S.Diag(Loc, diag::err_coroutine_constexpr) << Keyword;
-  } else if (FD->isVariadic()) {
-    S.Diag(Loc, diag::err_coroutine_varargs) << Keyword;
-  } else if (FD->isMain()) {
-    S.Diag(FD->getLocStart(), diag::err_coroutine_main);
-    S.Diag(Loc, diag::note_declared_coroutine_here)
-        << (Keyword == "co_await" ? 0 : 
-            Keyword == "co_yield" ? 1 : 2);
-  } else {
-    auto *ScopeInfo = S.getCurFunction();
-    assert(ScopeInfo && "missing function scope for function");
-
-    // If we don't have a promise variable, build one now.
-    if (!ScopeInfo->CoroutinePromise) {
-      QualType T =
-          FD->getType()->isDependentType()
-              ? S.Context.DependentTy
-              : lookupPromiseType(S, FD->getType()->castAs<FunctionProtoType>(),
-                                  Loc);
-      if (T.isNull())
-        return nullptr;
-
-      // Create and default-initialize the promise.
-      ScopeInfo->CoroutinePromise =
-          VarDecl::Create(S.Context, FD, FD->getLocation(), FD->getLocation(),
-                          &S.PP.getIdentifierTable().get("__promise"), T,
-                          S.Context.getTrivialTypeSourceInfo(T, Loc), SC_None);
-      S.CheckVariableDeclarationType(ScopeInfo->CoroutinePromise);
-      if (!ScopeInfo->CoroutinePromise->isInvalidDecl())
-        S.ActOnUninitializedDecl(ScopeInfo->CoroutinePromise, false);
-    }
+    return false;
+  }
+
+  // An enumeration for mapping the diagnostic type to the correct diagnostic
+  // selection index.
+  enum InvalidFuncDiag {
+    DiagCtor = 0,
+    DiagDtor,
+    DiagCopyAssign,
+    DiagMoveAssign,
+    DiagMain,
+    DiagConstexpr,
+    DiagAutoRet,
+    DiagVarargs,
+  };
+  bool Diagnosed = false;
+  auto DiagInvalid = [&](InvalidFuncDiag ID) {
+    S.Diag(Loc, diag::err_coroutine_invalid_func_context) << ID << Keyword;
+    Diagnosed = true;
+    return false;
+  };
+
+  // Diagnose when a constructor, destructor, copy/move assignment operator,
+  // or the function 'main' are declared as a coroutine.
+  auto *MD = dyn_cast<CXXMethodDecl>(FD);
+  if (MD && isa<CXXConstructorDecl>(MD))
+    return DiagInvalid(DiagCtor);
+  else if (MD && isa<CXXDestructorDecl>(MD))
+    return DiagInvalid(DiagDtor);
+  else if (MD && MD->isCopyAssignmentOperator())
+    return DiagInvalid(DiagCopyAssign);
+  else if (MD && MD->isMoveAssignmentOperator())
+    return DiagInvalid(DiagMoveAssign);
+  else if (FD->isMain())
+    return DiagInvalid(DiagMain);
+
+  // Emit a diagnostics for each of the following conditions which is not met.
+  if (FD->isConstexpr())
+    DiagInvalid(DiagConstexpr);
+  if (FD->getReturnType()->isUndeducedType())
+    DiagInvalid(DiagAutoRet);
+  if (FD->isVariadic())
+    DiagInvalid(DiagVarargs);
+
+  return !Diagnosed;
+}
+
+/// Check that this is a context in which a coroutine suspension can appear.
+static FunctionScopeInfo *checkCoroutineContext(Sema &S, SourceLocation Loc,
+                                                StringRef Keyword) {
+  if (!isValidCoroutineContext(S, Loc, Keyword))
+    return nullptr;
 
-    return ScopeInfo;
+  assert(isa<FunctionDecl>(S.CurContext) && "not in a function scope");
+  auto *FD = cast<FunctionDecl>(S.CurContext);
+  auto *ScopeInfo = S.getCurFunction();
+  assert(ScopeInfo && "missing function scope for function");
+
+  // If we don't have a promise variable, build one now.
+  if (!ScopeInfo->CoroutinePromise) {
+    QualType T = FD->getType()->isDependentType()
+                     ? S.Context.DependentTy
+                     : lookupPromiseType(
+                           S, FD->getType()->castAs<FunctionProtoType>(), Loc);
+    if (T.isNull())
+      return nullptr;
+
+    // Create and default-initialize the promise.
+    ScopeInfo->CoroutinePromise =
+        VarDecl::Create(S.Context, FD, FD->getLocation(), FD->getLocation(),
+                        &S.PP.getIdentifierTable().get("__promise"), T,
+                        S.Context.getTrivialTypeSourceInfo(T, Loc), SC_None);
+    S.CheckVariableDeclarationType(ScopeInfo->CoroutinePromise);
+    if (!ScopeInfo->CoroutinePromise->isInvalidDecl())
+      S.ActOnUninitializedDecl(ScopeInfo->CoroutinePromise, false);
   }
 
-  return nullptr;
+  return ScopeInfo;
 }
 
 static Expr *buildBuiltinCall(Sema &S, SourceLocation Loc, Builtin::ID Id,
diff --git a/clang/test/SemaCXX/coroutines.cpp b/clang/test/SemaCXX/coroutines.cpp
index f170201a2bc..a22383cd566 100644
--- a/clang/test/SemaCXX/coroutines.cpp
+++ b/clang/test/SemaCXX/coroutines.cpp
@@ -178,7 +178,19 @@ struct CtorDtor {
   }
   // FIXME: The spec says this is ill-formed.
   void operator=(CtorDtor&) {
-    co_yield 0;
+    co_yield 0; // expected-error {{'co_yield' cannot be used in a copy assignment operator}}
+  }
+  void operator=(CtorDtor const &) {
+    co_yield 0; // expected-error {{'co_yield' cannot be used in a copy assignment operator}}
+  }
+  void operator=(CtorDtor &&) {
+    co_await a; // expected-error {{'co_await' cannot be used in a move assignment operator}}
+  }
+  void operator=(CtorDtor const &&) {
+    co_await a; // expected-error {{'co_await' cannot be used in a move assignment operator}}
+  }
+  void operator=(int) {
+    co_await a; // OK. Not a special member
   }
 };
 
@@ -191,14 +203,19 @@ void unevaluated() {
   typeid(co_yield a); // expected-error {{cannot be used in an unevaluated context}}
 }
 
-constexpr void constexpr_coroutine() {
+constexpr auto constexpr_deduced_return_coroutine() {
   co_yield 0; // expected-error {{'co_yield' cannot be used in a constexpr function}}
+  // expected-error@-1 {{'co_yield' cannot be used in a function with a deduced return type}}
 }
 
 void varargs_coroutine(const char *, ...) {
   co_await a; // expected-error {{'co_await' cannot be used in a varargs function}}
 }
 
+auto deduced_return_coroutine() {
+  co_await a; // expected-error {{'co_await' cannot be used in a function with a deduced return type}}
+}
+
 struct outer {};
 
 namespace dependent_operator_co_await_lookup {
@@ -355,6 +372,6 @@ coro<bad_promise_8> calls_set_exception() {
 template<> struct std::experimental::coroutine_traits<int, int, const char**>
 { using promise_type = promise; };
 
-int main(int, const char**) { // expected-error {{'main' cannot be a coroutine}}
-  co_await a; // expected-note {{function is a coroutine due to use of 'co_await' here}}
+int main(int, const char**) {
+  co_await a; // expected-error {{'co_await' cannot be used in the 'main' function}}
 }