[coroutines] Build implicit return_value / return_void calls for co_return.

llvm-svn: 253816

3 files changed, 51 insertions, 25 deletions

diff --git a/clang/include/clang/AST/StmtCXX.h b/clang/include/clang/AST/StmtCXX.h
index 982181e3717..7f754ac4448 100644
--- a/clang/include/clang/AST/StmtCXX.h
+++ b/clang/include/clang/AST/StmtCXX.h
@@ -328,50 +328,47 @@ public:
 ///
 /// This statament models the initialization of the coroutine promise
 /// (encapsulating the eventual notional return value) from an expression
-/// (or braced-init-list).
+/// (or braced-init-list), followed by termination of the coroutine.
 ///
-/// This initialization is modeled by a call to one of:
+/// This initialization is modeled by the evaluation of the operand
+/// followed by a call to one of:
 ///   <promise>.return_value(<operand>)
 ///   <promise>.return_void()
 /// which we name the "promise call".
 class CoreturnStmt : public Stmt {
   SourceLocation CoreturnLoc;
 
-  /// The operand of the 'co_return' statement.
-  Stmt *Operand;
-  /// The implied call to the promise object. May be null if the
-  /// coroutine has not yet been finalized.
-  Stmt *PromiseCall;
+  enum SubStmt { Operand, PromiseCall, Count };
+  Stmt *SubStmts[SubStmt::Count];
 
   friend class ASTStmtReader;
 public:
-  CoreturnStmt(SourceLocation CoreturnLoc, Stmt *Operand)
-      : Stmt(CoreturnStmtClass), CoreturnLoc(CoreturnLoc),
-        Operand(Operand), PromiseCall(nullptr) {}
+  CoreturnStmt(SourceLocation CoreturnLoc, Stmt *Operand, Stmt *PromiseCall)
+      : Stmt(CoreturnStmtClass), CoreturnLoc(CoreturnLoc) {
+    SubStmts[SubStmt::Operand] = Operand;
+    SubStmts[SubStmt::PromiseCall] = PromiseCall;
+  }
 
   SourceLocation getKeywordLoc() const { return CoreturnLoc; }
 
   /// \brief Retrieve the operand of the 'co_return' statement. Will be nullptr
   /// if none was specified.
-  Expr *getOperand() const { return static_cast<Expr*>(Operand); }
+  Expr *getOperand() const { return static_cast<Expr*>(SubStmts[Operand]); }
 
   /// \brief Retrieve the promise call that results from this 'co_return'
   /// statement. Will be nullptr if either the coroutine has not yet been
   /// finalized or the coroutine has no eventual return type.
-  Expr *getPromiseCall() const { return static_cast<Expr*>(PromiseCall); }
-
-  /// \brief Set the resolved promise call. This is delayed until the
-  /// complete coroutine body has been parsed and the promise type is known.
-  void finalize(Stmt *PC) { PromiseCall = PC; }
+  Expr *getPromiseCall() const {
+    return static_cast<Expr*>(SubStmts[PromiseCall]);
+  }
 
   SourceLocation getLocStart() const LLVM_READONLY { return CoreturnLoc; }
   SourceLocation getLocEnd() const LLVM_READONLY {
-    return Operand->getLocEnd();
+    return getOperand()->getLocEnd();
   }
 
   child_range children() {
-    Stmt **Which = PromiseCall ? &PromiseCall : &Operand;
-    return child_range(Which, Which + 1);
+    return child_range(SubStmts, SubStmts + SubStmt::Count);
   }
 
   static bool classof(const Stmt *T) {
diff --git a/clang/lib/Sema/SemaCoroutine.cpp b/clang/lib/Sema/SemaCoroutine.cpp
index 20c9c551d46..b1dfd669383 100644
--- a/clang/lib/Sema/SemaCoroutine.cpp
+++ b/clang/lib/Sema/SemaCoroutine.cpp
@@ -253,8 +253,9 @@ ExprResult Sema::BuildCoawaitExpr(SourceLocation Loc, Expr *E) {
   return Res;
 }
 
-static ExprResult buildYieldValueCall(Sema &S, FunctionScopeInfo *Coroutine,
-                                      SourceLocation Loc, Expr *E) {
+static ExprResult buildPromiseCall(Sema &S, FunctionScopeInfo *Coroutine,
+                                   SourceLocation Loc, StringRef Name,
+                                   MutableArrayRef<Expr *> Args) {
   assert(Coroutine->CoroutinePromise && "no promise for coroutine");
 
   // Form a reference to the promise.
@@ -265,7 +266,7 @@ static ExprResult buildYieldValueCall(Sema &S, FunctionScopeInfo *Coroutine,
     return ExprError();
 
   // Call 'yield_value', passing in E.
-  return buildMemberCall(S, PromiseRef.get(), Loc, "yield_value", E);
+  return buildMemberCall(S, PromiseRef.get(), Loc, Name, Args);
 }
 
 ExprResult Sema::ActOnCoyieldExpr(Scope *S, SourceLocation Loc, Expr *E) {
@@ -280,7 +281,8 @@ ExprResult Sema::ActOnCoyieldExpr(Scope *S, SourceLocation Loc, Expr *E) {
     return ExprError();
 
   // Build yield_value call.
-  ExprResult Awaitable = buildYieldValueCall(*this, Coroutine, Loc, E);
+  ExprResult Awaitable =
+      buildPromiseCall(*this, Coroutine, Loc, "yield_value", E);
   if (Awaitable.isInvalid())
     return ExprError();
 
@@ -338,8 +340,22 @@ StmtResult Sema::BuildCoreturnStmt(SourceLocation Loc, Expr *E) {
   if (!Coroutine)
     return StmtError();
 
-  // FIXME: Build return_* calls.
-  Stmt *Res = new (Context) CoreturnStmt(Loc, E);
+  // FIXME: If the operand is a reference to a variable that's about to go out
+  // ot scope, we should treat the operand as an xvalue for this overload
+  // resolution.
+  ExprResult PC;
+  if (E && !E->getType()->isVoidType()) {
+    PC = buildPromiseCall(*this, Coroutine, Loc, "return_value", E);
+  } else {
+    E = MakeFullDiscardedValueExpr(E).get();
+    PC = buildPromiseCall(*this, Coroutine, Loc, "return_void", None);
+  }
+  if (PC.isInvalid())
+    return StmtError();
+
+  Expr *PCE = ActOnFinishFullExpr(PC.get()).get();
+
+  Stmt *Res = new (Context) CoreturnStmt(Loc, E, PCE);
   Coroutine->CoroutineStmts.push_back(Res);
   return Res;
 }
diff --git a/clang/test/SemaCXX/coroutines.cpp b/clang/test/SemaCXX/coroutines.cpp
index d5968fcd303..85cc919c6d1 100644
--- a/clang/test/SemaCXX/coroutines.cpp
+++ b/clang/test/SemaCXX/coroutines.cpp
@@ -52,6 +52,8 @@ struct promise {
   awaitable yield_value(int); // expected-note {{candidate}}
   awaitable yield_value(yielded_thing); // expected-note {{candidate}}
   not_awaitable yield_value(void()); // expected-note {{candidate}}
+  void return_void();
+  void return_value(int); // expected-note {{here}}
 };
 
 void yield() {
@@ -65,6 +67,17 @@ void yield() {
   co_yield yield; // expected-error {{no member named 'await_ready' in 'not_awaitable'}}
 }
 
+void coreturn(int n) {
+  co_await a;
+  if (n == 0)
+    co_return 3;
+  if (n == 1)
+    co_return {4};
+  if (n == 2)
+    co_return "foo"; // expected-error {{cannot initialize a parameter of type 'int' with an lvalue of type 'const char [4]'}}
+  co_return;
+}
+
 void mixed_yield() {
   co_yield 0; // expected-note {{use of 'co_yield'}}
   return; // expected-error {{not allowed in coroutine}}