Implement CapturedStmt AST

CapturedStmt can be used to implement generic function outlining as described in
http://lists.cs.uiuc.edu/pipermail/cfe-dev/2013-January/027540.html.

CapturedStmt is not exposed to the C api.

Serialization and template support are pending.

Author: Wei Pan <wei.pan@intel.com>

Differential Revision: http://llvm-reviews.chandlerc.com/D370

llvm-svn: 179615

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -2218,6 +2218,7 @@ DEF_TRAVERSE_STMT(UnresolvedMemberExpr, {
 DEF_TRAVERSE_STMT(SEHTryStmt, {})
 DEF_TRAVERSE_STMT(SEHExceptStmt, {})
 DEF_TRAVERSE_STMT(SEHFinallyStmt,{})
+DEF_TRAVERSE_STMT(CapturedStmt, {})
 
 DEF_TRAVERSE_STMT(CXXOperatorCallExpr, { })
 DEF_TRAVERSE_STMT(OpaqueValueExpr, { })

clang/include/clang/AST/Stmt.h

@@ -33,12 +33,14 @@ namespace clang {
   class Attr;
   class Decl;
   class Expr;
+  class FunctionDecl;
   class IdentifierInfo;
   class LabelDecl;
   class ParmVarDecl;
   class PrinterHelper;
   struct PrintingPolicy;
   class QualType;
+  class RecordDecl;
   class SourceManager;
   class StringLiteral;
   class SwitchStmt;
@@ -1882,6 +1884,170 @@ public:
   }
 };
 
+/// \brief This captures a statement into a function. For example, the following
+/// pragma annotated compound statement can be represented as a CapturedStmt,
+/// and this compound statement is the body of an anonymous outlined function.
+/// @code
+/// #pragma omp parallel
+/// {
+///   compute();
+/// }
+/// @endcode
+class CapturedStmt : public Stmt {
+public:
+  /// \brief The different capture forms: by 'this' or by reference, etc.
+  enum VariableCaptureKind {
+    VCK_This,
+    VCK_ByRef
+  };
+
+  /// \brief Describes the capture of either a variable or 'this'.
+  class Capture {
+    VarDecl *Var;
+    SourceLocation Loc;
+
+  public:
+    /// \brief Create a new capture.
+    ///
+    /// \param Loc The source location associated with this capture.
+    ///
+    /// \param Kind The kind of capture (this, ByRef, ...).
+    ///
+    /// \param Var The variable being captured, or null if capturing this.
+    ///
+    Capture(SourceLocation Loc, VariableCaptureKind Kind, VarDecl *Var = 0)
+      : Var(Var), Loc(Loc) {
+      switch (Kind) {
+      case VCK_This:
+        assert(Var == 0 && "'this' capture cannot have a variable!");
+        break;
+      case VCK_ByRef:
+        assert(Var && "capturing by reference must have a variable!");
+        break;
+      }
+    }
+
+    /// \brief Determine the kind of capture.
+    VariableCaptureKind getCaptureKind() const {
+      if (capturesThis())
+        return VCK_This;
+
+      return VCK_ByRef;
+    }
+
+    /// \brief Retrieve the source location at which the variable or 'this' was
+    /// first used.
+    SourceLocation getLocation() const { return Loc; }
+
+    /// \brief Determine whether this capture handles the C++ 'this' pointer.
+    bool capturesThis() const { return Var == 0; }
+
+    /// \brief Determine whether this capture handles a variable.
+    bool capturesVariable() const { return Var != 0; }
+
+    /// \brief Retrieve the declaration of the variable being captured.
+    ///
+    /// This operation is only valid if this capture does not capture 'this'.
+    VarDecl *getCapturedVar() const {
+      assert(!capturesThis() && "No variable available for 'this' capture");
+      return Var;
+    }
+  };
+
+private:
+  /// \brief The number of variable captured, including 'this'.
+  unsigned NumCaptures;
+
+  /// \brief The implicit outlined function.
+  FunctionDecl *TheFuncDecl;
+
+  /// \brief The record for captured variables, a RecordDecl or CXXRecordDecl.
+  RecordDecl *TheRecordDecl;
+
+  /// \brief Construct a captured statement.
+  CapturedStmt(Stmt *S, ArrayRef<Capture> Captures,
+               ArrayRef<Expr *> CaptureInits,
+               FunctionDecl *FD, RecordDecl *RD);
+
+  /// \brief Construct an empty captured statement.
+  CapturedStmt(EmptyShell Empty, unsigned NumCaptures);
+
+  Stmt **getStoredStmts() const {
+    return reinterpret_cast<Stmt **>(const_cast<CapturedStmt *>(this) + 1);
+  }
+
+  Capture *getStoredCaptures() const;
+
+public:
+  static CapturedStmt *Create(ASTContext &Context, Stmt *S,
+                              ArrayRef<Capture> Captures,
+                              ArrayRef<Expr *> CaptureInits,
+                              FunctionDecl *FD, RecordDecl *RD);
+
+  static CapturedStmt *CreateDeserialized(ASTContext &Context,
+                                          unsigned NumCaptures);
+
+  /// \brief Retrieve the statement being captured.
+  Stmt *getCapturedStmt() { return getStoredStmts()[NumCaptures]; }
+  const Stmt *getCapturedStmt() const {
+    return const_cast<CapturedStmt *>(this)->getCapturedStmt();
+  }
+
+  /// \brief Retrieve the outlined function declaration.
+  const FunctionDecl *getCapturedFunctionDecl() const { return TheFuncDecl; }
+
+  /// \brief Retrieve the record declaration for captured variables.
+  const RecordDecl *getCapturedRecordDecl() const { return TheRecordDecl; }
+
+  /// \brief True if this variable has been captured.
+  bool capturesVariable(const VarDecl *Var) const;
+
+  /// \brief An iterator that walks over the captures.
+  typedef const Capture *capture_iterator;
+
+  /// \brief Retrieve an iterator pointing to the first capture.
+  capture_iterator capture_begin() const { return getStoredCaptures(); }
+
+  /// \brief Retrieve an iterator pointing past the end of the sequence of
+  /// captures.
+  capture_iterator capture_end() const {
+    return getStoredCaptures() + NumCaptures;
+  }
+
+  /// \brief Retrieve the number of captures, including 'this'.
+  unsigned capture_size() const { return NumCaptures; }
+
+  /// \brief Iterator that walks over the capture initialization arguments.
+  typedef Expr **capture_init_iterator;
+
+  /// \brief Retrieve the first initialization argument.
+  capture_init_iterator capture_init_begin() const {
+    return reinterpret_cast<Expr **>(getStoredStmts());
+  }
+
+  /// \brief Retrieve the iterator pointing one past the last initialization
+  /// argument.
+  capture_init_iterator capture_init_end() const {
+    return capture_init_begin() + NumCaptures;
+  }
+
+  SourceLocation getLocStart() const LLVM_READONLY {
+    return getCapturedStmt()->getLocStart();
+  }
+  SourceLocation getLocEnd() const LLVM_READONLY {
+    return getCapturedStmt()->getLocEnd();
+  }
+  SourceRange getSourceRange() const LLVM_READONLY {
+    return getCapturedStmt()->getSourceRange();
+  }
+
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == CapturedStmtClass;
+  }
+
+  child_range children();
+};
+
 }  // end namespace clang
 
 #endif

