[C++20] [Coroutines] Conform the updates for CWG issue 2585

According to the updates in CWG issue 2585
https://cplusplus.github.io/CWG/issues/2585.html, we shouldn't find an
allocation function with (size, p0, …, pn) in global scope.

Reviewed By: erichkeane

Differential Revision: https://reviews.llvm.org/D126187

clang/docs/ReleaseNotes.rst

@@ -149,8 +149,11 @@ Bug Fixes
   because there is no way to fully qualify the enumerator name, so this
   "extension" was unintentional and useless. This fixes
   `Issue 42372 <https://github.com/llvm/llvm-project/issues/42372>`_.
-- Clang shouldn't lookup allocation function in global scope for coroutines
-  in case it found the allocation function name in the promise_type body.
+- Clang will now find and emit a call to an allocation function in a 
+  promise_type body for coroutines if there is any allocation function 
+  declaration in the scope of promise_type. Additionally, to implement CWG2585,
+  a coroutine will no longer generate a call to a global allocation function
+  with the signature (std::size_t, p0, ..., pn).
   This fixes Issue `Issue 54881 <https://github.com/llvm/llvm-project/issues/54881>`_.
 
 Improvements to Clang's diagnostics

clang/lib/Sema/SemaCoroutine.cpp

@@ -1239,6 +1239,41 @@ bool CoroutineStmtBuilder::makeReturnOnAllocFailure() {
   return true;
 }
 
