[coroutines] Introduce [[clang::coro_return_type]] and [[clang::coro_wrapper]]

[usx95]

First step in the implementation of RFC (final approved doc).

This introduces the concepts of a coroutine return type and explicit coroutine wrapper functions.

[llvmbot]

@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-clang

Author: Utkarsh Saxena (usx95)

Changes

First step in the implementation of RFC (final proposal doc).

This introduces the concepts of a coroutine return type and explicit coroutine wrapper functions.

---

Full diff: https://github.com/llvm/llvm-project/pull/71945.diff

7 Files Affected:

• (modified) clang/include/clang/Basic/Attr.td (+16)
• (modified) clang/include/clang/Basic/AttrDocs.td (+67)
• (modified) clang/include/clang/Basic/DiagnosticSemaKinds.td (+4)
• (modified) clang/include/clang/Sema/Sema.h (+1)
• (modified) clang/lib/Sema/SemaDecl.cpp (+23-2)
• (modified) clang/test/Misc/pragma-attribute-supported-attributes-list.test (+2)
• (added) clang/test/SemaCXX/coro-return-type-and-wrapper.cpp (+56)

diff --git a/clang/include/clang/Basic/Attr.td b/clang/include/clang/Basic/Attr.td
index 31434565becaec6..f7a2b83b15ef5bc 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -1094,6 +1094,22 @@ def CoroOnlyDestroyWhenComplete : InheritableAttr {
   let SimpleHandler = 1;
 }
 
+def CoroReturnType : InheritableAttr {
+  let Spellings = [Clang<"coro_return_type">];
+  let Subjects = SubjectList<[CXXRecord]>;
+  let LangOpts = [CPlusPlus];
+  let Documentation = [CoroReturnTypeAndWrapperDoc];
+  let SimpleHandler = 1;
+}
+
+def CoroWrapper : InheritableAttr {
+  let Spellings = [Clang<"coro_wrapper">];
+  let Subjects = SubjectList<[Function]>;
+  let LangOpts = [CPlusPlus];
+  let Documentation = [CoroReturnTypeAndWrapperDoc];
+  let SimpleHandler = 1;
+}
+
 // OSObject-based attributes.
 def OSConsumed : InheritableParamAttr {
   let Spellings = [Clang<"os_consumed">];
diff --git a/clang/include/clang/Basic/AttrDocs.td b/clang/include/clang/Basic/AttrDocs.td
index fa6f6acd0c30e88..66c92bcaa2d4a4a 100644
--- a/clang/include/clang/Basic/AttrDocs.td
+++ b/clang/include/clang/Basic/AttrDocs.td
@@ -7482,3 +7482,70 @@ generation of the other destruction cases, optimizing the above `foo.destroy` to
 
   }];
 }
+
+
+def CoroReturnTypeAndWrapperDoc : Documentation {
+  let Category = DocCatDecl;
+  let Content = [{
+The `coro_only_destroy_when_complete` attribute should be marked on a C++ class. The coroutines
+whose return type is marked with the attribute are assumed to be destroyed only after the coroutine has
+reached the final suspend point.
+
+This is helpful for the optimizers to reduce the size of the destroy function for the coroutines.
+
+For example,
+
+.. code-block:: c++
+
+  A foo() {
+    dtor d;
+    co_await something();
+    dtor d1;
+    co_await something();
+    dtor d2;
+    co_return 43;
+  }
+
+The compiler may generate the following pseudocode:
+
+.. code-block:: c++
+
+  void foo.destroy(foo.Frame *frame) {
+    switch(frame->suspend_index()) {
+      case 1:
+        frame->d.~dtor();
+        break;
+      case 2:
+        frame->d.~dtor();
+        frame->d1.~dtor();
+        break;
+      case 3:
+        frame->d.~dtor();
+        frame->d1.~dtor();
+        frame->d2.~dtor();
+        break;
+      default: // coroutine completed or haven't started
+        break;
+    }
+
+    frame->promise.~promise_type();
+    delete frame;
+  }
+
+The `foo.destroy()` function's purpose is to release all of the resources
+initialized for the coroutine when it is destroyed in a suspended state.
+However, if the coroutine is only ever destroyed at the final suspend state,
+the rest of the conditions are superfluous.
+
+The user can use the `coro_only_destroy_when_complete` attributo suppress
+generation of the other destruction cases, optimizing the above `foo.destroy` to:
+
+.. code-block:: c++
+
+  void foo.destroy(foo.Frame *frame) {
+    frame->promise.~promise_type();
+    delete frame;
+  }
+
+  }];
+}
diff --git a/clang/include/clang/Basic/DiagnosticSemaKinds.td b/clang/include/clang/Basic/DiagnosticSemaKinds.td
index 4614324babb1c91..0200457b39ce5eb 100644
--- a/clang/include/clang/Basic/DiagnosticSemaKinds.td
+++ b/clang/include/clang/Basic/DiagnosticSemaKinds.td
@@ -11591,6 +11591,10 @@ def err_conflicting_aligned_options : Error <
 def err_coro_invalid_addr_of_label : Error<
   "the GNU address of label extension is not allowed in coroutines."
 >;
