[coroutines] Handle get_return_object_on_allocation_failure

[GorNishanov]

Summary

If promise_type has get_return_object_on_allocation_failure defined,
check if an allocation function returns nullptr, and if so,
return the result of get_return_object_on_allocation_failure().

[EricWF]

This LGTM minus inline comments.

lib/Sema/SemaCoroutine.cpp

793

Maybe add a test for this in SemaCXX/coroutines.cpp?

test/CodeGenCoroutines/coro-alloc.cpp

131

FYI this test doesn't merge with trunk seemingly due to whitespace issues.

[GorNishanov]

Added diagnostic test to SemaCXX/coroutines.cpp
Preparing to Land

[alekseyshl]

Seems like this bot is not happy about this change: http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/3771

[EricWF]

Woops. this patch also needed changes to TreeTransform.h. I'll add those tonight.

include/clang/AST/StmtCXX.h

enum SubStmt {
     Promise,       ///< The promise statement.
     InitSuspend,   ///< The initial suspend statement, run before the body.
     FinalSuspend,  ///< The final suspend statement, run after the body.
     OnException,   ///< Handler for exceptions thrown in the body.
     OnFallthrough, ///< Handler for control flow falling off the body.
     Allocate,      ///< Coroutine frame memory allocation.
     Deallocate,    ///< Coroutine frame memory deallocation.
     ReturnValue,   ///< Return value for thunk function.
+    ReturnStmtOnAllocFailure, ///< Return statement if allocation failed.
     FirstParamMove ///< First offset for move construction of parameter copies.
   };
   unsigned NumParams;
 
   friend class ASTStmtReader;
   friend TrailingObjects;
 
   Stmt **getStoredStmts() { return getTrailingObjects<Stmt *>(); }
 
   Stmt *const *getStoredStmts() const { return getTrailingObjects<Stmt *>(); }
 
 public:
 
   struct CtorArgs {
     Stmt *Body = nullptr;
     Stmt *Promise = nullptr;
     Expr *InitialSuspend = nullptr;
     Expr *FinalSuspend = nullptr;
     Stmt *OnException = nullptr;
     Stmt *OnFallthrough = nullptr;
     Expr *Allocate = nullptr;
     Expr *Deallocate = nullptr;
     Stmt *ReturnValue = nullptr;
+    Stmt *ReturnStmtOnAllocFailure = nullptr;
     ArrayRef<Stmt *> ParamMoves;
   };
 
 private:
 
   CoroutineBodyStmt(CtorArgs const& Args);
 
 public:
public:
   }
   Expr *getDeallocate() const {
     return cast_or_null<Expr>(getStoredStmts()[SubStmt::Deallocate]);
   }
 
   Expr *getReturnValueInit() const {
     return cast_or_null<Expr>(getStoredStmts()[SubStmt::ReturnValue]);
   }
+  Stmt *getReturnStmtOnAllocFailure() const {
+    return getStoredStmts()[SubStmt::ReturnStmtOnAllocFailure];
+  }
   ArrayRef<Stmt const *> getParamMoves() const {
     return {getStoredStmts() + SubStmt::FirstParamMove, NumParams};
   }
 
   SourceLocation getLocStart() const LLVM_READONLY {
     return getBody() ? getBody()->getLocStart()
             : getPromiseDecl()->getLocStart();
   }

lib/AST/StmtCXX.cpp

CoroutineBodyStmt::CoroutineBodyStmt(CoroutineBodyStmt::CtorArgs const &Args)
   SubStmts[CoroutineBodyStmt::Promise] = Args.Promise;
   SubStmts[CoroutineBodyStmt::InitSuspend] = Args.InitialSuspend;
   SubStmts[CoroutineBodyStmt::FinalSuspend] = Args.FinalSuspend;
   SubStmts[CoroutineBodyStmt::OnException] = Args.OnException;
   SubStmts[CoroutineBodyStmt::OnFallthrough] = Args.OnFallthrough;
   SubStmts[CoroutineBodyStmt::Allocate] = Args.Allocate;
   SubStmts[CoroutineBodyStmt::Deallocate] = Args.Deallocate;
   SubStmts[CoroutineBodyStmt::ReturnValue] = Args.ReturnValue;
+  SubStmts[CoroutineBodyStmt::ReturnStmtOnAllocFailure] =
+      Args.ReturnStmtOnAllocFailure;
   std::copy(Args.ParamMoves.begin(), Args.ParamMoves.end(),
             const_cast<Stmt **>(getParamMoves().data()));
 }
 No newline at end of file