clang/include/clang/Basic/StmtNodes.td

@@ -27,6 +27,7 @@ def DeclStmt  : Stmt;
 def SwitchCase : Stmt<1>;
 def CaseStmt : DStmt<SwitchCase>;
 def DefaultStmt : DStmt<SwitchCase>;
+def CapturedStmt : Stmt;
 
 // Asm statements
 def AsmStmt : Stmt<1>;

clang/include/clang/Serialization/ASTBitCodes.h

@@ -1103,6 +1103,8 @@ namespace clang {
       STMT_RETURN,
       /// \brief A DeclStmt record.
       STMT_DECL,
+      /// \brief A CapturedStmt record.
+      STMT_CAPTURED,
       /// \brief A GCC-style AsmStmt record.
       STMT_GCCASM,
       /// \brief A MS-style AsmStmt record.

clang/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;
+}

clang/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())) {

clang/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);
 }

clang/lib/CodeGen/CGStmt.cpp

@@ -134,7 +134,9 @@ void CodeGenFunction::EmitStmt(const Stmt *S) {
   case Stmt::SwitchStmtClass:   EmitSwitchStmt(cast<SwitchStmt>(*S));     break;
   case Stmt::GCCAsmStmtClass:   // Intentional fall-through.
   case Stmt::MSAsmStmtClass:    EmitAsmStmt(cast<AsmStmt>(*S));           break;
-
+  case Stmt::CapturedStmtClass:
+    EmitCapturedStmt(cast<CapturedStmt>(*S));
+    break;
   case Stmt::ObjCAtTryStmtClass:
     EmitObjCAtTryStmt(cast<ObjCAtTryStmt>(*S));
     break;
@@ -1735,3 +1737,7 @@ void CodeGenFunction::EmitAsmStmt(const AsmStmt &S) {
     EmitStoreThroughLValue(RValue::get(Tmp), ResultRegDests[i]);
   }
 }
