[coroutines] Add allocation and deallocation substatements.

Summary:
SemaCoroutine: Add allocation / deallocation substatements.
CGCoroutine/Test: Emit allocation and deallocation + test.

Reviewers: rsmith

Subscribers: ABataev, EricWF, llvm-commits, mehdi_amini

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

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@285306 91177308-0d34-0410-b5e6-96231b3b80d8

include/clang/AST/StmtCXX.h

@@ -304,6 +304,8 @@ class CoroutineBodyStmt : public Stmt {
     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.
     FirstParamMove ///< First offset for move construction of parameter copies.
   };
@@ -313,6 +315,7 @@ class CoroutineBodyStmt : public Stmt {
 public:
   CoroutineBodyStmt(Stmt *Body, Stmt *Promise, Stmt *InitSuspend,
                     Stmt *FinalSuspend, Stmt *OnException, Stmt *OnFallthrough,
+                    Expr *Allocate, Stmt *Deallocate,
                     Expr *ReturnValue, ArrayRef<Expr *> ParamMoves)
       : Stmt(CoroutineBodyStmtClass) {
     SubStmts[CoroutineBodyStmt::Body] = Body;
@@ -321,6 +324,8 @@ public:
     SubStmts[CoroutineBodyStmt::FinalSuspend] = FinalSuspend;
     SubStmts[CoroutineBodyStmt::OnException] = OnException;
     SubStmts[CoroutineBodyStmt::OnFallthrough] = OnFallthrough;
+    SubStmts[CoroutineBodyStmt::Allocate] = Allocate;
+    SubStmts[CoroutineBodyStmt::Deallocate] = Deallocate;
     SubStmts[CoroutineBodyStmt::ReturnValue] = ReturnValue;
     // FIXME: Tail-allocate space for parameter move expressions and store them.
     assert(ParamMoves.empty() && "not implemented yet");
@@ -345,6 +350,9 @@ public:
     return SubStmts[SubStmt::OnFallthrough];
   }
 
+  Expr *getAllocate() const { return cast<Expr>(SubStmts[SubStmt::Allocate]); }
+  Stmt *getDeallocate() const { return SubStmts[SubStmt::Deallocate]; }
+
   Expr *getReturnValueInit() const {
     return cast<Expr>(SubStmts[SubStmt::ReturnValue]);
   }

lib/CodeGen/CGCoroutine.cpp

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "CodeGenFunction.h"
+#include "clang/AST/StmtCXX.h"
 
 using namespace clang;
 using namespace CodeGen;
@@ -36,9 +37,10 @@ struct CGCoroData {
 clang::CodeGen::CodeGenFunction::CGCoroInfo::CGCoroInfo() {}
 CodeGenFunction::CGCoroInfo::~CGCoroInfo() {}
 
-static bool createCoroData(CodeGenFunction &CGF,
+static void createCoroData(CodeGenFunction &CGF,
                            CodeGenFunction::CGCoroInfo &CurCoro,
-                           llvm::CallInst *CoroId, CallExpr const *CoroIdExpr) {
+                           llvm::CallInst *CoroId,
+                           CallExpr const *CoroIdExpr = nullptr) {
   if (CurCoro.Data) {
     if (CurCoro.Data->CoroIdExpr)
       CGF.CGM.Error(CoroIdExpr->getLocStart(),
@@ -49,13 +51,27 @@ static bool createCoroData(CodeGenFunction &CGF,
     else
       llvm_unreachable("EmitCoroutineBodyStatement called twice?");
 
-    return false;
+    return;
   }
 
   CurCoro.Data = std::unique_ptr<CGCoroData>(new CGCoroData);
   CurCoro.Data->CoroId = CoroId;
   CurCoro.Data->CoroIdExpr = CoroIdExpr;
-  return true;
+}
+
+void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
+  auto *NullPtr = llvm::ConstantPointerNull::get(Builder.getInt8PtrTy());
+  auto &TI = CGM.getContext().getTargetInfo();
+  unsigned NewAlign = TI.getNewAlign() / TI.getCharWidth();
+
+  auto *CoroId = Builder.CreateCall(
+      CGM.getIntrinsic(llvm::Intrinsic::coro_id),
+      {Builder.getInt32(NewAlign), NullPtr, NullPtr, NullPtr});
+  createCoroData(*this, CurCoro, CoroId);
+
+  EmitScalarExpr(S.getAllocate());
+  // FIXME: Emit the rest of the coroutine.
+  EmitStmt(S.getDeallocate());
 }
 
 // Emit coroutine intrinsic and patch up arguments of the token type.

lib/CodeGen/CGStmt.cpp

@@ -142,6 +142,8 @@ void CodeGenFunction::EmitStmt(const Stmt *S) {
   case Stmt::GCCAsmStmtClass:   // Intentional fall-through.
   case Stmt::MSAsmStmtClass:    EmitAsmStmt(cast<AsmStmt>(*S));           break;
   case Stmt::CoroutineBodyStmtClass:
+    EmitCoroutineBody(cast<CoroutineBodyStmt>(*S));
+    break;
   case Stmt::CoreturnStmtClass:
     CGM.ErrorUnsupported(S, "coroutine");
     break;

lib/CodeGen/CodeGenFunction.h

@@ -2311,6 +2311,7 @@ public:
   void EmitObjCAtSynchronizedStmt(const ObjCAtSynchronizedStmt &S);
   void EmitObjCAutoreleasePoolStmt(const ObjCAutoreleasePoolStmt &S);
 
+  void EmitCoroutineBody(const CoroutineBodyStmt &S);
   RValue EmitCoroutineIntrinsic(const CallExpr *E, unsigned int IID);
 
   void EnterCXXTryStmt(const CXXTryStmt &S, bool IsFnTryBlock = false);

lib/Sema/SemaCoroutine.cpp

@@ -78,7 +78,7 @@ static QualType lookupPromiseType(Sema &S, const FunctionProtoType *FnType,
   auto *Promise = R.getAsSingle<TypeDecl>();
   if (!Promise) {
     S.Diag(Loc, diag::err_implied_std_coroutine_traits_promise_type_not_found)
-      << RD;
+        << RD;
     return QualType();
   }
 
@@ -92,7 +92,7 @@ static QualType lookupPromiseType(Sema &S, const FunctionProtoType *FnType,
     PromiseType = S.Context.getElaboratedType(ETK_None, NNS, PromiseType);
 
     S.Diag(Loc, diag::err_implied_std_coroutine_traits_promise_type_not_class)
-      << PromiseType;
+        << PromiseType;
     return QualType();
   }
 
@@ -121,7 +121,7 @@ checkCoroutineContext(Sema &S, SourceLocation Loc, StringRef Keyword) {
     //
     // FIXME: We assume that this really means that a coroutine cannot
     //        be a constructor or destructor.
-    S.Diag(Loc, diag::err_coroutine_ctor_dtor)
+    S.Diag(Loc, diag::err_coroutine_ctor_dtor) 
       << isa<CXXDestructorDecl>(FD) << Keyword;
   } else if (FD->isConstexpr()) {
     S.Diag(Loc, diag::err_coroutine_constexpr) << Keyword;
@@ -130,8 +130,8 @@ checkCoroutineContext(Sema &S, SourceLocation Loc, StringRef Keyword) {
   } else if (FD->isMain()) {
     S.Diag(FD->getLocStart(), diag::err_coroutine_main);
     S.Diag(Loc, diag::note_declared_coroutine_here)
-      << (Keyword == "co_await" ? 0 :
-          Keyword == "co_yield" ? 1 : 2);
+        << (Keyword == "co_await" ? 0 : 
+            Keyword == "co_yield" ? 1 : 2);
   } else {
     auto *ScopeInfo = S.getCurFunction();
     assert(ScopeInfo && "missing function scope for function");
@@ -162,6 +162,26 @@ checkCoroutineContext(Sema &S, SourceLocation Loc, StringRef Keyword) {
   return nullptr;
 }
 
+static Expr *buildBuiltinCall(Sema &S, SourceLocation Loc, Builtin::ID Id,
+                              MutableArrayRef<Expr *> CallArgs) {
+  StringRef Name = S.Context.BuiltinInfo.getName(Id);
+  LookupResult R(S, &S.Context.Idents.get(Name), Loc, Sema::LookupOrdinaryName);
+  S.LookupName(R, S.TUScope, /*AllowBuiltinCreation=*/true);
+
+  auto *BuiltInDecl = R.getAsSingle<FunctionDecl>();
+  assert(BuiltInDecl && "failed to find builtin declaration");
+
+  ExprResult DeclRef =
+      S.BuildDeclRefExpr(BuiltInDecl, BuiltInDecl->getType(), VK_LValue, Loc);
+  assert(DeclRef.isUsable() && "Builtin reference cannot fail");
+
+  ExprResult Call =
+      S.ActOnCallExpr(/*Scope=*/nullptr, DeclRef.get(), Loc, CallArgs, Loc);
+
+  assert(!Call.isInvalid() && "Call to builtin cannot fail!");
+  return Call.get();
+}
+
 /// Build a call to 'operator co_await' if there is a suitable operator for
 /// the given expression.
 static ExprResult buildOperatorCoawaitCall(Sema &SemaRef, Scope *S,
@@ -204,7 +224,7 @@ static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, SourceLocation Loc,
   const StringRef Funcs[] = {"await_ready", "await_suspend", "await_resume"};
   for (size_t I = 0, N = llvm::array_lengthof(Funcs); I != N; ++I) {
     Expr *Operand = new (S.Context) OpaqueValueExpr(
-        Loc, E->getType(), VK_LValue, E->getObjectKind(), E);
+      Loc, E->getType(), VK_LValue, E->getObjectKind(), E);
 
     // FIXME: Pass coroutine handle to await_suspend.
     ExprResult Result = buildMemberCall(S, Operand, Loc, Funcs[I], None);
@@ -406,6 +426,113 @@ static ExprResult buildStdCurrentExceptionCall(Sema &S, SourceLocation Loc) {
   return Res;
 }
 
+// Find an appropriate delete for the promise.
+static FunctionDecl *findDeleteForPromise(Sema &S, SourceLocation Loc,
+                                          QualType PromiseType) {
+  FunctionDecl *OperatorDelete = nullptr;
+
+  DeclarationName DeleteName =
+      S.Context.DeclarationNames.getCXXOperatorName(OO_Delete);
+
+  auto *PointeeRD = PromiseType->getAsCXXRecordDecl();
+  assert(PointeeRD && "PromiseType must be a CxxRecordDecl type");
+
+  if (S.FindDeallocationFunction(Loc, PointeeRD, DeleteName, OperatorDelete))
+    return nullptr;
+
+  if (!OperatorDelete) {
+    // Look for a global declaration.
+    const bool CanProvideSize = S.isCompleteType(Loc, PromiseType);
+    const bool Overaligned = false;
+    OperatorDelete = S.FindUsualDeallocationFunction(Loc, CanProvideSize,
+                                                     Overaligned, DeleteName);
+  }
+  S.MarkFunctionReferenced(Loc, OperatorDelete);
+  return OperatorDelete;
+}
+
+// Builds allocation and deallocation for the coroutine. Returns false on
+// failure.
+static bool buildAllocationAndDeallocation(Sema &S, SourceLocation Loc,
+                                           FunctionScopeInfo *Fn,
+                                           Expr *&Allocation,
+                                           Stmt *&Deallocation) {
+  TypeSourceInfo *TInfo = Fn->CoroutinePromise->getTypeSourceInfo();
+  QualType PromiseType = TInfo->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 support for stateful allocators.
+
+  FunctionDecl *OperatorNew = nullptr;
+  FunctionDecl *OperatorDelete = nullptr;
+  FunctionDecl *UnusedResult = nullptr;
+  bool PassAlignment = false;
+
+  S.FindAllocationFunctions(Loc, SourceRange(),
+                            /*UseGlobal*/ false, PromiseType,
+                            /*isArray*/ false, PassAlignment,
+                            /*PlacementArgs*/ None, OperatorNew, UnusedResult);
+
+  OperatorDelete = findDeleteForPromise(S, Loc, PromiseType);
+
+  if (!OperatorDelete || !OperatorNew)
+    return false;
+
+  Expr *FramePtr =
+      buildBuiltinCall(S, Loc, Builtin::BI__builtin_coro_frame, {});
+
+  Expr *FrameSize =
+      buildBuiltinCall(S, Loc, Builtin::BI__builtin_coro_size, {});
+
+  // Make new call.
+
+  ExprResult NewRef =
+      S.BuildDeclRefExpr(OperatorNew, OperatorNew->getType(), VK_LValue, Loc);
+  if (NewRef.isInvalid())
+    return false;
+
+  ExprResult NewExpr =
+      S.ActOnCallExpr(S.getCurScope(), NewRef.get(), Loc, FrameSize, Loc);
+  if (NewExpr.isInvalid())
+    return false;
+
+  Allocation = NewExpr.get();
+
+  // Make delete call.
+
+  QualType OpDeleteQualType = OperatorDelete->getType();
+
+  ExprResult DeleteRef =
+      S.BuildDeclRefExpr(OperatorDelete, OpDeleteQualType, VK_LValue, Loc);
+  if (DeleteRef.isInvalid())
+    return false;
+
+  Expr *CoroFree =
+      buildBuiltinCall(S, Loc, Builtin::BI__builtin_coro_free, {FramePtr});
+
+  SmallVector<Expr *, 2> DeleteArgs{CoroFree};
+
+  // Check if we need to pass the size.
+  const auto *OpDeleteType =
+      OpDeleteQualType.getTypePtr()->getAs<FunctionProtoType>();
+  if (OpDeleteType->getNumParams() > 1)
+    DeleteArgs.push_back(FrameSize);
+
+  ExprResult DeleteExpr =
+      S.ActOnCallExpr(S.getCurScope(), DeleteRef.get(), Loc, DeleteArgs, Loc);
+  if (DeleteExpr.isInvalid())
+    return false;
+
+  Deallocation = DeleteExpr.get();
+
+  return true;
+}
+
 void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
   FunctionScopeInfo *Fn = getCurFunction();
   assert(Fn && !Fn->CoroutineStmts.empty() && "not a coroutine");
@@ -416,8 +543,8 @@ void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
     Diag(Fn->FirstReturnLoc, diag::err_return_in_coroutine);
     auto *First = Fn->CoroutineStmts[0];
     Diag(First->getLocStart(), diag::note_declared_coroutine_here)
-      << (isa<CoawaitExpr>(First) ? 0 :
-          isa<CoyieldExpr>(First) ? 1 : 2);
+        << (isa<CoawaitExpr>(First) ? 0 :
+            isa<CoyieldExpr>(First) ? 1 : 2);
   }
 
   bool AnyCoawaits = false;
@@ -460,7 +587,12 @@ void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
   if (FinalSuspend.isInvalid())
     return FD->setInvalidDecl();
 
-  // Try to form 'p.return_void();' expression statement to handle
+  // Form and check allocation and deallocation calls.
+  Expr *Allocation = nullptr;
+  Stmt *Deallocation = nullptr;
+  if (!buildAllocationAndDeallocation(*this, Loc, Fn, Allocation, Deallocation))
+    return FD->setInvalidDecl();
+
   // control flowing off the end of the coroutine.
   // Also try to form 'p.set_exception(std::current_exception());' to handle
   // uncaught exceptions.
@@ -517,7 +649,7 @@ void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
   // Build implicit 'p.get_return_object()' expression and form initialization
   // of return type from it.
   ExprResult ReturnObject =
-    buildPromiseCall(*this, Fn, Loc, "get_return_object", None);
+      buildPromiseCall(*this, Fn, Loc, "get_return_object", None);
   if (ReturnObject.isInvalid())
     return FD->setInvalidDecl();
   QualType RetType = FD->getReturnType();
@@ -539,5 +671,6 @@ void Sema::CheckCompletedCoroutineBody(FunctionDecl *FD, Stmt *&Body) {
   // Build body for the coroutine wrapper statement.
   Body = new (Context) CoroutineBodyStmt(
       Body, PromiseStmt.get(), InitialSuspend.get(), FinalSuspend.get(),
-      SetException.get(), Fallthrough.get(), ReturnObject.get(), ParamMoves);
+      SetException.get(), Fallthrough.get(), Allocation, Deallocation,
+      ReturnObject.get(), ParamMoves);
 }

test/CodeGenCoroutines/coro-alloc.cpp

@@ -0,0 +1,114 @@
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fcoroutines-ts -std=c++14 -emit-llvm %s -o - -disable-llvm-passes | FileCheck %s
+
+namespace std {
+namespace experimental {
+template <typename... T>
+struct coroutine_traits; // expected-note {{declared here}}
+}
+}
+
+struct suspend_always {
+  bool await_ready() { return false; }
+  void await_suspend() {}
+  void await_resume() {}
+};
+
+struct global_new_delete_tag {};
+
+template<>
+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( 
+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_await suspend_always{};
+}
+
+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( 
+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_await suspend_always{};
+}
+
+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( 
+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]])
+  co_await suspend_always{};
+}
+
+struct promise_sized_delete_tag {};
+
+template<>
+struct std::experimental::coroutine_traits<void, promise_sized_delete_tag> {
+  struct promise_type {
+    void operator delete(void*, unsigned long);
+    void get_return_object() {}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+  };
+};
+
+// CHECK-LABEL: f3( 
+extern "C" void f3(promise_sized_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: %[[SIZE2:.+]] = call i64 @llvm.coro.size.i64()
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJv24promise_sized_delete_tagEE12promise_typedlEPvm(i8* %[[MEM]], i64 %[[SIZE2]])
+  co_await suspend_always{};
+}