+def err_coroutine_return_type : Error<
+  "function returns a coroutine return type %0 but is neither a coroutine nor a coroutine wrapper; "
+  "non-coroutines should be marked with [[clang::coro_wrapper]] to allow returning coroutine return type"
+>;
 } // end of coroutines issue category
 
 let CategoryName = "Documentation Issue" in {
diff --git a/clang/include/clang/Sema/Sema.h b/clang/include/clang/Sema/Sema.h
index f69f366c1750918..4d45698e5786740 100644
--- a/clang/include/clang/Sema/Sema.h
+++ b/clang/include/clang/Sema/Sema.h
@@ -11183,6 +11183,7 @@ class Sema final {
   bool buildCoroutineParameterMoves(SourceLocation Loc);
   VarDecl *buildCoroutinePromise(SourceLocation Loc);
   void CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body);
+  void CheckCoroutineWrapper(FunctionDecl *FD);
   /// Lookup 'coroutine_traits' in std namespace and std::experimental
   /// namespace. The namespace found is recorded in Namespace.
   ClassTemplateDecl *lookupCoroutineTraits(SourceLocation KwLoc,
diff --git a/clang/lib/Sema/SemaDecl.cpp b/clang/lib/Sema/SemaDecl.cpp
index 3876eb501083acb..3e0674a2c79466d 100644
--- a/clang/lib/Sema/SemaDecl.cpp
+++ b/clang/lib/Sema/SemaDecl.cpp
@@ -17,6 +17,7 @@
 #include "clang/AST/CXXInheritance.h"
 #include "clang/AST/CharUnits.h"
 #include "clang/AST/CommentDiagnostic.h"
+#include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclTemplate.h"
@@ -15811,6 +15812,23 @@ static void diagnoseImplicitlyRetainedSelf(Sema &S) {
           << FixItHint::CreateInsertion(P.first, "self->");
 }
 
+void Sema::CheckCoroutineWrapper(FunctionDecl *FD) {
+  if (!getLangOpts().Coroutines || !FD || getCurFunction()->isCoroutine())
+    return;
+  RecordDecl *RD = FD->getReturnType()->getAsRecordDecl();
+  if (!RD || !RD->getUnderlyingDecl()->hasAttr<CoroReturnTypeAttr>())
+    return;
+  // Allow `promise_type::get_return_object`.
+  if (FD->getName() == "get_return_object") {
+    if (auto *GRT = dyn_cast<CXXMethodDecl>(FD)) {
+      if (auto *PT = GRT->getParent(); PT && PT->getName() == "promise_type")
+        return;
+    }
+  }
+  if (!FD->hasAttr<CoroWrapperAttr>())
+    Diag(FD->getLocation(), diag::err_coroutine_return_type) << RD;
+}
+
 Decl *Sema::ActOnFinishFunctionBody(Decl *dcl, Stmt *Body,
                                     bool IsInstantiation) {
   FunctionScopeInfo *FSI = getCurFunction();
@@ -15822,8 +15840,11 @@ Decl *Sema::ActOnFinishFunctionBody(Decl *dcl, Stmt *Body,
   sema::AnalysisBasedWarnings::Policy WP = AnalysisWarnings.getDefaultPolicy();
   sema::AnalysisBasedWarnings::Policy *ActivePolicy = nullptr;
 
-  if (getLangOpts().Coroutines && FSI->isCoroutine())
-    CheckCompletedCoroutineBody(FD, Body);
+  if (getLangOpts().Coroutines) {
+    if (FSI->isCoroutine())
+      CheckCompletedCoroutineBody(FD, Body);
+    CheckCoroutineWrapper(FD);
+  }
 
   {
     // Do not call PopExpressionEvaluationContext() if it is a lambda because
diff --git a/clang/test/Misc/pragma-attribute-supported-attributes-list.test b/clang/test/Misc/pragma-attribute-supported-attributes-list.test
index 969794073a5f2e8..a57bc011c059483 100644
--- a/clang/test/Misc/pragma-attribute-supported-attributes-list.test
+++ b/clang/test/Misc/pragma-attribute-supported-attributes-list.test
@@ -57,6 +57,8 @@
 // CHECK-NEXT: ConsumableSetOnRead (SubjectMatchRule_record)
 // CHECK-NEXT: Convergent (SubjectMatchRule_function)
 // CHECK-NEXT: CoroOnlyDestroyWhenComplete (SubjectMatchRule_record)
+// CHECK-NEXT: CoroReturnType (SubjectMatchRule_record)
+// CHECK-NEXT: CoroWrapper
 // CHECK-NEXT: CountedBy (SubjectMatchRule_field)
 // CHECK-NEXT: DLLExport (SubjectMatchRule_function, SubjectMatchRule_variable, SubjectMatchRule_record, SubjectMatchRule_objc_interface)
 // CHECK-NEXT: DLLImport (SubjectMatchRule_function, SubjectMatchRule_variable, SubjectMatchRule_record, SubjectMatchRule_objc_interface)
diff --git a/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp b/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp
new file mode 100644
index 000000000000000..a54bbdcbc335279
--- /dev/null
+++ b/clang/test/SemaCXX/coro-return-type-and-wrapper.cpp
@@ -0,0 +1,56 @@
+// RUN: %clang_cc1 -triple x86_64-apple-darwin9 %s -std=c++20 -fsyntax-only -verify -Wall -Wextra
+#include "Inputs/std-coroutine.h"
+
+using std::suspend_always;
+using std::suspend_never;
+
+
+template <typename T> struct [[clang::coro_return_type]] Gen {
+  struct promise_type {
+    Gen<T> get_return_object() {
+      return {};
+    }
+    suspend_always initial_suspend();
+    suspend_always final_suspend() noexcept;
+    void unhandled_exception();
+    void return_value(const T &t);
+
+    template <typename U>
+    auto await_transform(const Gen<U> &) {
+      struct awaitable {
+        bool await_ready() noexcept { return false; }
+        void await_suspend(std::coroutine_handle<>) noexcept {}
+        U await_resume() noexcept { return {}; }
+      };
+      return awaitable{};
+    }
+  };
+};
+
+Gen<int> foo_coro(int b);
+Gen<int> foo_coro(int b) { co_return b; }
+
+[[clang::coro_wrapper]] Gen<int> marked_wrapper1(int b) { return foo_coro(b); }
+
+// expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+Gen<int> non_marked_wrapper(int b) { return foo_coro(b); }
+
+namespace using_decl {
+template <typename T> using Co = Gen<T>;
+
+[[clang::coro_wrapper]] Co<int> marked_wrapper1(int b) { return foo_coro(b); }
+
+// expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+Co<int> non_marked_wrapper(int b) { return foo_coro(b); }
+} // namespace using_decl
+
+namespace lambdas {
+void foo() {
+  auto coro_lambda = []() -> Gen<int> {
+    co_return 1;
+  };
+  auto wrapper_lambda = [&]() -> Gen<int> {
+    return coro_lambda();
+  }
+}
+}
\ No newline at end of file

[hokein]

I think we should document this extension in the release note.

[hokein]

Thanks, the change looks good from my side. I will leave the final stamp to @ChuanqiXu9 and @ilya-biryukov.

[ChuanqiXu9]

LGTM. Please wait for 3~5 days in case @ilya-biryukov has more comments.

[ilya-biryukov]

Thanks, I have a few small NIT (but feel free to ignore any number of them).

I do have a major question about the use of getCurFunction()->isCoroutine(), see the corresponding comment.

[ilya-biryukov]

LGTM assuming the getCurFunction() comment is addressed in some form (moving it to callers or adding an assertion)

Thanks for the change!

[usx95]

Thank you everyone for the review.