lib/CodeGen/CGCoroutine.cpp

void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   auto *RetBB = createBasicBlock("coro.ret");
 
   auto *CoroId = Builder.CreateCall(
       CGM.getIntrinsic(llvm::Intrinsic::coro_id),
       {Builder.getInt32(NewAlign), NullPtr, NullPtr, NullPtr});
   createCoroData(*this, CurCoro, CoroId);
   CurCoro.Data->SuspendBB = RetBB;
 
-  EmitScalarExpr(S.getAllocate());
+  auto *AllocateCall = EmitScalarExpr(S.getAllocate());
+
+  // Handle allocation failure if 'ReturnStmtOnAllocFailure' was provided.
+  if (auto *RetOnAllocFailure = S.getReturnStmtOnAllocFailure()) {
+    auto *RetOnFailureBB = createBasicBlock("coro.ret.on.failure");
+    auto *InitBB = createBasicBlock("coro.init");
+
+    // See if allocation was successful.
+    auto *NullPtr = llvm::ConstantPointerNull::get(Int8PtrTy);
+    auto *Cond = Builder.CreateICmpNE(AllocateCall, NullPtr);
+    Builder.CreateCondBr(Cond, InitBB, RetOnFailureBB);
+
+    // If not, return OnAllocFailure object.
+    EmitBlock(RetOnFailureBB);
+    EmitStmt(RetOnAllocFailure);
+
+    EmitBlock(InitBB);
+  }
 
   // FIXME: Setup cleanup scopes.
 
   EmitStmt(S.getPromiseDeclStmt());
 
   CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
   CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 

lib/Sema/SemaCoroutine.cpp

SubStmtBuilder(Sema &S, FunctionDecl &FD, FunctionScopeInfo &Fn, Stmt *Body)
             Fn.CoroutinePromise->getType()->isDependentType()) {
     this->Body = Body;
     if (!IsPromiseDependentType) {
       PromiseRecordDecl = Fn.CoroutinePromise->getType()->getAsCXXRecordDecl();
       assert(PromiseRecordDecl && "Type should have already been checked");
     }
     this->IsValid = makePromiseStmt() && makeInitialAndFinalSuspend() &&
                     makeOnException() && makeOnFallthrough() &&
-                    makeNewAndDeleteExpr() && makeReturnObject() &&
-                    makeParamMoves();
+                    makeReturnOnAllocFailure() && makeNewAndDeleteExpr() &&
+                    makeReturnObject() && makeParamMoves();
   }
 
   bool isInvalid() const { return !this->IsValid; }
 
   bool makePromiseStmt();
   bool makeInitialAndFinalSuspend();
   bool makeNewAndDeleteExpr();
   bool makeOnFallthrough();
   bool makeOnException();
   bool makeReturnObject();
+  bool makeReturnOnAllocFailure();
   bool makeParamMoves();
 };
 }
 
 void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
   FunctionScopeInfo *Fn = getCurFunction();
   assert(Fn && Fn->CoroutinePromise && "not a coroutine");
 

 bool SubStmtBuilder::makeInitialAndFinalSuspend() {
   if (Fn.hasInvalidCoroutineSuspends())
     return false;
   this->InitialSuspend = cast<Expr>(Fn.CoroutineSuspends.first);
   this->FinalSuspend = cast<Expr>(Fn.CoroutineSuspends.second);
   return true;
 }
 