+// Collect placement arguments for allocation function of coroutine FD.
+// Return true if we collect placement arguments succesfully. Return false,
+// otherwise.
+static bool collectPlacementArgs(Sema &S, FunctionDecl &FD, SourceLocation Loc,
+                                 SmallVectorImpl<Expr *> &PlacementArgs) {
+  if (auto *MD = dyn_cast<CXXMethodDecl>(&FD)) {
+    if (MD->isInstance() && !isLambdaCallOperator(MD)) {
+      ExprResult ThisExpr = S.ActOnCXXThis(Loc);
+      if (ThisExpr.isInvalid())
+        return false;
+      ThisExpr = S.CreateBuiltinUnaryOp(Loc, UO_Deref, ThisExpr.get());
+      if (ThisExpr.isInvalid())
+        return false;
+      PlacementArgs.push_back(ThisExpr.get());
+    }
+  }
+
+  for (auto *PD : FD.parameters()) {
+    if (PD->getType()->isDependentType())
+      continue;
+
+    // Build a reference to the parameter.
+    auto PDLoc = PD->getLocation();
+    ExprResult PDRefExpr =
+        S.BuildDeclRefExpr(PD, PD->getOriginalType().getNonReferenceType(),
+                           ExprValueKind::VK_LValue, PDLoc);
+    if (PDRefExpr.isInvalid())
+      return false;
+
+    PlacementArgs.push_back(PDRefExpr.get());
+  }
+
+  return true;
+}
+
 bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
   // Form and check allocation and deallocation calls.
   assert(!IsPromiseDependentType &&
@@ -1255,13 +1290,7 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
   // allocated, followed by the coroutine function's arguments. If a matching
   // allocation function exists, use it. Otherwise, use an allocation function
   // that just takes the requested size.
-
-  FunctionDecl *OperatorNew = nullptr;
-  FunctionDecl *OperatorDelete = nullptr;
-  FunctionDecl *UnusedResult = nullptr;
-  bool PassAlignment = false;
-  SmallVector<Expr *, 1> PlacementArgs;
-
+  //
   // [dcl.fct.def.coroutine]p9
   //   An implementation may need to allocate additional storage for a
   //   coroutine.
@@ -1288,31 +1317,12 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
   // and p_i denotes the i-th function parameter otherwise. For a non-static
   // member function, q_1 is an lvalue that denotes *this; any other q_i is an
   // lvalue that denotes the parameter copy corresponding to p_i.
-  if (auto *MD = dyn_cast<CXXMethodDecl>(&FD)) {
-    if (MD->isInstance() && !isLambdaCallOperator(MD)) {
-      ExprResult ThisExpr = S.ActOnCXXThis(Loc);
-      if (ThisExpr.isInvalid())
-        return false;
-      ThisExpr = S.CreateBuiltinUnaryOp(Loc, UO_Deref, ThisExpr.get());
-      if (ThisExpr.isInvalid())
-        return false;
-      PlacementArgs.push_back(ThisExpr.get());
-    }
-  }
-  for (auto *PD : FD.parameters()) {
-    if (PD->getType()->isDependentType())
-      continue;
 
-    // Build a reference to the parameter.
-    auto PDLoc = PD->getLocation();
-    ExprResult PDRefExpr =
-        S.BuildDeclRefExpr(PD, PD->getOriginalType().getNonReferenceType(),
-                           ExprValueKind::VK_LValue, PDLoc);
-    if (PDRefExpr.isInvalid())
-      return false;
-
-    PlacementArgs.push_back(PDRefExpr.get());
-  }
+  FunctionDecl *OperatorNew = nullptr;
+  FunctionDecl *OperatorDelete = nullptr;
+  FunctionDecl *UnusedResult = nullptr;
+  bool PassAlignment = false;
+  SmallVector<Expr *, 1> PlacementArgs;
 
   bool PromiseContainNew = [this, &PromiseType]() -> bool {
     DeclarationName NewName =
@@ -1330,8 +1340,10 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
     //   The allocation function's name is looked up by searching for it in the
     // scope of the promise type.
     // - If any declarations are found, ...
-    // - Otherwise, a search is performed in the global scope.
-    Sema::AllocationFunctionScope NewScope = PromiseContainNew ? Sema::AFS_Class : Sema::AFS_Global;
+    // - If no declarations are found in the scope of the promise type, a search
+    // is performed in the global scope.
+    Sema::AllocationFunctionScope NewScope =
+        PromiseContainNew ? Sema::AFS_Class : Sema::AFS_Global;
     S.FindAllocationFunctions(Loc, SourceRange(),
                               NewScope,
                               /*DeleteScope*/ Sema::AFS_Both, PromiseType,
@@ -1339,13 +1351,19 @@ bool CoroutineStmtBuilder::makeNewAndDeleteExpr() {
                               OperatorNew, UnusedResult, /*Diagnose*/ false);
   };
 
+  // We don't expect to call to global operator new with (size, p0, …, pn).
+  // So if we choose to lookup the allocation function in global scope, we
+  // shouldn't lookup placement arguments.
+  if (PromiseContainNew && !collectPlacementArgs(S, FD, Loc, PlacementArgs))
+    return false;
+
   LookupAllocationFunction();
 
   // [dcl.fct.def.coroutine]p9
   //   If no viable function is found ([over.match.viable]), overload resolution
   // is performed again on a function call created by passing just the amount of
   // space required as an argument of type std::size_t.
-  if (!OperatorNew && !PlacementArgs.empty()) {
+  if (!OperatorNew && !PlacementArgs.empty() && PromiseContainNew) {
     PlacementArgs.clear();
     LookupAllocationFunction();
   }

clang/test/CodeGenCoroutines/coro-alloc-2.cpp

@@ -0,0 +1,30 @@
+// Tests that we wouldn't generate an allocation call in global scope with (std::size_t, p0, ..., pn)
+// RUN: %clang_cc1 %s -std=c++20 -S -triple x86_64-unknown-linux-gnu -emit-llvm -disable-llvm-passes %s -o - | FileCheck %s
+#include "Inputs/coroutine.h"
+
+namespace std {
+typedef decltype(sizeof(int)) size_t;
+}
+
+struct Allocator {};
+
+struct resumable {
+  struct promise_type {
+
+    resumable get_return_object() { return {}; }
+    auto initial_suspend() { return std::suspend_always(); }
+    auto final_suspend() noexcept { return std::suspend_always(); }
+    void unhandled_exception() {}
+    void return_void(){};
+  };
+};
+
+void *operator new(std::size_t, void *);
+
+resumable f1(void *) {
+  co_return;
+}
+
+// CHECK: coro.alloc:
+// CHECK-NEXT: [[SIZE:%.+]] = call [[BITWIDTH:.+]] @llvm.coro.size.[[BITWIDTH]]()
+// CHECK-NEXT: call {{.*}} ptr @_Znwm([[BITWIDTH]] noundef [[SIZE]])