+
+void CodeGenFunction::EmitCapturedStmt(const CapturedStmt &S) {
+  llvm_unreachable("not implemented yet");
+}

clang/lib/CodeGen/CodeGenFunction.h

@@ -2133,6 +2133,7 @@ public:
   void EmitCaseStmt(const CaseStmt &S);
   void EmitCaseStmtRange(const CaseStmt &S);
   void EmitAsmStmt(const AsmStmt &S);
+  void EmitCapturedStmt(const CapturedStmt &S);
 
   void EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S);
   void EmitObjCAtTryStmt(const ObjCAtTryStmt &S);

clang/lib/Sema/TreeTransform.h

@@ -9377,6 +9377,12 @@ TreeTransform<Derived>::RebuildCXXPseudoDestructorExpr(Expr *Base,
                                             /*TemplateArgs*/ 0);
 }
 
+template<typename Derived>
+StmtResult
+TreeTransform<Derived>::TransformCapturedStmt(CapturedStmt *S) {
+  llvm_unreachable("not implement yet");
+}
+
 } // end namespace clang
 
 #endif // LLVM_CLANG_SEMA_TREETRANSFORM_H

clang/lib/Serialization/ASTReaderStmt.cpp

@@ -324,6 +324,10 @@ void ASTStmtReader::VisitMSAsmStmt(MSAsmStmt *S) {
   VisitStmt(S);
 }
 
+void ASTStmtReader::VisitCapturedStmt(CapturedStmt *S) {
+  llvm_unreachable("not implemented yet");
+}
+
 void ASTStmtReader::VisitExpr(Expr *E) {
   VisitStmt(E);
   E->setType(Reader.readType(F, Record, Idx));
@@ -1724,6 +1728,10 @@ Stmt *ASTReader::ReadStmtFromStream(ModuleFile &F) {
       S = new (Context) MSAsmStmt(Empty);
       break;
 
+    case STMT_CAPTURED:
+      llvm_unreachable("not implemented yet");
+      break;
+
     case EXPR_PREDEFINED:
       S = new (Context) PredefinedExpr(Empty);
       break;

clang/lib/Serialization/ASTWriterStmt.cpp

@@ -255,6 +255,13 @@ void ASTStmtWriter::VisitMSAsmStmt(MSAsmStmt *S) {
   Code = serialization::STMT_MSASM;
 }
 
+void ASTStmtWriter::VisitCapturedStmt(CapturedStmt *S) {
+  VisitStmt(S);
+  Code = serialization::STMT_CAPTURED;
+
+  llvm_unreachable("not implemented yet");
+}
+
 void ASTStmtWriter::VisitExpr(Expr *E) {
   VisitStmt(E);
   Writer.AddTypeRef(E->getType(), Record);

clang/lib/StaticAnalyzer/Core/ExprEngine.cpp

@@ -656,6 +656,7 @@ void ExprEngine::Visit(const Stmt *S, ExplodedNode *Pred,
     case Stmt::SwitchStmtClass:
     case Stmt::WhileStmtClass:
     case Expr::MSDependentExistsStmtClass:
+    case Stmt::CapturedStmtClass:
       llvm_unreachable("Stmt should not be in analyzer evaluation loop");
 
     case Stmt::ObjCSubscriptRefExprClass:

clang/tools/libclang/CXCursor.cpp

@@ -270,6 +270,10 @@ CXCursor cxcursor::MakeCXCursor(const Stmt *S, const Decl *Parent,
     K = CXCursor_DeclStmt;
     break;
 
+  case Stmt::CapturedStmtClass:
+    K = CXCursor_UnexposedStmt;
+    break;
+
   case Stmt::IntegerLiteralClass:
     K = CXCursor_IntegerLiteral;
     break;

clang/tools/libclang/RecursiveASTVisitor.h

@@ -1839,7 +1839,7 @@ DEF_TRAVERSE_STMT(MSDependentExistsStmt, {
 DEF_TRAVERSE_STMT(ReturnStmt, { })
 DEF_TRAVERSE_STMT(SwitchStmt, { })
 DEF_TRAVERSE_STMT(WhileStmt, { })
-
+DEF_TRAVERSE_STMT(CapturedStmt, { })
 
 DEF_TRAVERSE_STMT(CXXDependentScopeMemberExpr, {
     TRY_TO(TraverseNestedNameSpecifierLoc(S->getQualifierLoc()));