3 files changed, 110 insertions, 0 deletions

diff --git a/lib/AST/Stmt.cpp b/lib/AST/Stmt.cpp
index 2ae5a1266c..e120c6a1f8 100644
--- a/lib/AST/Stmt.cpp
+++ b/lib/AST/Stmt.cpp
@@ -1023,3 +1023,105 @@ SEHFinallyStmt* SEHFinallyStmt::Create(ASTContext &C,
                                        Stmt *Block) {
   return new(C)SEHFinallyStmt(Loc,Block);
 }
+
+CapturedStmt::Capture *CapturedStmt::getStoredCaptures() const {
+  unsigned Size = sizeof(CapturedStmt) + sizeof(Stmt *) * (NumCaptures + 1);
+
+  // Offset of the first Capture object.
+  unsigned FirstCaptureOffset =
+    llvm::RoundUpToAlignment(Size, llvm::alignOf<Capture>());
+
+  return reinterpret_cast<Capture *>(
+      reinterpret_cast<char *>(const_cast<CapturedStmt *>(this))
+      + FirstCaptureOffset);
+}
+
+CapturedStmt::CapturedStmt(Stmt *S, ArrayRef<Capture> Captures,
+                           ArrayRef<Expr *> CaptureInits,
+                           FunctionDecl *FD,
+                           RecordDecl *RD)
+  : Stmt(CapturedStmtClass), NumCaptures(Captures.size()),
+    TheFuncDecl(FD), TheRecordDecl(RD) {
+  assert( S && "null captured statement");
+  assert(FD && "null function declaration for captured statement");
+  assert(RD && "null record declaration for captured statement");
+
+  // Copy initialization expressions.
+  Stmt **Stored = getStoredStmts();
+  for (unsigned I = 0, N = NumCaptures; I != N; ++I)
+    *Stored++ = CaptureInits[I];
+
+  // Copy the statement being captured.
+  *Stored = S;
+
+  // Copy all Capture objects.
+  Capture *Buffer = getStoredCaptures();
+  std::copy(Captures.begin(), Captures.end(), Buffer);
+}
+
+CapturedStmt::CapturedStmt(EmptyShell Empty, unsigned NumCaptures)
+  : Stmt(CapturedStmtClass, Empty), NumCaptures(NumCaptures),
+    TheFuncDecl(0), TheRecordDecl(0) {
+  getStoredStmts()[NumCaptures] = 0;
+}
+
+CapturedStmt *CapturedStmt::Create(ASTContext &Context, Stmt *S,
+                                   ArrayRef<Capture> Captures,
+                                   ArrayRef<Expr *> CaptureInits,
+                                   FunctionDecl *FD,
+                                   RecordDecl *RD) {
+  // The layout is
+  //
+  // -----------------------------------------------------------
+  // | CapturedStmt, Init, ..., Init, S, Capture, ..., Capture |
+  // ----------------^-------------------^----------------------
+  //                 getStoredStmts()    getStoredCaptures()
+  //
+  // where S is the statement being captured.
+  //
+  assert(CaptureInits.size() == Captures.size() && "wrong number of arguments");
+
+  unsigned Size = sizeof(CapturedStmt) + sizeof(Stmt *) * (Captures.size() + 1);
+  if (!Captures.empty()) {
+    // Realign for the following Capture array.
+    Size = llvm::RoundUpToAlignment(Size, llvm::alignOf<Capture>());
+    Size += sizeof(Capture) * Captures.size();
+  }
+
+  void *Mem = Context.Allocate(Size);
+  return new (Mem) CapturedStmt(S, Captures, CaptureInits, FD, RD);
+}
+
+CapturedStmt *CapturedStmt::CreateDeserialized(ASTContext &Context,
+                                               unsigned NumCaptures) {
+  unsigned Size = sizeof(CapturedStmt) + sizeof(Stmt *) * (NumCaptures + 1);
+  if (NumCaptures > 0) {
+    // Realign for the following Capture array.
+    Size = llvm::RoundUpToAlignment(Size, llvm::alignOf<Capture>());
+    Size += sizeof(Capture) * NumCaptures;
+  }
+
+  void *Mem = Context.Allocate(Size);
+  return new (Mem) CapturedStmt(EmptyShell(), NumCaptures);
+}
+
+Stmt::child_range CapturedStmt::children() {
+  // Children are captured field initilizers and the statement being captured.
+  return child_range(getStoredStmts(), getStoredStmts() + NumCaptures + 1);
+}
+
+bool CapturedStmt::capturesVariable(const VarDecl *Var) const {
+  for (capture_iterator I = capture_begin(),
+                        E = capture_end(); I != E; ++I) {
+    if (I->capturesThis())
+      continue;
+
+    // This does not handle variable redeclarations. This should be
+    // extended to capture variables with redeclarations, for example
+    // a thread-private variable in OpenMP.
+    if (I->getCapturedVar() == Var)
+      return true;
+  }
+
+  return false;
+}
diff --git a/lib/AST/StmtPrinter.cpp b/lib/AST/StmtPrinter.cpp
index a86159f49d..469c2846a6 100644
--- a/lib/AST/StmtPrinter.cpp
+++ b/lib/AST/StmtPrinter.cpp
@@ -450,6 +450,10 @@ void StmtPrinter::VisitMSAsmStmt(MSAsmStmt *Node) {
     Indent() << "}\n";
 }
 
+void StmtPrinter::VisitCapturedStmt(CapturedStmt *Node) {
+  PrintStmt(Node->getCapturedStmt());
+}
+
 void StmtPrinter::VisitObjCAtTryStmt(ObjCAtTryStmt *Node) {
   Indent() << "@try";
   if (CompoundStmt *TS = dyn_cast<CompoundStmt>(Node->getTryBody())) {
diff --git a/lib/AST/StmtProfile.cpp b/lib/AST/StmtProfile.cpp
index 5525018f79..d99400c603 100644
--- a/lib/AST/StmtProfile.cpp
+++ b/lib/AST/StmtProfile.cpp
@@ -215,6 +215,10 @@ void StmtProfiler::VisitSEHExceptStmt(const SEHExceptStmt *S) {
   VisitStmt(S);
 }
 
+void StmtProfiler::VisitCapturedStmt(const CapturedStmt *S) {
+  VisitStmt(S);
+}
+
 void StmtProfiler::VisitObjCForCollectionStmt(const ObjCForCollectionStmt *S) {
   VisitStmt(S);
 }