[coroutines] Build GRO declaration and return GRO statement

Summary:
1. build declaration of the gro local variable that keeps the result of get_return_object.
2. build return statement returning the gro variable
3. emit them during CodeGen
4. sema and CodeGen tests updated

Reviewers: EricWF, rsmith

Reviewed By: rsmith

Subscribers: cfe-commits

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

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

9 files changed, 208 insertions, 29 deletions

diff --git a/include/clang/AST/StmtCXX.h b/include/clang/AST/StmtCXX.h
index 56bfce987f..d6c9654fef 100644
--- a/include/clang/AST/StmtCXX.h
+++ b/include/clang/AST/StmtCXX.h
@@ -308,7 +308,9 @@ class CoroutineBodyStmt final
     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.
+    ReturnValue,   ///< Return value for thunk function: p.get_return_object().
+    ResultDecl,    ///< Declaration holding the result of get_return_object.
+    ReturnStmt,    ///< Return statement for the thunk function.
     ReturnStmtOnAllocFailure, ///< Return statement if allocation failed.
     FirstParamMove ///< First offset for move construction of parameter copies.
   };
@@ -332,7 +334,9 @@ public:
     Stmt *OnFallthrough = nullptr;
     Expr *Allocate = nullptr;
     Expr *Deallocate = nullptr;
-    Stmt *ReturnValue = nullptr;
+    Expr *ReturnValue = nullptr;
+    Stmt *ResultDecl = nullptr;
+    Stmt *ReturnStmt = nullptr;
     Stmt *ReturnStmtOnAllocFailure = nullptr;
     ArrayRef<Stmt *> ParamMoves;
   };
@@ -381,10 +385,11 @@ public:
   Expr *getDeallocate() const {
     return cast_or_null<Expr>(getStoredStmts()[SubStmt::Deallocate]);
   }
-
   Expr *getReturnValueInit() const {
-    return cast_or_null<Expr>(getStoredStmts()[SubStmt::ReturnValue]);
+    return cast<Expr>(getStoredStmts()[SubStmt::ReturnValue]);
   }
+  Stmt *getResultDecl() const { return getStoredStmts()[SubStmt::ResultDecl]; }
+  Stmt *getReturnStmt() const { return getStoredStmts()[SubStmt::ReturnStmt]; }
   Stmt *getReturnStmtOnAllocFailure() const {
     return getStoredStmts()[SubStmt::ReturnStmtOnAllocFailure];
   }
diff --git a/include/clang/Basic/DiagnosticSemaKinds.td b/include/clang/Basic/DiagnosticSemaKinds.td
index 463d057d8f..734a810d05 100644
--- a/include/clang/Basic/DiagnosticSemaKinds.td
+++ b/include/clang/Basic/DiagnosticSemaKinds.td
@@ -8910,6 +8910,8 @@ def err_return_in_coroutine : Error<
   "return statement not allowed in coroutine; did you mean 'co_return'?">;
 def note_declared_coroutine_here : Note<
   "function is a coroutine due to use of '%0' here">;
+def note_promise_member_declared_here : Note<
+  "'%0' is declared here">;
 def err_coroutine_objc_method : Error<
   "Objective-C methods as coroutines are not yet supported">;
 def err_coroutine_unevaluated_context : Error<
