[coroutines] Introduce [[clang::coro_return_type]] and [[clang::coro_wrapper]] (#71945)

First step in the implementation of
[RFC](https://discourse.llvm.org/t/rfc-lifetime-bound-check-for-parameters-of-coroutines/74253)
([final approved
doc](https://docs.google.com/document/d/1hkfXHuvIW1Yv5LI-EIkpWzdWgIoUlzO6Zv_KJpknQzM/edit)).

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

---------

Co-authored-by: Chuanqi Xu <[email protected]>

Author: usx95

clang/docs/ReleaseNotes.rst

@@ -306,6 +306,11 @@ Attribute Changes in Clang
   to reduce the size of the destroy functions for coroutines which are known to
   be destroyed after having reached the final suspend point.
 
+- Clang now introduced ``[[clang::coro_return_type]]`` and ``[[clang::coro_wrapper]]``
+  attributes. A function returning a type marked with ``[[clang::coro_return_type]]``
+  should be a coroutine. A non-coroutine function marked with ``[[clang::coro_wrapper]]``
+  is still allowed to return the such a type. This is helpful for analyzers to recognize coroutines from the function signatures.
+
 Improvements to Clang's diagnostics
 -----------------------------------
 - Clang constexpr evaluator now prints template arguments when displaying

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">];

clang/include/clang/Basic/AttrDocs.td

@@ -7482,3 +7482,61 @@ generation of the other destruction cases, optimizing the above `foo.destroy` to
 
   }];
 }
+
+
+def CoroReturnTypeAndWrapperDoc : Documentation {
+  let Category = DocCatDecl;
+  let Content = [{
+The ``[[clang::coro_return_type]]`` attribute is used to help static analyzers to recognize
+coroutines from the function signatures.
+
+The ``coro_return_type`` attribute should be marked on a C++ class to mark it as
+a **coroutine return type (CRT)**.
+
+A function ``R func(P1, .., PN)`` has a coroutine return type (CRT) ``R`` if ``R``
+is marked by ``[[clang::coro_return_type]]`` and  ``R`` has a promise type associated to it
+(i.e., std​::​coroutine_traits<R, P1, .., PN>​::​promise_type is a valid promise type).
+
+If the return type of a function is a ``CRT`` then the function must be a coroutine.
+Otherwise the program is invalid. It is allowed for a non-coroutine to return a ``CRT``
+if the function is marked with ``[[clang::coro_wrapper]]``.
+
+The ``[[clang::coro_wrapper]]`` attribute should be marked on a C++ function to mark it as
+a **coroutine wrapper**. A coroutine wrapper is a function which returns a ``CRT``,
+is not a coroutine itself and is marked with ``[[clang::coro_wrapper]]``.
+
+Clang will enforce that all functions that return a ``CRT`` are either coroutines or marked
+with ``[[clang::coro_wrapper]]``. Clang will enforce this with an error.
+
+From a language perspective, it is not possible to differentiate between a coroutine and a
+function returning a CRT by merely looking at the function signature.
+
+Coroutine wrappers, in particular, are susceptible to capturing
+references to temporaries and other lifetime issues. This allows to avoid such lifetime
+issues with coroutine wrappers.
+
+For example,
+
+.. code-block:: c++
+
+  // This is a CRT.
+  template <typename T> struct [[clang::coro_return_type]] Task {
+    using promise_type = some_promise_type;
+  };
+
+  Task<int> increment(int a) { co_return a + 1; } // Fine. This is a coroutine.
+  Task<int> foo() { return increment(1); } // Error. foo is not a coroutine.
+
+  // Fine for a coroutine wrapper to return a CRT.
+  Task<int> [[clang::coro_wrapper]] foo() { return increment(1); }
+
+  void bar() {
+    // Invalid. This intantiates a function which returns a CRT but is not marked as
+    // a coroutine wrapper.
+    std::function<Task<int>(int)> f = increment;
+  }
+
+Note: ``a_promise_type::get_return_object`` is exempted from this analysis as it is a necessary
+implementation detail of any coroutine library.
+}];
+}

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 type %0 makred with [[clang::coro_return_type]] 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 {

clang/include/clang/Sema/Sema.h

@@ -11189,6 +11189,12 @@ class Sema final {
   bool buildCoroutineParameterMoves(SourceLocation Loc);
   VarDecl *buildCoroutinePromise(SourceLocation Loc);
   void CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body);
+
+  // As a clang extension, enforces that a non-coroutine function must be marked
+  // with [[clang::coro_wrapper]] if it returns a type marked with
+  // [[clang::coro_return_type]].
+  // Expects that FD is not a coroutine.
+  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,

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"
@@ -26,6 +27,7 @@
 #include "clang/AST/NonTrivialTypeVisitor.h"
 #include "clang/AST/Randstruct.h"
 #include "clang/AST/StmtCXX.h"
+#include "clang/AST/Type.h"
 #include "clang/Basic/Builtins.h"
 #include "clang/Basic/HLSLRuntime.h"
 #include "clang/Basic/PartialDiagnostic.h"
@@ -15814,6 +15816,20 @@ static void diagnoseImplicitlyRetainedSelf(Sema &S) {
           << FixItHint::CreateInsertion(P.first, "self->");
 }
 
+void Sema::CheckCoroutineWrapper(FunctionDecl *FD) {
+  if (!FD)
+    return;
+  RecordDecl *RD = FD->getReturnType()->getAsRecordDecl();
+  if (!RD || !RD->getUnderlyingDecl()->hasAttr<CoroReturnTypeAttr>())
+    return;
+  // Allow `get_return_object()`.
+  if (FD->getDeclName().isIdentifier() &&
+      FD->getName().equals("get_return_object") && FD->param_empty())
+    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();
@@ -15825,8 +15841,12 @@ 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);
+    else
+      CheckCoroutineWrapper(FD);
+  }
 
   {
     // Do not call PopExpressionEvaluationContext() if it is a lambda because

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)

clang/test/SemaCXX/coro-return-type-and-wrapper.cpp

@@ -0,0 +1,134 @@
+// 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(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 {
+#define CORO_WRAPPER \
+  _Pragma("clang diagnostic push") \
+  _Pragma("clang diagnostic ignored \"-Wc++23-extensions\"") \
+  [[clang::coro_wrapper]] \
+  _Pragma("clang diagnostic pop")
+
+void foo() {
+  auto coro_lambda = []() -> Gen<int> {
+    co_return 1;
+  };
+  // expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+  auto not_allowed_wrapper = []() -> Gen<int> {
+    return foo_coro(1);
+  };
+  auto allowed_wrapper = [] CORO_WRAPPER() -> Gen<int> {
+    return foo_coro(1);
+  };
+}
+
+Gen<int> coro_containing_lambda() {
+  // expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+  auto wrapper_lambda = []() -> Gen<int> {
+    return foo_coro(1);
+  };
+  co_return co_await wrapper_lambda();
+}
+} // namespace lambdas
+
+namespace std_function {
+namespace std {
+template <typename> class function;
+
+template <typename ReturnValue, typename... Args>
+class function<ReturnValue(Args...)> {
+public:
+  template <typename T> function &operator=(T) {}
+  template <typename T> function(T) {}
+  // expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+  ReturnValue operator()(Args... args) const {
+    return callable_->Invoke(args...);  // expected-note {{in instantiation of member}}
+  }
+
+private:
+  class Callable {
+  public:
+    // expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+    ReturnValue Invoke(Args...) const { return {}; }
+  };
+  Callable* callable_;
+};
+} // namespace std
+
+void use_std_function() {
+  std::function<int(bool)> foo = [](bool b) { return b ? 1 : 2; };
+  // expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+  std::function<Gen<int>(bool)> test1 = [](bool b) {
+    return foo_coro(b);
+  };
+  std::function<Gen<int>(bool)> test2 = [](bool) -> Gen<int> {
+    co_return 1;
+  };
+  std::function<Gen<int>(bool)> test3 = foo_coro;
+
+  foo(true);   // Fine.
+  test1(true); // expected-note 2 {{in instantiation of member}}
+  test2(true);
+  test3(true);
+}
+} // namespace std_function
+
+// different_promise_type
+class [[clang::coro_return_type]] Task{};
+struct my_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<Task, int> {
+    using promise_type = my_promise_type;
+};
+} // namespace std
+// expected-error@+1 {{neither a coroutine nor a coroutine wrapper}}
+Task foo(int) { return Task{}; }