1 files changed, 76 insertions, 32 deletions

diff --git a/include/clang/AST/StmtCXX.h b/include/clang/AST/StmtCXX.h
index 8eef34cb21b3c..56bfce987f5bd 100644
--- a/include/clang/AST/StmtCXX.h
+++ b/include/clang/AST/StmtCXX.h
@@ -296,7 +296,9 @@ public:
 /// \brief Represents the body of a coroutine. This wraps the normal function
 /// body and holds the additional semantic context required to set up and tear
 /// down the coroutine frame.
-class CoroutineBodyStmt : public Stmt {
+class CoroutineBodyStmt final
+    : public Stmt,
+      private llvm::TrailingObjects<CoroutineBodyStmt, Stmt *> {
   enum SubStmt {
     Body,          ///< The body of the coroutine.
     Promise,       ///< The promise statement.
@@ -307,65 +309,100 @@ class CoroutineBodyStmt : public Stmt {
     Allocate,      ///< Coroutine frame memory allocation.
     Deallocate,    ///< Coroutine frame memory deallocation.
     ReturnValue,   ///< Return value for thunk function.
+    ReturnStmtOnAllocFailure, ///< Return statement if allocation failed.
     FirstParamMove ///< First offset for move construction of parameter copies.
   };
-  Stmt *SubStmts[SubStmt::FirstParamMove];
+  unsigned NumParams;
 
   friend class ASTStmtReader;
+  friend TrailingObjects;
+
+  Stmt **getStoredStmts() { return getTrailingObjects<Stmt *>(); }
+
+  Stmt *const *getStoredStmts() const { return getTrailingObjects<Stmt *>(); }
+
+public:
+
+  struct CtorArgs {
+    Stmt *Body = nullptr;
+    Stmt *Promise = nullptr;
+    Expr *InitialSuspend = nullptr;
+    Expr *FinalSuspend = nullptr;
+    Stmt *OnException = nullptr;
+    Stmt *OnFallthrough = nullptr;
+    Expr *Allocate = nullptr;
+    Expr *Deallocate = nullptr;
+    Stmt *ReturnValue = nullptr;
+    Stmt *ReturnStmtOnAllocFailure = nullptr;
+    ArrayRef<Stmt *> ParamMoves;
+  };
+
+private:
+
+  CoroutineBodyStmt(CtorArgs const& Args);
+
 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;
-    SubStmts[CoroutineBodyStmt::Promise] = Promise;
-    SubStmts[CoroutineBodyStmt::InitSuspend] = InitSuspend;
-    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");
+  static CoroutineBodyStmt *Create(const ASTContext &C, CtorArgs const &Args);
+
+  bool hasDependentPromiseType() const {
+    return getPromiseDecl()->getType()->isDependentType();
   }
 
   /// \brief Retrieve the body of the coroutine as written. This will be either
   /// a CompoundStmt or a TryStmt.
   Stmt *getBody() const {
-    return SubStmts[SubStmt::Body];
+    return getStoredStmts()[SubStmt::Body];
   }
 
-  Stmt *getPromiseDeclStmt() const { return SubStmts[SubStmt::Promise]; }
+  Stmt *getPromiseDeclStmt() const {
+    return getStoredStmts()[SubStmt::Promise];
+  }
   VarDecl *getPromiseDecl() const {
     return cast<VarDecl>(cast<DeclStmt>(getPromiseDeclStmt())->getSingleDecl());
   }
 
-  Stmt *getInitSuspendStmt() const { return SubStmts[SubStmt::InitSuspend]; }
-  Stmt *getFinalSuspendStmt() const { return SubStmts[SubStmt::FinalSuspend]; }
+  Stmt *getInitSuspendStmt() const {
+    return getStoredStmts()[SubStmt::InitSuspend];
+  }
+  Stmt *getFinalSuspendStmt() const {
+    return getStoredStmts()[SubStmt::FinalSuspend];
+  }
 
-  Stmt *getExceptionHandler() const { return SubStmts[SubStmt::OnException]; }
+  Stmt *getExceptionHandler() const {
+    return getStoredStmts()[SubStmt::OnException];
+  }
   Stmt *getFallthroughHandler() const {
-    return SubStmts[SubStmt::OnFallthrough];
+    return getStoredStmts()[SubStmt::OnFallthrough];
   }
 
-  Expr *getAllocate() const { return cast<Expr>(SubStmts[SubStmt::Allocate]); }
-  Stmt *getDeallocate() const { return SubStmts[SubStmt::Deallocate]; }
+  Expr *getAllocate() const {
+    return cast_or_null<Expr>(getStoredStmts()[SubStmt::Allocate]);
+  }
+  Expr *getDeallocate() const {
+    return cast_or_null<Expr>(getStoredStmts()[SubStmt::Deallocate]);
+  }
 
   Expr *getReturnValueInit() const {
-    return cast<Expr>(SubStmts[SubStmt::ReturnValue]);
+    return cast_or_null<Expr>(getStoredStmts()[SubStmt::ReturnValue]);
+  }
+  Stmt *getReturnStmtOnAllocFailure() const {
+    return getStoredStmts()[SubStmt::ReturnStmtOnAllocFailure];
+  }
+  ArrayRef<Stmt const *> getParamMoves() const {
+    return {getStoredStmts() + SubStmt::FirstParamMove, NumParams};
   }
 
   SourceLocation getLocStart() const LLVM_READONLY {
-    return getBody()->getLocStart();
+    return getBody() ? getBody()->getLocStart()
+            : getPromiseDecl()->getLocStart();
   }
   SourceLocation getLocEnd() const LLVM_READONLY {
-    return getBody()->getLocEnd();
+    return getBody() ? getBody()->getLocEnd() : getPromiseDecl()->getLocEnd();
   }
 
   child_range children() {
-    return child_range(SubStmts, SubStmts + SubStmt::FirstParamMove);
+    return child_range(getStoredStmts(),
+                       getStoredStmts() + SubStmt::FirstParamMove + NumParams);
   }
 
   static bool classof(const Stmt *T) {
@@ -390,10 +427,14 @@ class CoreturnStmt : public Stmt {
   enum SubStmt { Operand, PromiseCall, Count };
   Stmt *SubStmts[SubStmt::Count];
 
+  bool IsImplicit : 1;
+
   friend class ASTStmtReader;
 public:
-  CoreturnStmt(SourceLocation CoreturnLoc, Stmt *Operand, Stmt *PromiseCall)
-      : Stmt(CoreturnStmtClass), CoreturnLoc(CoreturnLoc) {
+  CoreturnStmt(SourceLocation CoreturnLoc, Stmt *Operand, Stmt *PromiseCall,
+               bool IsImplicit = false)
+      : Stmt(CoreturnStmtClass), CoreturnLoc(CoreturnLoc),
+        IsImplicit(IsImplicit) {
     SubStmts[SubStmt::Operand] = Operand;
     SubStmts[SubStmt::PromiseCall] = PromiseCall;
   }
@@ -411,6 +452,9 @@ public:
     return static_cast<Expr*>(SubStmts[PromiseCall]);
   }
 
+  bool isImplicit() const { return IsImplicit; }
+  void setIsImplicit(bool value = true) { IsImplicit = value; }
+
   SourceLocation getLocStart() const LLVM_READONLY { return CoreturnLoc; }
   SourceLocation getLocEnd() const LLVM_READONLY {
     return getOperand() ? getOperand()->getLocEnd() : getLocStart();