[Coroutines] Use LookupAllocationFunction to find allocation functions for coroutines consistently

Previously we may call Sema::FindAllocationFunctions directly to lookup
allocation functions directly instead of using our wrapped lambda
LookupAllocationFunction, which is slightly incosnsistent. It will be
helpful to refactor this for further changes.

Also previously, when we lookup 'operator new(std::size_t, std::nothrow_t)' in
case we found `get_­return_­object_­on_­allocation_­failure` in the
promise_type, the compiler will try to look at the allocation function
in promise_type. However, this is not wanted actually. According to
[dcl.fct.def.coroutine]p10:

> if a global allocation function is selected, the
> ::operator new(size_­t, nothrow_­t) form is used.

So we should only lookup for `::operator (size_t, nothrow_t)` for the
global allocation function. For the allocation function in the
promise_type, the requirement is that it shouldn't throw, which has
already been checked.

Given users generally include headers from standard libs so it will
generally include the <new> header, so this change should be a trivial
one and shouldn't affect almost any user.

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -11272,6 +11272,9 @@ def warn_always_inline_coroutine : Warning<
 def err_coroutine_unusable_new : Error<
   "'operator new' provided by %0 is not usable with the function signature of %1"
 >;
+def err_coroutine_unfound_nothrow_new : Error <
+  "unable to find '::operator new(size_­t, nothrow_­t)' for %0"
+>;
 } // end of coroutines issue category
 
 let CategoryName = "Documentation Issue" in {

clang/lib/Sema/SemaCoroutine.cpp

@@ -1326,7 +1326,7 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
   bool PassAlignment = false;
   SmallVector<Expr *, 1> PlacementArgs;
 
-  bool PromiseContainsNew = [this, &PromiseType]() -> bool {
+  const bool PromiseContainsNew = [this, &PromiseType]() -> bool {
     DeclarationName NewName =
         S.getASTContext().DeclarationNames.getCXXOperatorName(OO_New);
     LookupResult R(S, NewName, Loc, Sema::LookupOrdinaryName);
@@ -1337,22 +1337,22 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
     return !R.empty() && !R.isAmbiguous();
   }();
 
-  auto LookupAllocationFunction = [&]() {
+  auto LookupAllocationFunction = [&](Sema::AllocationFunctionScope NewScope =
+                                          Sema::AFS_Both) {
     // [dcl.fct.def.coroutine]p9
     //   The allocation function's name is looked up by searching for it in the
     // scope of the promise type.
     // - If any declarations are found, ...
     // - If no declarations are found in the scope of the promise type, a search
     // is performed in the global scope.
-    Sema::AllocationFunctionScope NewScope =
-        PromiseContainsNew ? Sema::AFS_Class : Sema::AFS_Global;
+    if (NewScope == Sema::AFS_Both)
+      NewScope = PromiseContainsNew ? Sema::AFS_Class : Sema::AFS_Global;
+
     FunctionDecl *UnusedResult = nullptr;
-    S.FindAllocationFunctions(Loc, SourceRange(),
-                              NewScope,
+    S.FindAllocationFunctions(Loc, SourceRange(), NewScope,
                               /*DeleteScope*/ Sema::AFS_Both, PromiseType,
                               /*isArray*/ false, PassAlignment, PlacementArgs,
-                              OperatorNew, UnusedResult,
-                              /*Diagnose*/ false);
+                              OperatorNew, UnusedResult, /*Diagnose*/ false);
   };
 
   // We don't expect to call to global operator new with (size, p0, …, pn).
@@ -1383,16 +1383,14 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
       return false;
     PlacementArgs = {StdNoThrow};
     OperatorNew = nullptr;
-    FunctionDecl *UnusedResult = nullptr;
-    S.FindAllocationFunctions(Loc, SourceRange(), /*NewScope*/ Sema::AFS_Both,
-                              /*DeleteScope*/ Sema::AFS_Both, PromiseType,
-                              /*isArray*/ false, PassAlignment, PlacementArgs,
-                              OperatorNew, UnusedResult);
+    LookupAllocationFunction(Sema::AFS_Global);
   }
 
   if (!OperatorNew) {
     if (PromiseContainsNew)
       S.Diag(Loc, diag::err_coroutine_unusable_new) << PromiseType << &FD;
+    else if (RequiresNoThrowAlloc)
+      S.Diag(Loc, diag::err_coroutine_unfound_nothrow_new) << &FD;
 
     return false;
   }

clang/test/SemaCXX/coroutine-alloc-2.cpp

@@ -0,0 +1,52 @@
+// Tests that we wouldn't generate an allocation call in promise_type with (std::size_t, std::nothrow_t) in case we find promsie_type::get_return_object_on_allocation_failure;
+// RUN: %clang_cc1 %s -std=c++20 %s -fsyntax-only -verify
+
+namespace std {
+template <typename... T>
+struct coroutine_traits;
+
+template <class Promise = void>
+struct coroutine_handle {
+  coroutine_handle() = default;
+  static coroutine_handle from_address(void *) noexcept { return {}; }
+};
+
+template <>
+struct coroutine_handle<void> {
+  static coroutine_handle from_address(void *) { return {}; }
+  coroutine_handle() = default;
+  template <class PromiseType>
+  coroutine_handle(coroutine_handle<PromiseType>) noexcept {}
+};
+
+struct suspend_always {
+  bool await_ready() noexcept { return false; }
+  void await_suspend(std::coroutine_handle<>) noexcept {}
+  void await_resume() noexcept {}
+};
+
+struct nothrow_t {};
+constexpr nothrow_t nothrow = {};
+
+} // end namespace std
+
+using SizeT = decltype(sizeof(int));
+
+struct promise_on_alloc_failure_tag {};
+
+template <>
+struct std::coroutine_traits<int, promise_on_alloc_failure_tag> {
+  struct promise_type {
+    int get_return_object() { return 0; }
+    std::suspend_always initial_suspend() { return {}; }
+    std::suspend_always final_suspend() noexcept { return {}; }
+    void return_void() {}
+    void unhandled_exception() {}
+    static int get_return_object_on_allocation_failure() { return -1; }
+    void *operator new(SizeT, std::nothrow_t) noexcept;
+  };
+};
+
+extern "C" int f(promise_on_alloc_failure_tag) { // expected-error 1+{{is not usable with the function signature of 'f'}}
+    co_return;
+}

clang/test/SemaCXX/coroutine-alloc-3.cpp

@@ -0,0 +1,51 @@
+// Tests that we'll emit the proper diagnostic message if we failed to find `::operator new(size_­t, nothrow_­t)`.
+// RUN: %clang_cc1 %s -std=c++20 %s -fsyntax-only -verify
+
+namespace std {
+template <typename... T>
+struct coroutine_traits;
+
+template <class Promise = void>
+struct coroutine_handle {
+  coroutine_handle() = default;
+  static coroutine_handle from_address(void *) noexcept { return {}; }
+};
+
+template <>
+struct coroutine_handle<void> {
+  static coroutine_handle from_address(void *) { return {}; }
+  coroutine_handle() = default;
+  template <class PromiseType>
+  coroutine_handle(coroutine_handle<PromiseType>) noexcept {}
+};
+
+struct suspend_always {
+  bool await_ready() noexcept { return false; }
+  void await_suspend(std::coroutine_handle<>) noexcept {}
+  void await_resume() noexcept {}
+};
+
+struct nothrow_t {};
+constexpr nothrow_t nothrow = {};
+
+} // end namespace std
+
+using SizeT = decltype(sizeof(int));
+
+struct promise_on_alloc_failure_tag {};
+
+template <>
+struct std::coroutine_traits<int, promise_on_alloc_failure_tag> {
+  struct promise_type {
+    int get_return_object() { return 0; }
+    std::suspend_always initial_suspend() { return {}; }
+    std::suspend_always final_suspend() noexcept { return {}; }
+    void return_void() {}
+    void unhandled_exception() {}
+    static int get_return_object_on_allocation_failure() { return -1; }
+  };
+};
+
+extern "C" int f(promise_on_alloc_failure_tag) { // expected-error 1+{{unable to find '::operator new(size_­t, nothrow_­t)' for 'f'}}
+    co_return;
+}