diff --git a/lib/AST/StmtCXX.cpp b/lib/AST/StmtCXX.cpp
index aade13ed3b..8466cd61f0 100644
--- a/lib/AST/StmtCXX.cpp
+++ b/lib/AST/StmtCXX.cpp
@@ -88,7 +88,7 @@ const VarDecl *CXXForRangeStmt::getLoopVariable() const {
 }
 
 CoroutineBodyStmt *CoroutineBodyStmt::Create(
-    const ASTContext &C, CoroutineBodyStmt::CtorArgs const& Args) {
+    const ASTContext &C, CoroutineBodyStmt::CtorArgs const &Args) {
   std::size_t Size = totalSizeToAlloc<Stmt *>(
       CoroutineBodyStmt::FirstParamMove + Args.ParamMoves.size());
 
@@ -108,6 +108,8 @@ CoroutineBodyStmt::CoroutineBodyStmt(CoroutineBodyStmt::CtorArgs const &Args)
   SubStmts[CoroutineBodyStmt::Allocate] = Args.Allocate;
   SubStmts[CoroutineBodyStmt::Deallocate] = Args.Deallocate;
   SubStmts[CoroutineBodyStmt::ReturnValue] = Args.ReturnValue;
+  SubStmts[CoroutineBodyStmt::ResultDecl] = Args.ResultDecl;
+  SubStmts[CoroutineBodyStmt::ReturnStmt] = Args.ReturnStmt;
   SubStmts[CoroutineBodyStmt::ReturnStmtOnAllocFailure] =
       Args.ReturnStmtOnAllocFailure;
   std::copy(Args.ParamMoves.begin(), Args.ParamMoves.end(),
diff --git a/lib/CodeGen/CGCoroutine.cpp b/lib/CodeGen/CGCoroutine.cpp
index 0ef680ef66..f01237205f 100644
--- a/lib/CodeGen/CGCoroutine.cpp
+++ b/lib/CodeGen/CGCoroutine.cpp
@@ -309,6 +309,7 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, S.getDeallocate());
 
     EmitStmt(S.getPromiseDeclStmt());
+    EmitStmt(S.getResultDecl()); // FIXME: Gro lifetime is wrong.
 
     EHStack.pushCleanup<CallCoroEnd>(EHCleanup);
 
@@ -329,10 +330,13 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   }
 
   EmitBlock(RetBB);
+  // Emit coro.end before getReturnStmt (and parameter destructors), since
+  // resume and destroy parts of the coroutine should not include them.
   llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
   Builder.CreateCall(CoroEnd, {NullPtr, Builder.getFalse()});
 
-  // FIXME: Emit return for the coroutine return object.
+  if (Stmt *Ret = S.getReturnStmt())
+    EmitStmt(Ret);
 }
 
 // Emit coroutine intrinsic and patch up arguments of the token type.
