[coroutines] Allow co_await and co_yield expressions that return an lvalue to compile

Summary:
The title says it all.


Reviewers: GorNishanov, rsmith

Reviewed By: GorNishanov

Subscribers: rjmccall, cfe-commits

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

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

5 files changed, 106 insertions, 8 deletions

diff --git a/lib/AST/ExprClassification.cpp b/lib/AST/ExprClassification.cpp
index c035a42439..d149bdd0cd 100644
--- a/lib/AST/ExprClassification.cpp
+++ b/lib/AST/ExprClassification.cpp
@@ -190,7 +190,6 @@ static Cl::Kinds ClassifyInternal(ASTContext &Ctx, const Expr *E) {
   case Expr::ArrayInitIndexExprClass:
   case Expr::NoInitExprClass:
   case Expr::DesignatedInitUpdateExprClass:
-  case Expr::CoyieldExprClass:
     return Cl::CL_PRValue;
 
     // Next come the complicated cases.
@@ -414,7 +413,8 @@ static Cl::Kinds ClassifyInternal(ASTContext &Ctx, const Expr *E) {
     return ClassifyInternal(Ctx, cast<InitListExpr>(E)->getInit(0));
 
   case Expr::CoawaitExprClass:
-    return ClassifyInternal(Ctx, cast<CoawaitExpr>(E)->getResumeExpr());
+  case Expr::CoyieldExprClass:
+    return ClassifyInternal(Ctx, cast<CoroutineSuspendExpr>(E)->getResumeExpr());
   }
 
   llvm_unreachable("unhandled expression kind in classification");
diff --git a/lib/CodeGen/CGCoroutine.cpp b/lib/CodeGen/CGCoroutine.cpp
index bc5f6327c9..daefb0edcb 100644
--- a/lib/CodeGen/CGCoroutine.cpp
+++ b/lib/CodeGen/CGCoroutine.cpp
@@ -148,10 +148,16 @@ static SmallString<32> buildSuspendPrefixStr(CGCoroData &Coro, AwaitKind Kind) {
 //
 //  See llvm's docs/Coroutines.rst for more details.
 //
-static RValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
+namespace {
+  struct LValueOrRValue {
+    LValue LV;
+    RValue RV;
+  };
+}
+static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
                                     CoroutineSuspendExpr const &S,
                                     AwaitKind Kind, AggValueSlot aggSlot,
-                                    bool ignoreResult) {
+                                    bool ignoreResult, bool forLValue) {
   auto *E = S.getCommonExpr();
 
   // FIXME: rsmith 5/22/2017. Does it still make sense for us to have a 
@@ -217,7 +223,12 @@ static RValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
 
   // Emit await_resume expression.
   CGF.EmitBlock(ReadyBlock);
-  return CGF.EmitAnyExpr(S.getResumeExpr(), aggSlot, ignoreResult);
+  LValueOrRValue Res;
+  if (forLValue)
+    Res.LV = CGF.EmitLValue(S.getResumeExpr());
+  else
+    Res.RV = CGF.EmitAnyExpr(S.getResumeExpr(), aggSlot, ignoreResult);
+  return Res;
 }
 
 RValue CodeGenFunction::EmitCoawaitExpr(const CoawaitExpr &E,
@@ -225,13 +236,13 @@ RValue CodeGenFunction::EmitCoawaitExpr(const CoawaitExpr &E,
                                         bool ignoreResult) {
   return emitSuspendExpression(*this, *CurCoro.Data, E,
                                CurCoro.Data->CurrentAwaitKind, aggSlot,
-                               ignoreResult);
+                               ignoreResult, /*forLValue*/false).RV;
 }
 RValue CodeGenFunction::EmitCoyieldExpr(const CoyieldExpr &E,
                                         AggValueSlot aggSlot,
                                         bool ignoreResult) {
   return emitSuspendExpression(*this, *CurCoro.Data, E, AwaitKind::Yield,
-                               aggSlot, ignoreResult);
+                               aggSlot, ignoreResult, /*forLValue*/false).RV;
 }
 
 void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
@@ -240,6 +251,38 @@ void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
   EmitBranchThroughCleanup(CurCoro.Data->FinalJD);
 }
 