+bool SubStmtBuilder::makeReturnOnAllocFailure() {
+  if (!PromiseRecordDecl) return true;
+
+  // [dcl.fct.def.coroutine]/8
+  // The unqualified-id get_return_object_on_allocation_failure is looked up in
+  // the scope of class P by class member access lookup (3.4.5). ...
+  // If an allocation function returns nullptr, ... the coroutine return value
+  // is obtained by a call to ... get_return_object_on_allocation_failure().
+
+  DeclarationName DN =
+      S.PP.getIdentifierInfo("get_return_object_on_allocation_failure");
+  LookupResult Found(S, DN, Loc, Sema::LookupMemberName);
+  // Suppress diagnostics when a private member is selected. The same warnings
+  // will be produced again when building the call.
+  Found.suppressDiagnostics();
+  if (!S.LookupQualifiedName(Found, PromiseRecordDecl)) return true;
+
+  CXXScopeSpec SS;
+  ExprResult DeclNameExpr =
+      S.BuildDeclarationNameExpr(SS, Found, /*NeedsADL=*/false);
+  if (DeclNameExpr.isInvalid()) return false;
+
+  ExprResult ReturnObjectOnAllocationFailure =
+      S.ActOnCallExpr(nullptr, DeclNameExpr.get(), Loc, {}, Loc);
+  if (ReturnObjectOnAllocationFailure.isInvalid()) return false;
+
+  StmtResult ReturnStmt = S.ActOnReturnStmt(
+      Loc, ReturnObjectOnAllocationFailure.get(), S.getCurScope());
+  if (ReturnStmt.isInvalid()) return false;
+
+  this->ReturnStmtOnAllocFailure = ReturnStmt.get();
+  return true;
+}
+
 bool SubStmtBuilder::makeNewAndDeleteExpr() {
   // Form and check allocation and deallocation calls.
   QualType PromiseType = Fn.CoroutinePromise->getType();
   if (PromiseType->isDependentType())
     return true;
 
   if (S.RequireCompleteType(Loc, PromiseType, diag::err_incomplete_type))
     return false;
 
-  // FIXME: Add support for get_return_object_on_allocation failure.
+  // FIXME: Add nothrow_t placement arg for global alloc
+  //        if ReturnStmtOnAllocFailure != nullptr.
   // FIXME: Add support for stateful allocators.
 
   FunctionDecl *OperatorNew = nullptr;
   FunctionDecl *OperatorDelete = nullptr;
   FunctionDecl *UnusedResult = nullptr;
   bool PassAlignment = false;
 
   S.FindAllocationFunctions(Loc, SourceRange(),

test/CodeGenCoroutines/coro-alloc.cpp

struct std::experimental::coroutine_traits<void, global_new_delete_tag> {
   struct promise_type {
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }
     suspend_always final_suspend() { return {}; }
     void return_void() {}
   };
 };
 
-// CHECK-LABEL: f0( 
+// CHECK-LABEL: f0(
 extern "C" void f0(global_new_delete_tag) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
   co_return;
 }
 
 struct promise_new_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_new_tag> {
   struct promise_type {
     void *operator new(unsigned long);
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }
     suspend_always final_suspend() { return {}; }
     void return_void() {}
   };
 };
 
-// CHECK-LABEL: f1( 
+// CHECK-LABEL: f1(
 extern "C" void f1(promise_new_tag ) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_ZNSt12experimental16coroutine_traitsIJv15promise_new_tagEE12promise_typenwEm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
   co_return;
 }
 
 struct promise_delete_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_delete_tag> {
   struct promise_type {
     void operator delete(void*);
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }
     suspend_always final_suspend() { return {}; }
     void return_void() {}
   };
 };
 
-// CHECK-LABEL: f2( 
+// CHECK-LABEL: f2(
 extern "C" void f2(promise_delete_tag) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJv18promise_delete_tagEE12promise_typedlEPv(i8* %[[MEM]])
extern "C" void f3(promise_sized_delete_tag) {
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: %[[SIZE2:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJv24promise_sized_delete_tagEE12promise_typedlEPvm(i8* %[[MEM]], i64 %[[SIZE2]])
   co_return;
 }
+
+struct promise_on_alloc_failure_tag {};
+
+template<>
+struct std::experimental::coroutine_traits<int, promise_on_alloc_failure_tag> {
+  struct promise_type {
+    int get_return_object() {}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+    static int get_return_object_on_allocation_failure() { return -1; }
+  };
+};
+
+// CHECK-LABEL: f4(
+extern "C" int f4(promise_on_alloc_failure_tag) {
+  // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
+  // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
+  // CHECK: %[[MEM:.+]] = call i8* @_Znwm(i64 %[[SIZE]])
+  // CHECK: %[[OK:.+]] = icmp ne i8* %[[MEM]], null
+  // CHECK: br i1 %[[OK]], label %[[OKBB:.+]], label %[[ERRBB:.+]]
+
+  // CHECK: [[ERRBB]]:
+  // CHECK: %[[RETVAL:.+]] = call i32 @_ZNSt12experimental16coroutine_traitsIJi28promise_on_alloc_failure_tagEE12promise_type39get_return_object_on_allocation_failureEv(
+  // CHECK: ret i32 %[[RETVAL]]
+  co_return;
+}