[OpenACC] Implement 'wait' construct

The arguments to this are the same as for the 'wait' clause, so this
reuses all of that infrastructure. So all this has to do is support a
pair of clauses that are already implemented (if and async), plus create
an AST node.  This patch does so, and adds proper testing.

Author: erichkeane

clang/include/clang-c/Index.h

@@ -2186,7 +2186,11 @@ enum CXCursorKind {
    */
   CXCursor_OpenACCHostDataConstruct = 326,
 
-  CXCursor_LastStmt = CXCursor_OpenACCHostDataConstruct,
+  /** OpenACC wait Construct.
+   */
+  CXCursor_OpenACCWaitConstruct = 327,
+
+  CXCursor_LastStmt = CXCursor_OpenACCWaitConstruct,
 
   /**
    * Cursor that represents the translation unit itself.

clang/include/clang/AST/ASTNodeTraverser.h

@@ -159,7 +159,7 @@ class ASTNodeTraverser
 
       // Some statements have custom mechanisms for dumping their children.
       if (isa<DeclStmt>(S) || isa<GenericSelectionExpr>(S) ||
-          isa<RequiresExpr>(S))
+          isa<RequiresExpr>(S) || isa<OpenACCWaitConstruct>(S))
         return;
 
       if (Traversal == TK_IgnoreUnlessSpelledInSource &&
@@ -825,6 +825,16 @@ class ASTNodeTraverser
       Visit(C);
   }
 
+  void VisitOpenACCWaitConstruct(const OpenACCWaitConstruct *Node) {
+    // Needs custom child checking to put clauses AFTER the children, which are
+    // the expressions in the 'wait' construct. Others likely need this as well,
+    // and might need to do the associated statement after it.
+    for (const Stmt *S : Node->children())
+      Visit(S);
+    for (const auto *C : Node->clauses())
+      Visit(C);
+  }
+
   void VisitInitListExpr(const InitListExpr *ILE) {
     if (auto *Filler = ILE->getArrayFiller()) {
       Visit(Filler, "array_filler");

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -4063,10 +4063,19 @@ DEF_TRAVERSE_STMT(OpenACCCombinedConstruct,
                   { TRY_TO(TraverseOpenACCAssociatedStmtConstruct(S)); })
 DEF_TRAVERSE_STMT(OpenACCDataConstruct,
                   { TRY_TO(TraverseOpenACCAssociatedStmtConstruct(S)); })
-DEF_TRAVERSE_STMT(OpenACCEnterDataConstruct, {})
-DEF_TRAVERSE_STMT(OpenACCExitDataConstruct, {})
+DEF_TRAVERSE_STMT(OpenACCEnterDataConstruct,
+                  { TRY_TO(VisitOpenACCClauseList(S->clauses())); })
+DEF_TRAVERSE_STMT(OpenACCExitDataConstruct,
+                  { TRY_TO(VisitOpenACCClauseList(S->clauses())); })
 DEF_TRAVERSE_STMT(OpenACCHostDataConstruct,
                   { TRY_TO(TraverseOpenACCAssociatedStmtConstruct(S)); })
+DEF_TRAVERSE_STMT(OpenACCWaitConstruct, {
+  if (S->hasDevNumExpr())
+    TRY_TO(TraverseStmt(S->getDevNumExpr()));
+  for (auto *E : S->getQueueIdExprs())
+    TRY_TO(TraverseStmt(E));
+  TRY_TO(VisitOpenACCClauseList(S->clauses()));
+})
 
 // Traverse HLSL: Out argument expression
 DEF_TRAVERSE_STMT(HLSLOutArgExpr, {})

clang/include/clang/AST/StmtOpenACC.h

@@ -469,5 +469,128 @@ class OpenACCHostDataConstruct final
     return const_cast<OpenACCHostDataConstruct *>(this)->getStructuredBlock();
   }
 };
+
+// This class represents a 'wait' construct, which has some expressions plus a
+// clause list.
+class OpenACCWaitConstruct final
+    : public OpenACCConstructStmt,
+      private llvm::TrailingObjects<OpenACCWaitConstruct, Expr *,
+                                    OpenACCClause *> {
+  // FIXME: We should be storing a `const OpenACCClause *` to be consistent with
+  // the rest of the constructs, but TrailingObjects doesn't allow for mixing
+  // constness in its implementation of `getTrailingObjects`.
+
+  friend TrailingObjects;
+  friend class ASTStmtWriter;
+  friend class ASTStmtReader;
+  // Locations of the left and right parens of the 'wait-argument'
+  // expression-list.
+  SourceLocation LParenLoc, RParenLoc;
+  // Location of the 'queues' keyword, if present.
+  SourceLocation QueuesLoc;
+
+  // Number of the expressions being represented.  Index '0' is always the
+  // 'devnum' expression, even if it not present.
+  unsigned NumExprs = 0;
+
+  OpenACCWaitConstruct(unsigned NumExprs, unsigned NumClauses)
+      : OpenACCConstructStmt(OpenACCWaitConstructClass,
+                             OpenACCDirectiveKind::Wait, SourceLocation{},
+                             SourceLocation{}, SourceLocation{}),
+        NumExprs(NumExprs) {
+    assert(NumExprs >= 1 &&
+           "NumExprs should always be >= 1 because the 'devnum' "
+           "expr is represented by a null if necessary");
+    std::uninitialized_value_construct(getExprPtr(),
+                                       getExprPtr() + NumExprs);
+    std::uninitialized_value_construct(getTrailingObjects<OpenACCClause *>(),
+                                       getTrailingObjects<OpenACCClause *>() +
+                                           NumClauses);
+    setClauseList(MutableArrayRef(const_cast<const OpenACCClause **>(
+                                      getTrailingObjects<OpenACCClause *>()),
+                                  NumClauses));
+  }
+
+  OpenACCWaitConstruct(SourceLocation Start, SourceLocation DirectiveLoc,
+                       SourceLocation LParenLoc, Expr *DevNumExpr,
+                       SourceLocation QueuesLoc, ArrayRef<Expr *> QueueIdExprs,
+                       SourceLocation RParenLoc, SourceLocation End,
+                       ArrayRef<const OpenACCClause *> Clauses)
+      : OpenACCConstructStmt(OpenACCWaitConstructClass,
+                             OpenACCDirectiveKind::Wait, Start, DirectiveLoc,
+                             End),
+        LParenLoc(LParenLoc), RParenLoc(RParenLoc), QueuesLoc(QueuesLoc),
+        NumExprs(QueueIdExprs.size() + 1) {
+    assert(NumExprs >= 1 &&
+           "NumExprs should always be >= 1 because the 'devnum' "
+           "expr is represented by a null if necessary");
+
+    std::uninitialized_copy(&DevNumExpr, &DevNumExpr + 1,
+                            getExprPtr());
+    std::uninitialized_copy(QueueIdExprs.begin(), QueueIdExprs.end(),
+                            getExprPtr() + 1);
+
+    std::uninitialized_copy(const_cast<OpenACCClause **>(Clauses.begin()),
+                            const_cast<OpenACCClause **>(Clauses.end()),
+                            getTrailingObjects<OpenACCClause *>());
+    setClauseList(MutableArrayRef(const_cast<const OpenACCClause **>(
+                                      getTrailingObjects<OpenACCClause *>()),
+                                  Clauses.size()));
+  }
+
+  size_t numTrailingObjects(OverloadToken<Expr *>) const { return NumExprs; }
+  size_t numTrailingObjects(OverloadToken<const OpenACCClause *>) const {
+    return clauses().size();
+  }
+
+  Expr **getExprPtr() const {
+    return const_cast<Expr**>(getTrailingObjects<Expr *>());
+  }
+
+  llvm::ArrayRef<Expr *> getExprs() const {
+    return llvm::ArrayRef<Expr *>(getExprPtr(), NumExprs);
+  }
+
+  llvm::ArrayRef<Expr *> getExprs() {
+    return llvm::ArrayRef<Expr *>(getExprPtr(), NumExprs);
+  }
+
+public:
+  static bool classof(const Stmt *T) {
+    return T->getStmtClass() == OpenACCWaitConstructClass;
+  }
+
+  static OpenACCWaitConstruct *
+  CreateEmpty(const ASTContext &C, unsigned NumExprs, unsigned NumClauses);
+
+  static OpenACCWaitConstruct *
+  Create(const ASTContext &C, SourceLocation Start, SourceLocation DirectiveLoc,
+         SourceLocation LParenLoc, Expr *DevNumExpr, SourceLocation QueuesLoc,
+         ArrayRef<Expr *> QueueIdExprs, SourceLocation RParenLoc,
+         SourceLocation End, ArrayRef<const OpenACCClause *> Clauses);
+
+  SourceLocation getLParenLoc() const { return LParenLoc; }
+  SourceLocation getRParenLoc() const { return RParenLoc; }
+  bool hasQueuesTag() const { return !QueuesLoc.isInvalid(); }
+  SourceLocation getQueuesLoc() const { return QueuesLoc; }
+
+  bool hasDevNumExpr() const { return getExprs()[0]; }
+  Expr *getDevNumExpr() const { return getExprs()[0]; }
+  llvm::ArrayRef<Expr *> getQueueIdExprs() { return getExprs().drop_front(); }
+  llvm::ArrayRef<Expr *> getQueueIdExprs() const {
+    return getExprs().drop_front();
+  }
+
+  child_range children() {
+    Stmt **Begin = reinterpret_cast<Stmt **>(getExprPtr());
+    return child_range(Begin, Begin + NumExprs);
+  }
+
+  const_child_range children() const {
+    Stmt *const *Begin =
+        reinterpret_cast<Stmt *const *>(getExprPtr());
+    return const_child_range(Begin, Begin + NumExprs);
+  }
+};
 } // namespace clang
 #endif // LLVM_CLANG_AST_STMTOPENACC_H

clang/include/clang/AST/TextNodeDumper.h

@@ -415,6 +415,7 @@ class TextNodeDumper
   void VisitOpenACCEnterDataConstruct(const OpenACCEnterDataConstruct *S);
   void VisitOpenACCExitDataConstruct(const OpenACCExitDataConstruct *S);
   void VisitOpenACCHostDataConstruct(const OpenACCHostDataConstruct *S);
+  void VisitOpenACCWaitConstruct(const OpenACCWaitConstruct *S);
   void VisitOpenACCAsteriskSizeExpr(const OpenACCAsteriskSizeExpr *S);
   void VisitEmbedExpr(const EmbedExpr *S);
   void VisitAtomicExpr(const AtomicExpr *AE);

clang/include/clang/Basic/StmtNodes.td

@@ -312,6 +312,7 @@ def OpenACCDataConstruct : StmtNode<OpenACCAssociatedStmtConstruct>;
 def OpenACCEnterDataConstruct : StmtNode<OpenACCConstructStmt>;
 def OpenACCExitDataConstruct : StmtNode<OpenACCConstructStmt>;
 def OpenACCHostDataConstruct : StmtNode<OpenACCAssociatedStmtConstruct>;
+def OpenACCWaitConstruct : StmtNode<OpenACCConstructStmt>;
 
 // OpenACC Additional Expressions.
 def OpenACCAsteriskSizeExpr : StmtNode<Expr>;

clang/include/clang/Parse/Parser.h

@@ -3706,17 +3706,28 @@ class Parser : public CodeCompletionHandler {
     OpenACCDirectiveKind DirKind;
     SourceLocation StartLoc;
     SourceLocation DirLoc;
+    SourceLocation LParenLoc;
+    SourceLocation RParenLoc;
     SourceLocation EndLoc;
+    SourceLocation MiscLoc;
+    SmallVector<Expr *> Exprs;
     SmallVector<OpenACCClause *> Clauses;
-    // TODO OpenACC: As we implement support for the Atomic, Routine, Cache, and
-    // Wait constructs, we likely want to put that information in here as well.
+    // TODO OpenACC: As we implement support for the Atomic, Routine, and Cache
+    // constructs, we likely want to put that information in here as well.
   };
 
   struct OpenACCWaitParseInfo {
     bool Failed = false;
     Expr *DevNumExpr = nullptr;
     SourceLocation QueuesLoc;
     SmallVector<Expr *> QueueIdExprs;
+
+    SmallVector<Expr *> getAllExprs() {
+      SmallVector<Expr *> Out;
+      Out.push_back(DevNumExpr);
+      Out.insert(Out.end(), QueueIdExprs.begin(), QueueIdExprs.end());
+      return Out;
+    }
   };
 
   /// Represents the 'error' state of parsing an OpenACC Clause, and stores

clang/include/clang/Sema/SemaOpenACC.h

@@ -679,12 +679,18 @@ class SemaOpenACC : public SemaBase {
 
   /// Called after the directive has been completely parsed, including the
   /// declaration group or associated statement.
-  StmtResult ActOnEndStmtDirective(OpenACCDirectiveKind K,
-                                   SourceLocation StartLoc,
-                                   SourceLocation DirLoc,
-                                   SourceLocation EndLoc,
-                                   ArrayRef<OpenACCClause *> Clauses,
-                                   StmtResult AssocStmt);
+  /// LParenLoc: Location of the left paren, if it exists (not on all
+  /// constructs).
+  /// MiscLoc: First misc location, if necessary (not all constructs).
+  /// Exprs: List of expressions on the construct itself, if necessary (not all
+  /// constructs).
+  /// RParenLoc: Location of the right paren, if it exists (not on all
+  /// constructs).
+  StmtResult ActOnEndStmtDirective(
+      OpenACCDirectiveKind K, SourceLocation StartLoc, SourceLocation DirLoc,
+      SourceLocation LParenLoc, SourceLocation MiscLoc, ArrayRef<Expr *> Exprs,
+      SourceLocation RParenLoc, SourceLocation EndLoc,
+      ArrayRef<OpenACCClause *> Clauses, StmtResult AssocStmt);
 
   /// Called after the directive has been completely parsed, including the
   /// declaration group or associated statement.

clang/include/clang/Serialization/ASTBitCodes.h

@@ -2021,6 +2021,7 @@ enum StmtCode {
   STMT_OPENACC_ENTER_DATA_CONSTRUCT,
   STMT_OPENACC_EXIT_DATA_CONSTRUCT,
   STMT_OPENACC_HOST_DATA_CONSTRUCT,
+  STMT_OPENACC_WAIT_CONSTRUCT,
 
   // HLSL Constructs
   EXPR_HLSL_OUT_ARG,

clang/lib/AST/StmtOpenACC.cpp

@@ -196,3 +196,32 @@ OpenACCHostDataConstruct *OpenACCHostDataConstruct::Create(
                                                   Clauses, StructuredBlock);
   return Inst;
 }
+
+OpenACCWaitConstruct *OpenACCWaitConstruct::CreateEmpty(const ASTContext &C,
+                                                        unsigned NumExprs,
+                                                        unsigned NumClauses) {
+  void *Mem = C.Allocate(
+      OpenACCWaitConstruct::totalSizeToAlloc<Expr *, OpenACCClause *>(
+          NumExprs, NumClauses));
+
+  auto *Inst = new (Mem) OpenACCWaitConstruct(NumExprs, NumClauses);
+  return Inst;
+}
+
+OpenACCWaitConstruct *OpenACCWaitConstruct::Create(
+    const ASTContext &C, SourceLocation Start, SourceLocation DirectiveLoc,
+    SourceLocation LParenLoc, Expr *DevNumExpr, SourceLocation QueuesLoc,
+    ArrayRef<Expr *> QueueIdExprs, SourceLocation RParenLoc, SourceLocation End,
+    ArrayRef<const OpenACCClause *> Clauses) {
+
+  assert(llvm::all_of(QueueIdExprs, [](Expr *E) { return E != nullptr; }));
+
+  void *Mem = C.Allocate(
+      OpenACCWaitConstruct::totalSizeToAlloc<Expr *, OpenACCClause *>(
+          QueueIdExprs.size() + 1, Clauses.size()));
+
+  auto *Inst = new (Mem)
+      OpenACCWaitConstruct(Start, DirectiveLoc, LParenLoc, DevNumExpr,
+                           QueuesLoc, QueueIdExprs, RParenLoc, End, Clauses);
+  return Inst;
+}