Surface error for plain return statement in coroutine earlier

Author: ivanaivanovska

clang/lib/Sema/SemaCoroutine.cpp

@@ -1120,7 +1120,7 @@ void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
 
   // [stmt.return.coroutine]p1:
   //   A coroutine shall not enclose a return statement ([stmt.return]).
-  if (Fn->FirstReturnLoc.isValid()) {
+  if (Fn->FirstReturnLoc.isValid() && Fn->FirstReturnLoc < Fn->FirstCoroutineStmtLoc) {
     assert(Fn->FirstCoroutineStmtLoc.isValid() &&
                    "first coroutine location not set");
     Diag(Fn->FirstReturnLoc, diag::err_return_in_coroutine);

clang/lib/Sema/SemaStmt.cpp

@@ -3747,6 +3747,16 @@ Sema::ActOnReturnStmt(SourceLocation ReturnLoc, Expr *RetValExp,
         Diag(ReturnLoc, diag::err_acc_branch_in_out_compute_construct)
         << /*return*/ 1 << /*out of */ 0);
 
+  // using plain return in a coroutine is not allowed.
+  FunctionScopeInfo *FSI = getCurFunction();
+  if (getLangOpts().Coroutines && FSI->isCoroutine()) {
+    assert(FSI->FirstCoroutineStmtLoc.isValid() &&
+                  "first coroutine location not set");
+    Diag(ReturnLoc, diag::err_return_in_coroutine);
+    Diag(FSI->FirstCoroutineStmtLoc, diag::note_declared_coroutine_here)
+            << FSI->getFirstCoroutineStmtKeyword();
+  }
+
   StmtResult R =
       BuildReturnStmt(ReturnLoc, RetVal.get(), /*AllowRecovery=*/true);
   if (R.isInvalid() || ExprEvalContexts.back().isDiscardedStatementContext())

clang/test/SemaCXX/coroutines.cpp

@@ -154,12 +154,15 @@ namespace std {
 template <class PromiseType = void>
 struct coroutine_handle {
   static coroutine_handle from_address(void *) noexcept;
+  static coroutine_handle from_promise(PromiseType &promise);
 };
 template <>
 struct coroutine_handle<void> {
   template <class PromiseType>
   coroutine_handle(coroutine_handle<PromiseType>) noexcept;
   static coroutine_handle from_address(void *) noexcept;
+  template <class PromiseType>
+  static coroutine_handle from_promise(PromiseType &promise);
 };
 } // namespace std
 
@@ -291,6 +294,38 @@ void mixed_coreturn_template2(bool b, T) {
     return; // expected-error {{not allowed in coroutine}}
 }
 
+struct promise_handle;
+
+struct Handle : std::coroutine_handle<promise_handle> { // expected-note 2{{candidate constructor (the implicit copy constructor) not viable}}
+    // expected-note@-1 2{{candidate constructor (the implicit move constructor) not viable}}
+    using promise_type = promise_handle;
+};
+
+struct promise_handle {
+    Handle get_return_object() noexcept {
+      { return Handle(std::coroutine_handle<Handle::promise_type>::from_promise(*this)); }
+    }
+    suspend_never initial_suspend() const noexcept { return {}; }
+    suspend_never final_suspend() const noexcept { return {}; }
+    void return_void() const noexcept {}
+    void unhandled_exception() const noexcept {}
+};
+
+Handle mixed_return_value() {
+  co_await a; // expected-note {{function is a coroutine due to use of 'co_await' here}}
+  return 0; // expected-error {{return statement not allowed in coroutine}}
+  // expected-error@-1 {{no viable conversion from returned value of type}}
+}
+
+Handle mixed_return_value_return_first(bool b) {
+   if (b) {
+        return 0; // expected-error {{no viable conversion from returned value of type}}
+        // expected-error@-1 {{return statement not allowed in coroutine}}
+    }
+    co_await a; // expected-note {{function is a coroutine due to use of 'co_await' here}}
+    co_return 0; // expected-error {{no member named 'return_value' in 'promise_handle'}}
+}
+
 struct CtorDtor {
   CtorDtor() {
     co_yield 0; // expected-error {{'co_yield' cannot be used in a constructor}}