[coroutines] Add support for deallocation elision

[GorNishanov]

Summary

Wrap deallocation code with:

if (auto *mem = coro.free()) Deallocate

When backend decides to elide allocations it will replace coro.free with nullptr to suppress deallocation code.

[GorNishanov]

Gentle ping

[GorNishanov]

LGTM

[GorNishanov]

Commited: r303599 = abbd4f85db96b0ed1478e2ee52e9f1eec7404131

lib/CodeGen/CGCoroutine.cpp

struct clang::CodeGen::CGCoroData {
   // builtin.
   llvm::CallInst *CoroId = nullptr;
 
   // Stores the llvm.coro.begin emitted in the function so that we can replace
   // all coro.frame intrinsics with direct SSA value of coro.begin that returns
   // the address of the coroutine frame of the current coroutine.
   llvm::CallInst *CoroBegin = nullptr;
 
+  // Stores the last emitted coro.free for the deallocate expressions, we use it
+  // to wrap dealloc code with if(auto mem = coro.free) dealloc(mem).
+  llvm::CallInst *LastCoroFree = nullptr;
+
   // If coro.id came from the builtin, remember the expression to give better
   // diagnostic. If CoroIdExpr is nullptr, the coro.id was created by
   // EmitCoroutineBody.
   CallExpr const *CoroIdExpr = nullptr;
 };
 
 // Defining these here allows to keep CGCoroData private to this file.
 clang::CodeGen::CodeGenFunction::CGCoroInfo::CGCoroInfo() {}
void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
   EmitBranchThroughCleanup(CurCoro.Data->FinalJD);
 }
 
 namespace {
 // Make sure to call coro.delete on scope exit.
 struct CallCoroDelete final : public EHScopeStack::Cleanup {
   Stmt *Deallocate;
 
-  // TODO: Wrap deallocate in if(coro.free(...)) Deallocate.
-  void Emit(CodeGenFunction &CGF, Flags) override {
-    // Note: That deallocation will be emitted twice: once for a normal exit and
-    // once for exceptional exit. This usage is safe because Deallocate does not
-    // contain any declarations. The SubStmtBuilder::makeNewAndDeleteExpr()
-    // builds a single call to a deallocation function which is safe to emit
-    // multiple times.
+  // Emit "if (coro.free(CoroId, CoroBegin)) Deallocate;"
+
+  void Emit(CodeGenFunction &CGF, Flags) override {
+    // Remember the current point, as we are going to emit deallocation code
+    // first to get to coro.free instruction that is an argument to a delete
+    // call.
+    BasicBlock *SaveInsertBlock = CGF.Builder.GetInsertBlock();
+
+    auto *FreeBB = CGF.createBasicBlock("coro.free");
+    CGF.EmitBlock(FreeBB);
     CGF.EmitStmt(Deallocate);
+
+    auto *AfterFreeBB = CGF.createBasicBlock("after.coro.free");
+    CGF.EmitBlock(AfterFreeBB);
+
+    // We should have captured coro.free from the emission of deallocate.
+    auto *CoroFree = CGF.CurCoro.Data->LastCoroFree;
+    if (!CoroFree) {
+      CGF.CGM.Error(Deallocate->getLocStart(),
+                    "Deallocation expressoin does not refer to coro.free");
+      return;
+    }
+
+    // Get back to the block we were originally and move coro.free there.
+    auto *InsertPt = SaveInsertBlock->getTerminator();
+    CoroFree->moveBefore(InsertPt);
+    CGF.Builder.SetInsertPoint(InsertPt);
+
+    // Add if (auto *mem = coro.free) Deallocate;
+    auto *NullPtr = llvm::ConstantPointerNull::get(CGF.Int8PtrTy);
+    auto *Cond = CGF.Builder.CreateICmpNE(CoroFree, NullPtr);
+    CGF.Builder.CreateCondBr(Cond, FreeBB, AfterFreeBB);
+
+    // No longer need old terminator.
+    InsertPt->eraseFromParent();
+    CGF.Builder.SetInsertPoint(AfterFreeBB);
   }
   explicit CallCoroDelete(Stmt *DeallocStmt) : Deallocate(DeallocStmt) {}
 };
 }
 
 void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   auto *NullPtr = llvm::ConstantPointerNull::get(Builder.getInt8PtrTy());
   auto &TI = CGM.getContext().getTargetInfo();
case llvm::Intrinsic::coro_suspend:
     break;
   }
   for (auto &Arg : E->arguments())
     Args.push_back(EmitScalarExpr(Arg));
 
   llvm::Value *F = CGM.getIntrinsic(IID);
   llvm::CallInst *Call = Builder.CreateCall(F, Args);
 
-  // Note: The following code is to enable to emit coroutine intrinsics by
+  // Note: The following code is to enable to emit coro.id and coro.begin by
   // hand to experiment with coroutines in C.
   // If we see @llvm.coro.id remember it in the CoroData. We will update
   // coro.alloc, coro.begin and coro.free intrinsics to refer to it.
   if (IID == llvm::Intrinsic::coro_id) {
     createCoroData(*this, CurCoro, Call, E);
   }
   else if (IID == llvm::Intrinsic::coro_begin) {
     if (CurCoro.Data)
       CurCoro.Data->CoroBegin = Call;
   }
-  return RValue::get(Call);
+  else if (IID == llvm::Intrinsic::coro_free) {
+    // Remember the last coro_free as we need it to build the conditional
+    // deletion of the coroutine frame.
+    if (CurCoro.Data)
+      CurCoro.Data->LastCoroFree = Call;
+  }  return RValue::get(Call);
 }

test/CodeGenCoroutines/coro-alloc.cpp

extern "C" void f0(global_new_delete_tag) {
   // CHECK: %[[MEM:.+]] = call i8* @_Znwm(i64 %[[SIZE]])
   // CHECK: br label %[[InitBB]]
 
   // CHECK: [[InitBB]]:
   // CHECK: %[[PHI:.+]] = phi i8* [ null, %{{.+}} ], [ %call, %[[AllocBB]] ]
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(token %[[ID]], i8* %[[PHI]])
 
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
+  // CHECK: %[[NeedDealloc:.+]] = icmp ne i8* %[[MEM]], null
+  // CHECK: br i1 %[[NeedDealloc]], label %[[FreeBB:.+]], label %[[Afterwards:.+]]
+
+  // CHECK: [[FreeBB]]:
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
+  // CHECK: br label %[[Afterwards]]
+
+  // CHECK: [[Afterwards]]:
+  // CHECK: ret void
   co_return;
 }
 
 struct promise_new_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_new_tag> {
   struct promise_type {