diff --git a/lib/Sema/AnalysisBasedWarnings.cpp b/lib/Sema/AnalysisBasedWarnings.cpp
index 50ad113fc8..895d00b3eb 100644
--- a/lib/Sema/AnalysisBasedWarnings.cpp
+++ b/lib/Sema/AnalysisBasedWarnings.cpp
@@ -334,6 +334,10 @@ static ControlFlowKind CheckFallThrough(AnalysisDeclContext &AC) {
   bool HasPlainEdge = false;
   bool HasAbnormalEdge = false;
 
+  // In a coroutine, only co_return statements count as normal returns. Remember
+  // if we are processing a coroutine or not.
+  const bool IsCoroutine = isa<CoroutineBodyStmt>(AC.getBody());
+
   // Ignore default cases that aren't likely to be reachable because all
   // enums in a switch(X) have explicit case statements.
   CFGBlock::FilterOptions FO;
@@ -375,7 +379,7 @@ static ControlFlowKind CheckFallThrough(AnalysisDeclContext &AC) {
 
     CFGStmt CS = ri->castAs<CFGStmt>();
     const Stmt *S = CS.getStmt();
-    if (isa<ReturnStmt>(S) || isa<CoreturnStmt>(S)) {
+    if ((isa<ReturnStmt>(S) && !IsCoroutine) || isa<CoreturnStmt>(S)) {
       HasLiveReturn = true;
       continue;
     }
diff --git a/lib/Sema/CoroutineStmtBuilder.h b/lib/Sema/CoroutineStmtBuilder.h
index 4958576219..954a0f100e 100644
--- a/lib/Sema/CoroutineStmtBuilder.h
+++ b/lib/Sema/CoroutineStmtBuilder.h
@@ -28,7 +28,6 @@ class CoroutineStmtBuilder : public CoroutineBodyStmt::CtorArgs {
   sema::FunctionScopeInfo &Fn;
   bool IsValid = true;
   SourceLocation Loc;
-  QualType RetType;
   SmallVector<Stmt *, 4> ParamMovesVector;
   const bool IsPromiseDependentType;
   CXXRecordDecl *PromiseRecordDecl = nullptr;
@@ -61,6 +60,7 @@ private:
   bool makeOnFallthrough();
   bool makeOnException();
   bool makeReturnObject();
+  bool makeGroDeclAndReturnStmt();
   bool makeReturnOnAllocFailure();
   bool makeParamMoves();
 };
diff --git a/lib/Sema/SemaCoroutine.cpp b/lib/Sema/SemaCoroutine.cpp
index c709a1a723..be89cd16de 100644
--- a/lib/Sema/SemaCoroutine.cpp
+++ b/lib/Sema/SemaCoroutine.cpp
@@ -373,7 +373,6 @@ static ExprResult buildPromiseCall(Sema &S, VarDecl *Promise,
   if (PromiseRef.isInvalid())
     return ExprError();
 
-  // Call 'yield_value', passing in E.
   return buildMemberCall(S, PromiseRef.get(), Loc, Name, Args);
 }
 
@@ -780,7 +779,8 @@ bool CoroutineStmtBuilder::buildDependentStatements() {
   assert(!this->IsPromiseDependentType &&
          "coroutine cannot have a dependent promise type");
   this->IsValid = makeOnException() && makeOnFallthrough() &&
-                  makeReturnOnAllocFailure() && makeNewAndDeleteExpr();
+                  makeGroDeclAndReturnStmt() && makeReturnOnAllocFailure() &&
+                  makeNewAndDeleteExpr();
   return this->IsValid;
 }
 
@@ -857,15 +857,16 @@ bool CoroutineStmtBuilder::makeReturnOnAllocFailure() {
   if (ReturnObjectOnAllocationFailure.isInvalid())
     return false;
 
-  // FIXME: ActOnReturnStmt expects a scope that is inside of the function, due
-  //   to CheckJumpOutOfSEHFinally(*this, ReturnLoc, *CurScope->getFnParent());
-  //   S.getCurScope()->getFnParent() == nullptr at ActOnFinishFunctionBody when
-  //   CoroutineBodyStmt is built. Figure it out and fix it.
-  //   Use BuildReturnStmt here to unbreak sanitized tests. (Gor:3/27/2017)
   StmtResult ReturnStmt =
       S.BuildReturnStmt(Loc, ReturnObjectOnAllocationFailure.get());
-  if (ReturnStmt.isInvalid())
+  if (ReturnStmt.isInvalid()) {
+    S.Diag(Found.getFoundDecl()->getLocation(),
+           diag::note_promise_member_declared_here)
+        << DN.getAsString();
+    S.Diag(Fn.FirstCoroutineStmtLoc, diag::note_declared_coroutine_here)
+        << Fn.getFirstCoroutineStmtKeyword();
     return false;
+  }
 
   this->ReturnStmtOnAllocFailure = ReturnStmt.get();
   return true;
@@ -1047,27 +1048,106 @@ bool CoroutineStmtBuilder::makeOnException() {
 }
 
 bool CoroutineStmtBuilder::makeReturnObject() {
-
   // Build implicit 'p.get_return_object()' expression and form initialization
   // of return type from it.
   ExprResult ReturnObject =
       buildPromiseCall(S, Fn.CoroutinePromise, Loc, "get_return_object", None);
   if (ReturnObject.isInvalid())
     return false;
-  QualType RetType = FD.getReturnType();
-  if (!RetType->isDependentType()) {
-    InitializedEntity Entity =
-        InitializedEntity::InitializeResult(Loc, RetType, false);
-    ReturnObject = S.PerformMoveOrCopyInitialization(Entity, nullptr, RetType,
-                                                   ReturnObject.get());
-    if (ReturnObject.isInvalid())
+
+  this->ReturnValue = ReturnObject.get();
+  return true;
+}
+
+static void noteMemberDeclaredHere(Sema &S, Expr *E, FunctionScopeInfo &Fn) {
+  if (auto *MbrRef = dyn_cast<CXXMemberCallExpr>(E)) {
+    auto *MethodDecl = MbrRef->getMethodDecl();
+    S.Diag(MethodDecl->getLocation(), diag::note_promise_member_declared_here)
+        << MethodDecl->getName();
+  }
+  S.Diag(Fn.FirstCoroutineStmtLoc, diag::note_declared_coroutine_here)
+      << Fn.getFirstCoroutineStmtKeyword();
+}
+
+bool CoroutineStmtBuilder::makeGroDeclAndReturnStmt() {
+  assert(!IsPromiseDependentType &&
+         "cannot make statement while the promise type is dependent");
+  assert(this->ReturnValue && "ReturnValue must be already formed");
+
+  QualType const GroType = this->ReturnValue->getType();
+  assert(!GroType->isDependentType() &&
+         "get_return_object type must no longer be dependent");
+
+  QualType const FnRetType = FD.getReturnType();
+  assert(!FnRetType->isDependentType() &&
+         "get_return_object type must no longer be dependent");
+
+  if (FnRetType->isVoidType()) {
+    ExprResult Res = S.ActOnFinishFullExpr(this->ReturnValue, Loc);
+    if (Res.isInvalid())
       return false;
+
+    this->ResultDecl = Res.get();
+    return true;
   }
-  ReturnObject = S.ActOnFinishFullExpr(ReturnObject.get(), Loc);
-  if (ReturnObject.isInvalid())
+
+  if (GroType->isVoidType()) {
+    // Trigger a nice error message.
+    InitializedEntity Entity =
+        InitializedEntity::InitializeResult(Loc, FnRetType, false);
+    S.PerformMoveOrCopyInitialization(Entity, nullptr, FnRetType, ReturnValue);
+    noteMemberDeclaredHere(S, ReturnValue, Fn);
     return false;
+  }
 
-  this->ReturnValue = ReturnObject.get();
+  auto *GroDecl = VarDecl::Create(
+      S.Context, &FD, FD.getLocation(), FD.getLocation(),
+      &S.PP.getIdentifierTable().get("__coro_gro"), GroType,
+      S.Context.getTrivialTypeSourceInfo(GroType, Loc), SC_None);
+
+  S.CheckVariableDeclarationType(GroDecl);
+  if (GroDecl->isInvalidDecl())
+    return false;
+
+  InitializedEntity Entity = InitializedEntity::InitializeVariable(GroDecl);
+  ExprResult Res = S.PerformMoveOrCopyInitialization(Entity, nullptr, GroType,
+                                                     this->ReturnValue);
+  if (Res.isInvalid())
+    return false;
+
+  Res = S.ActOnFinishFullExpr(Res.get());
+  if (Res.isInvalid())
+    return false;
+
+  if (GroType == FnRetType) {
+    GroDecl->setNRVOVariable(true);
+  }
+
+  S.AddInitializerToDecl(GroDecl, Res.get(),
+                         /*DirectInit=*/false);
+
+  S.FinalizeDeclaration(GroDecl);
+
+  // Form a declaration statement for the return declaration, so that AST
+  // visitors can more easily find it.
+  StmtResult GroDeclStmt =
+      S.ActOnDeclStmt(S.ConvertDeclToDeclGroup(GroDecl), Loc, Loc);
+  if (GroDeclStmt.isInvalid())
+    return false;
+
+  this->ResultDecl = GroDeclStmt.get();
+
+  ExprResult declRef = S.BuildDeclRefExpr(GroDecl, GroType, VK_LValue, Loc);
+  if (declRef.isInvalid())
+    return false;
+
+  StmtResult ReturnStmt = S.BuildReturnStmt(Loc, declRef.get());
+  if (ReturnStmt.isInvalid()) {
+    noteMemberDeclaredHere(S, ReturnValue, Fn);
+    return false;
+  }
+
+  this->ReturnStmt = ReturnStmt.get();
   return true;
 }
 
diff --git a/test/CodeGenCoroutines/coro-alloc.cpp b/test/CodeGenCoroutines/coro-alloc.cpp
index 3a65736f6c..b551257f14 100644
--- a/test/CodeGenCoroutines/coro-alloc.cpp
+++ b/test/CodeGenCoroutines/coro-alloc.cpp
@@ -156,6 +156,7 @@ struct std::experimental::coroutine_traits<int, promise_on_alloc_failure_tag> {
 
 // CHECK-LABEL: f4(
 extern "C" int f4(promise_on_alloc_failure_tag) {
+  // CHECK: %[[RetVal:.+]] = alloca i32
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: %[[MEM:.+]] = call i8* @_ZnwmRKSt9nothrow_t(i64 %[[SIZE]], %"struct.std::nothrow_t"* dereferenceable(1) @_ZStL7nothrow)
@@ -163,7 +164,16 @@ extern "C" int f4(promise_on_alloc_failure_tag) {
   // 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]]
+  // CHECK:   %[[FailRet:.+]] = call i32 @_ZNSt12experimental16coroutine_traitsIJi28promise_on_alloc_failure_tagEE12promise_type39get_return_object_on_allocation_failureEv(
+  // CHECK:   store i32 %[[FailRet]], i32* %[[RetVal]]
+  // CHECK:   br label %[[RetBB:.+]]
+
+  // CHECK: [[OKBB]]:
+  // CHECK:   %[[OkRet:.+]] = call i32 @_ZNSt12experimental16coroutine_traitsIJi28promise_on_alloc_failure_tagEE12promise_type17get_return_objectEv(
+  // CHECK:   store i32 %[[OkRet]], i32* %[[RetVal]]
+
+  // CHECK: [[RetBB]]:
+  // CHECK:   %[[LoadRet:.+]] = load i32, i32* %[[RetVal]], align 4
+  // CHECK:   ret i32 %[[LoadRet]]
   co_return;
 }
diff --git a/test/SemaCXX/coroutines.cpp b/test/SemaCXX/coroutines.cpp
index 856110333d..3e33c99881 100644
--- a/test/SemaCXX/coroutines.cpp
+++ b/test/SemaCXX/coroutines.cpp
@@ -746,3 +746,75 @@ coro<T> dependent_uses_nothrow_new(T) {
    co_return;
 }
 template coro<good_promise_13> dependent_uses_nothrow_new(good_promise_13);
+
+struct mismatch_gro_type_tag1 {};
+template<>
+struct std::experimental::coroutine_traits<int, mismatch_gro_type_tag1> {
+  struct promise_type {
+    void get_return_object() {} //expected-note {{'get_return_object' is declared here}}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+    void unhandled_exception();
+  };
+};
+
+extern "C" int f(mismatch_gro_type_tag1) { 
+  // expected-error@-1 {{cannot initialize return object of type 'int' with an rvalue of type 'void'}}
+  co_return; //expected-note {{function is a coroutine due to use of 'co_return' here}}
+}
+
+struct mismatch_gro_type_tag2 {};
+template<>
+struct std::experimental::coroutine_traits<int, mismatch_gro_type_tag2> {
+  struct promise_type {
+    void* get_return_object() {} //expected-note {{'get_return_object' is declared here}}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+    void unhandled_exception();
+  };
+};
+
+extern "C" int f(mismatch_gro_type_tag2) { 
+  // expected-error@-1 {{cannot initialize return object of type 'int' with an lvalue of type 'void *'}}
+  co_return; //expected-note {{function is a coroutine due to use of 'co_return' here}}
+}
+
+struct mismatch_gro_type_tag3 {};
+template<>
+struct std::experimental::coroutine_traits<int, mismatch_gro_type_tag3> {
+  struct promise_type {
+    int get_return_object() {}
+    static void get_return_object_on_allocation_failure() {} //expected-note {{'get_return_object_on_allocation_failure' is declared here}}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+    void unhandled_exception();
+  };
+};
+
+extern "C" int f(mismatch_gro_type_tag3) { 
+  // expected-error@-1 {{cannot initialize return object of type 'int' with an rvalue of type 'void'}}
+  co_return; //expected-note {{function is a coroutine due to use of 'co_return' here}}
+}
+
+
+struct mismatch_gro_type_tag4 {};
+template<>
+struct std::experimental::coroutine_traits<int, mismatch_gro_type_tag4> {
+  struct promise_type {
+    int get_return_object() {}
+    static char* get_return_object_on_allocation_failure() {} //expected-note {{'get_return_object_on_allocation_failure' is declared here}}
+    suspend_always initial_suspend() { return {}; }
+    suspend_always final_suspend() { return {}; }
+    void return_void() {}
+    void unhandled_exception();
+  };
+};
+
+extern "C" int f(mismatch_gro_type_tag4) { 
+  // expected-error@-1 {{cannot initialize return object of type 'int' with an rvalue of type 'char *'}}
+  co_return; //expected-note {{function is a coroutine due to use of 'co_return' here}}
+}
+