+
+#ifndef NDEBUG
+static QualType getCoroutineSuspendExprReturnType(const ASTContext &Ctx,
+  const CoroutineSuspendExpr *E) {
+  const auto *RE = E->getResumeExpr();
+  // Is it possible for RE to be a CXXBindTemporaryExpr wrapping
+  // a MemberCallExpr?
+  assert(isa<CallExpr>(RE) && "unexpected suspend expression type");
+  return cast<CallExpr>(RE)->getCallReturnType(Ctx);
+}
+#endif
+
+LValue
+CodeGenFunction::EmitCoawaitLValue(const CoawaitExpr *E) {
+  assert(getCoroutineSuspendExprReturnType(getContext(), E)->isReferenceType() &&
+         "Can't have a scalar return unless the return type is a "
+         "reference type!");
+  return emitSuspendExpression(*this, *CurCoro.Data, *E,
+                               CurCoro.Data->CurrentAwaitKind, AggValueSlot::ignored(),
+                               /*ignoreResult*/false, /*forLValue*/true).LV;
+}
+
+LValue
+CodeGenFunction::EmitCoyieldLValue(const CoyieldExpr *E) {
+  assert(getCoroutineSuspendExprReturnType(getContext(), E)->isReferenceType() &&
+         "Can't have a scalar return unless the return type is a "
+         "reference type!");
+  return emitSuspendExpression(*this, *CurCoro.Data, *E,
+                               AwaitKind::Yield, AggValueSlot::ignored(),
+                               /*ignoreResult*/false, /*forLValue*/true).LV;
+}
+
 // Hunts for the parameter reference in the parameter copy/move declaration.
 namespace {
 struct GetParamRef : public StmtVisitor<GetParamRef> {
diff --git a/lib/CodeGen/CGExpr.cpp b/lib/CodeGen/CGExpr.cpp
index 1b3103fd71..15829888ea 100644
--- a/lib/CodeGen/CGExpr.cpp
+++ b/lib/CodeGen/CGExpr.cpp
@@ -1158,6 +1158,11 @@ LValue CodeGenFunction::EmitLValue(const Expr *E) {
 
   case Expr::MaterializeTemporaryExprClass:
     return EmitMaterializeTemporaryExpr(cast<MaterializeTemporaryExpr>(E));
+
+  case Expr::CoawaitExprClass:
+    return EmitCoawaitLValue(cast<CoawaitExpr>(E));
+  case Expr::CoyieldExprClass:
+    return EmitCoyieldLValue(cast<CoyieldExpr>(E));
   }
 }
 
diff --git a/lib/CodeGen/CodeGenFunction.h b/lib/CodeGen/CodeGenFunction.h
index a179276a07..831eedf9e4 100644
--- a/lib/CodeGen/CodeGenFunction.h
+++ b/lib/CodeGen/CodeGenFunction.h
@@ -2550,9 +2550,11 @@ public:
   RValue EmitCoawaitExpr(const CoawaitExpr &E,
                          AggValueSlot aggSlot = AggValueSlot::ignored(),
                          bool ignoreResult = false);
+  LValue EmitCoawaitLValue(const CoawaitExpr *E);
   RValue EmitCoyieldExpr(const CoyieldExpr &E,
                          AggValueSlot aggSlot = AggValueSlot::ignored(),
                          bool ignoreResult = false);
+  LValue EmitCoyieldLValue(const CoyieldExpr *E);
   RValue EmitCoroutineIntrinsic(const CallExpr *E, unsigned int IID);
 
   void EnterCXXTryStmt(const CXXTryStmt &S, bool IsFnTryBlock = false);
diff --git a/test/CodeGenCoroutines/coro-await.cpp b/test/CodeGenCoroutines/coro-await.cpp
index 1e2deaa8f5..fc6559f1e0 100644
--- a/test/CodeGenCoroutines/coro-await.cpp
+++ b/test/CodeGenCoroutines/coro-await.cpp
@@ -1,4 +1,5 @@
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fcoroutines-ts -std=c++14 -emit-llvm %s -o - -disable-llvm-passes | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fcoroutines-ts -std=c++14 \
+// RUN:   -emit-llvm %s -o - -disable-llvm-passes -Wno-coroutine -Wno-unused | FileCheck %s
 
 namespace std {
 namespace experimental {
@@ -278,3 +279,50 @@ void AwaitLValue() {
   suspend_always lval;
   co_await lval;
 }
+
+struct RefTag { };
+
+struct AwaitResumeReturnsLValue {
+  bool await_ready();
+  void await_suspend(std::experimental::coroutine_handle<>);
+  RefTag& await_resume();
+};
+
+
+template<>
+struct std::experimental::coroutine_traits<void,double> {
+  struct promise_type {
+    void get_return_object();
+    init_susp initial_suspend();
+    final_susp final_suspend();
+    void return_void();
+    AwaitResumeReturnsLValue yield_value(int);
+  };
+};
+
+// Verifies that we don't crash when returning an lvalue from an await_resume()
+// expression.
+// CHECK-LABEL:  define void @_Z18AwaitReturnsLValued(double)
+void AwaitReturnsLValue(double) {
+  AwaitResumeReturnsLValue a;
+  // CHECK: %[[AVAR:.+]] = alloca %struct.AwaitResumeReturnsLValue,
+  // CHECK: %[[XVAR:.+]] = alloca %struct.RefTag*,
+
+  // CHECK: %[[YVAR:.+]] = alloca %struct.RefTag*,
+  // CHECK-NEXT: %[[TMP1:.+]] = alloca %struct.AwaitResumeReturnsLValue,
+
+  // CHECK: %[[ZVAR:.+]] = alloca %struct.RefTag*,
+  // CHECK-NEXT: %[[TMP2:.+]] = alloca %struct.AwaitResumeReturnsLValue,
+
+  // CHECK: %[[RES1:.+]] = call dereferenceable({{.*}}) %struct.RefTag* @_ZN24AwaitResumeReturnsLValue12await_resumeEv(%struct.AwaitResumeReturnsLValue* %[[AVAR]])
+  // CHECK-NEXT: store %struct.RefTag* %[[RES1]], %struct.RefTag** %[[XVAR]],
+  RefTag& x = co_await a;
+
+  // CHECK: %[[RES2:.+]] = call dereferenceable({{.*}}) %struct.RefTag* @_ZN24AwaitResumeReturnsLValue12await_resumeEv(%struct.AwaitResumeReturnsLValue* %[[TMP1]])
+  // CHECK-NEXT: store %struct.RefTag* %[[RES2]], %struct.RefTag** %[[YVAR]],
+
+  RefTag& y = co_await AwaitResumeReturnsLValue{};
+  // CHECK: %[[RES3:.+]] = call dereferenceable({{.*}}) %struct.RefTag* @_ZN24AwaitResumeReturnsLValue12await_resumeEv(%struct.AwaitResumeReturnsLValue* %[[TMP2]])
+  // CHECK-NEXT: store %struct.RefTag* %[[RES3]], %struct.RefTag** %[[ZVAR]],
+  RefTag& z = co_yield 42;
+}