[clang] [Coroutines] Allow [[clang::coro_wrapper]] for class (PR #93268)

[Chuanqi Xu via cfe-commits]

https://github.com/ChuanqiXu9 created https://github.com/llvm/llvm-project/pull/93268

Previously we allow `[[clang::coro_wrapper]]` to be marked with function to allow it to not be checked by `[[clang::coro_return_type]]`.

But in our internal practice, there are classes can be return type of coroutines and non-coroutines and there are too many uses. It is slightly hard to mark every use with `[[clang::coro_wrapper]]`. So it is a sugar to allow we mark the class as  `[[clang::coro_wrapper]]`.

>From 2e05a2871ea5beb3159ab946c040cb836201181f Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <yedeng.yd at linux.alibaba.com>
Date: Fri, 24 May 2024 11:02:02 +0800
Subject: [PATCH] [Coroutines] Allow [[clang::coro_wrapper]] for class

Previously we allow `[[clang::coro_wrapper]]` to be marked with function
to allow it to not be checked by `[[clang::coro_return_type]]`.

But in our internal practice, there are classes can be return type of
coroutines and non-coroutines and there are too many uses. It is
slightly hard to mark every use with `[[clang::coro_wrapper]]`. So it is
a sugar to allow we mark the class as  `[[clang::coro_wrapper]]`.
---
 clang/include/clang/Basic/Attr.td             |  2 +-
 clang/lib/Sema/SemaDecl.cpp                   |  3 ++-
 ...a-attribute-supported-attributes-list.test |  2 +-
 .../SemaCXX/coro-return-type-and-wrapper.cpp  | 19 +++++++++++++++++++
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/clang/include/clang/Basic/Attr.td b/clang/include/clang/Basic/Attr.td
index e59cccccdd369..7933fd5b38d87 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -1192,7 +1192,7 @@ def CoroReturnType : InheritableAttr {
 
 def CoroWrapper : InheritableAttr {
   let Spellings = [Clang<"coro_wrapper">];
-  let Subjects = SubjectList<[Function]>;
+  let Subjects = SubjectList<[Function, CXXRecord]>;
   let LangOpts = [CPlusPlus];
   let Documentation = [CoroReturnTypeAndWrapperDoc];
   let SimpleHandler = 1;
diff --git a/clang/lib/Sema/SemaDecl.cpp b/clang/lib/Sema/SemaDecl.cpp
index 2a87b26f17a2b..9cae3a0022dc2 100644
--- a/clang/lib/Sema/SemaDecl.cpp
+++ b/clang/lib/Sema/SemaDecl.cpp
@@ -15998,7 +15998,8 @@ void Sema::CheckCoroutineWrapper(FunctionDecl *FD) {
   // Allow some_promise_type::get_return_object().
   if (CanBeGetReturnObject(FD) || CanBeGetReturnTypeOnAllocFailure(FD))
     return;
-  if (!FD->hasAttr<CoroWrapperAttr>())
+  if (!FD->hasAttr<CoroWrapperAttr>() &&
+      !RD->getUnderlyingDecl()->hasAttr<CoroWrapperAttr>())
     Diag(FD->getLocation(), diag::err_coroutine_return_type) << RD;
 }
 
diff --git a/clang/test/Misc/pragma-attribute-supported-attributes-list.test b/clang/test/Misc/pragma-attribute-supported-attributes-list.test
index 99732694f72a5..04dc2962ac04e 100644
--- a/clang/test/Misc/pragma-attribute-supported-attributes-list.test
+++ b/clang/test/Misc/pragma-attribute-supported-attributes-list.test
@@ -62,7 +62,7 @@
 // CHECK-NEXT: CoroLifetimeBound (SubjectMatchRule_record)
 // CHECK-NEXT: CoroOnlyDestroyWhenComplete (SubjectMatchRule_record)
 // CHECK-NEXT: CoroReturnType (SubjectMatchRule_record)
-// CHECK-NEXT: CoroWrapper (SubjectMatchRule_function)
+// CHECK-NEXT: CoroWrapper (SubjectMatchRule_function, SubjectMatchRule_record)
 // CHECK-NEXT: DLLExport (SubjectMatchRule_function, SubjectMatchRule_variable, SubjectMatchRule_record, SubjectMatchRule_objc_interface)
 // CHECK-NEXT: DLLImport (SubjectMatchRule_function, SubjectMatchRule_variable, SubjectMatchRule_record, SubjectMatchRule_objc_interface)
 // CHECK-NEXT: Destructor (SubjectMatchRule_function)
diff --git a/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp b/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp
index b08e1c9c065a0..3c174f3ffaecf 100644
--- a/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp
+++ b/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp
@@ -139,3 +139,22 @@ template<> class coroutine_traits<Task, int> {
 } // namespace std
 // expected-error at +1 {{neither a coroutine nor a coroutine wrapper}}
 Task foo(int) { return Task{}; }
+
+// Testing that we can add a `[[clang::coro_wrapper]]` attribute to the class.
+class [[clang::coro_return_type]] [[clang::coro_wrapper]] WrappedTask{
+public:
+  struct promise_type {
+    Task get_return_object() {
+      return {};
+    }
+    suspend_always initial_suspend();
+    suspend_always final_suspend() noexcept;
+    void unhandled_exception();
+  };
+};
+namespace std {
+template<> class coroutine_traits<WrappedTask, int> {
+    using promise_type = WrappedTask::promise_type;
+};
+} // namespace std
+WrappedTask wrapped_class(int) { return WrappedTask{}; }