[flang][OpenMP] Convert AST node for ALLOCATORS to use Block as body (#148005)

The ALLOCATORS construct is one of the few constructs that require a
special form of the associated block.
Convert the AST node to use OmpDirectiveSpecification for the directive
and the optional end directive, and to use parser::Block as the body:
the form of the block is checked in the semantic checks (with a more
meaningful message).

Author: kparzysz

flang/examples/FeatureList/FeatureList.cpp

@@ -472,7 +472,6 @@ struct NodeVisitor {
   READ_FEATURE(OmpIteration)
   READ_FEATURE(OmpIterationOffset)
   READ_FEATURE(OmpIterationVector)
-  READ_FEATURE(OmpEndAllocators)
   READ_FEATURE(OmpEndBlockDirective)
   READ_FEATURE(OmpEndCriticalDirective)
   READ_FEATURE(OmpEndLoopDirective)

flang/include/flang/Parser/dump-parse-tree.h

@@ -578,7 +578,6 @@ class ParseTreeDumper {
   NODE(parser, OmpDetachClause)
   NODE(parser, OmpDoacrossClause)
   NODE(parser, OmpDestroyClause)
-  NODE(parser, OmpEndAllocators)
   NODE(parser, OmpEndBlockDirective)
   NODE(parser, OmpEndCriticalDirective)
   NODE(parser, OmpEndLoopDirective)

flang/include/flang/Parser/parse-tree.h

@@ -4598,8 +4598,11 @@ struct OmpClauseList {
 struct OmpDirectiveSpecification {
   ENUM_CLASS(Flags, None, DeprecatedSyntax);
   TUPLE_CLASS_BOILERPLATE(OmpDirectiveSpecification);
+  const OmpDirectiveName &DirName() const {
+    return std::get<OmpDirectiveName>(t);
+  }
   llvm::omp::Directive DirId() const { //
-    return std::get<OmpDirectiveName>(t).v;
+    return DirName().v;
   }
   const OmpArgumentList &Arguments() const;
   const OmpClauseList &Clauses() const;
@@ -4839,17 +4842,17 @@ struct OpenMPExecutableAllocate {
       t;
 };
 
-EMPTY_CLASS(OmpEndAllocators);
-
-// 6.7 Allocators construct [OpenMP 5.2]
-//     allocators-construct -> ALLOCATORS [allocate-clause [,]]
-//                                allocate-stmt
-//                             [omp-end-allocators-construct]
+// Ref: [5.2:180-181], [6.0:315]
+//
+// allocators-construct ->
+//    ALLOCATORS [allocate-clause...]
+//    block
+//    [END ALLOCATORS]
 struct OpenMPAllocatorsConstruct {
   TUPLE_CLASS_BOILERPLATE(OpenMPAllocatorsConstruct);
   CharBlock source;
-  std::tuple<Verbatim, OmpClauseList, Statement<AllocateStmt>,
-      std::optional<OmpEndAllocators>>
+  std::tuple<OmpDirectiveSpecification, Block,
+      std::optional<OmpDirectiveSpecification>>
       t;
 };
 

flang/lib/Parser/openmp-parsers.cpp

@@ -1276,6 +1276,32 @@ struct OmpEndDirectiveParser {
   llvm::omp::Directive dir_;
 };
 
+struct OmpAllocatorsConstructParser {
+  using resultType = OpenMPAllocatorsConstruct;
+
+  std::optional<resultType> Parse(ParseState &state) const {
+    auto dirSpec{Parser<OmpDirectiveSpecification>{}.Parse(state)};
+    if (!dirSpec || dirSpec->DirId() != llvm::omp::Directive::OMPD_allocators) {
+      return std::nullopt;
+    }
+
+    // This should be an allocate-stmt. That will be checked in semantics.
+    Block block;
+    if (auto stmt{attempt(Parser<ExecutionPartConstruct>{}).Parse(state)}) {
+      block.emplace_back(std::move(*stmt));
+    }
+    // Allow empty block. Check for this in semantics.
+
+    auto end{OmpEndDirectiveParser{llvm::omp::Directive::OMPD_allocators}};
+    return OpenMPAllocatorsConstruct{
+        std::move(*dirSpec), std::move(block), *maybe(end).Parse(state)};
+  }
+};
+
+TYPE_PARSER(sourced( //
+    construct<OpenMPAllocatorsConstruct>(
+        "ALLOCATORS"_tok >= OmpAllocatorsConstructParser{})))
+
 // Parser for an arbitrary OpenMP ATOMIC construct.
 //
 // Depending on circumstances, an ATOMIC construct applies to one or more
@@ -1622,16 +1648,6 @@ TYPE_PARSER(
         maybe(nonemptyList(Parser<OpenMPDeclarativeAllocate>{})) / endOmpLine,
         statement(allocateStmt))))
 
-// 6.7 Allocators construct [OpenMP 5.2]
-//     allocators-construct -> ALLOCATORS [allocate-clause [,]]
-//                                allocate-stmt
-//                             [omp-end-allocators-construct]
-TYPE_PARSER(sourced(construct<OpenMPAllocatorsConstruct>(
-    verbatim("ALLOCATORS"_tok), Parser<OmpClauseList>{} / endOmpLine,
-    statement(allocateStmt), maybe(Parser<OmpEndAllocators>{} / endOmpLine))))
-
-TYPE_PARSER(construct<OmpEndAllocators>(startOmpLine >> "END ALLOCATORS"_tok))
-
 // 2.8.2 Declare Simd construct
 TYPE_PARSER(sourced(construct<OpenMPDeclareSimdConstruct>(
     verbatim("DECLARE SIMD"_tok) || verbatim("DECLARE_SIMD"_tok),

flang/lib/Parser/unparse.cpp

@@ -2571,7 +2571,7 @@ class UnparseVisitor {
     Word(ToUpperCaseLetters(common::EnumToString(x)));
   }
 
-  void Unparse(const OpenMPAtomicConstruct &x) {
+  template <typename Construct> void UnparseBlockConstruct(const Construct &x) {
     BeginOpenMP();
     Word("!$OMP ");
     Walk(std::get<OmpDirectiveSpecification>(x.t));
@@ -2587,6 +2587,10 @@ class UnparseVisitor {
     }
   }
 
+  void Unparse(const OpenMPAtomicConstruct &x) { //
+    UnparseBlockConstruct(x);
+  }
+
   void Unparse(const OpenMPExecutableAllocate &x) {
     const auto &fields =
         std::get<std::optional<std::list<parser::OpenMPDeclarativeAllocate>>>(
@@ -2614,22 +2618,8 @@ class UnparseVisitor {
     Put("\n");
     EndOpenMP();
   }
-  void Unparse(const OmpEndAllocators &x) {
-    BeginOpenMP();
-    Word("!$OMP END ALLOCATE");
-    Put("\n");
-    EndOpenMP();
-  }
-  void Unparse(const OpenMPAllocatorsConstruct &x) {
-    BeginOpenMP();
-    Word("!$OMP ALLOCATE");
-    Walk(std::get<OmpClauseList>(x.t));
-    Put("\n");
-    EndOpenMP();
-    Walk(std::get<Statement<AllocateStmt>>(x.t));
-    if (const auto &end = std::get<std::optional<OmpEndAllocators>>(x.t)) {
-      Walk(*end);
-    }
+  void Unparse(const OpenMPAllocatorsConstruct &x) { //
+    UnparseBlockConstruct(x);
   }
   void Unparse(const OmpAssumeDirective &x) {
     BeginOpenMP();

flang/lib/Semantics/check-omp-atomic.cpp

@@ -47,43 +47,12 @@ static bool operator!=(const evaluate::Expr<T> &e, const evaluate::Expr<U> &f) {
   return !(e == f);
 }
 
-// There is no consistent way to get the source of a given ActionStmt, so
-// extract the source information from Statement<ActionStmt> when we can,
-// and keep it around for error reporting in further analyses.
-struct SourcedActionStmt {
-  const parser::ActionStmt *stmt{nullptr};
-  parser::CharBlock source;
-
-  operator bool() const { return stmt != nullptr; }
-};
-
 struct AnalyzedCondStmt {
   SomeExpr cond{evaluate::NullPointer{}}; // Default ctor is deleted
   parser::CharBlock source;
   SourcedActionStmt ift, iff;
 };
 
-static SourcedActionStmt GetActionStmt(
-    const parser::ExecutionPartConstruct *x) {
-  if (x == nullptr) {
-    return SourcedActionStmt{};
-  }
-  if (auto *exec{std::get_if<parser::ExecutableConstruct>(&x->u)}) {
-    using ActionStmt = parser::Statement<parser::ActionStmt>;
-    if (auto *stmt{std::get_if<ActionStmt>(&exec->u)}) {
-      return SourcedActionStmt{&stmt->statement, stmt->source};
-    }
-  }
-  return SourcedActionStmt{};
-}
-
-static SourcedActionStmt GetActionStmt(const parser::Block &block) {
-  if (block.size() == 1) {
-    return GetActionStmt(&block.front());
-  }
-  return SourcedActionStmt{};
-}
-
 // Compute the `evaluate::Assignment` from parser::ActionStmt. The assumption
 // is that the ActionStmt will be either an assignment or a pointer-assignment,
 // otherwise return std::nullopt.

flang/lib/Semantics/check-omp-structure.cpp

@@ -495,6 +495,12 @@ template <typename Checker> struct DirectiveSpellingVisitor {
   template <typename T> bool Pre(const T &) { return true; }
   template <typename T> void Post(const T &) {}
 
+  template <typename... Ts>
+  static const parser::OmpDirectiveName &GetDirName(
+      const std::tuple<Ts...> &t) {
+    return std::get<parser::OmpDirectiveSpecification>(t).DirName();
+  }
+
   bool Pre(const parser::OmpSectionsDirective &x) {
     checker_(x.source, x.v);
     return false;
@@ -520,8 +526,7 @@ template <typename Checker> struct DirectiveSpellingVisitor {
     return false;
   }
   bool Pre(const parser::OpenMPAllocatorsConstruct &x) {
-    checker_(
-        std::get<parser::Verbatim>(x.t).source, Directive::OMPD_allocators);
+    checker_(GetDirName(x.t).source, Directive::OMPD_allocators);
     return false;
   }
   bool Pre(const parser::OmpAssumeDirective &x) {
@@ -1657,26 +1662,45 @@ void OmpStructureChecker::Leave(const parser::OpenMPExecutableAllocate &x) {
 
 void OmpStructureChecker::Enter(const parser::OpenMPAllocatorsConstruct &x) {
   isPredefinedAllocator = true;
-  const auto &dir{std::get<parser::Verbatim>(x.t)};
-  PushContextAndClauseSets(dir.source, llvm::omp::Directive::OMPD_allocators);
-  const auto &clauseList{std::get<parser::OmpClauseList>(x.t)};
-  for (const auto &clause : clauseList.v) {
+
+  auto &dirSpec{std::get<parser::OmpDirectiveSpecification>(x.t)};
+  auto &block{std::get<parser::Block>(x.t)};
+  PushContextAndClauseSets(
+      dirSpec.DirName().source, llvm::omp::Directive::OMPD_allocators);
+
+  if (block.empty()) {
+    context_.Say(dirSpec.source,
+        "The ALLOCATORS construct should contain a single ALLOCATE statement"_err_en_US);
+    return;
+  }
+
+  omp::SourcedActionStmt action{omp::GetActionStmt(block)};
+  const auto *allocate{
+      action ? parser::Unwrap<parser::AllocateStmt>(action.stmt) : nullptr};
+
+  if (!allocate) {
+    const parser::CharBlock &source = action ? action.source : x.source;
+    context_.Say(source,
+        "The body of the ALLOCATORS construct should be an ALLOCATE statement"_err_en_US);
+  }
+
+  for (const auto &clause : dirSpec.Clauses().v) {
     if (const auto *allocClause{
             parser::Unwrap<parser::OmpClause::Allocate>(clause)}) {
       CheckVarIsNotPartOfAnotherVar(
-          dir.source, std::get<parser::OmpObjectList>(allocClause->v.t));
+          dirSpec.source, std::get<parser::OmpObjectList>(allocClause->v.t));
     }
   }
 }
 
 void OmpStructureChecker::Leave(const parser::OpenMPAllocatorsConstruct &x) {
-  const auto &dir{std::get<parser::Verbatim>(x.t)};
-  const auto &clauseList{std::get<parser::OmpClauseList>(x.t)};
-  for (const auto &clause : clauseList.v) {
+  auto &dirSpec{std::get<parser::OmpDirectiveSpecification>(x.t)};
+
+  for (const auto &clause : dirSpec.Clauses().v) {
     if (const auto *allocClause{
             std::get_if<parser::OmpClause::Allocate>(&clause.u)}) {
       CheckPredefinedAllocatorRestriction(
-          dir.source, std::get<parser::OmpObjectList>(allocClause->v.t));
+          dirSpec.source, std::get<parser::OmpObjectList>(allocClause->v.t));
     }
   }
   dirContext_.pop_back();

flang/lib/Semantics/openmp-utils.cpp

@@ -38,6 +38,26 @@
 
 namespace Fortran::semantics::omp {
 
+SourcedActionStmt GetActionStmt(const parser::ExecutionPartConstruct *x) {
+  if (x == nullptr) {
+    return SourcedActionStmt{};
+  }
+  if (auto *exec{std::get_if<parser::ExecutableConstruct>(&x->u)}) {
+    using ActionStmt = parser::Statement<parser::ActionStmt>;
+    if (auto *stmt{std::get_if<ActionStmt>(&exec->u)}) {
+      return SourcedActionStmt{&stmt->statement, stmt->source};
+    }
+  }
+  return SourcedActionStmt{};
+}
+
+SourcedActionStmt GetActionStmt(const parser::Block &block) {
+  if (block.size() == 1) {
+    return GetActionStmt(&block.front());
+  }
+  return SourcedActionStmt{};
+}
+
 std::string ThisVersion(unsigned version) {
   std::string tv{
       std::to_string(version / 10) + "." + std::to_string(version % 10)};

flang/lib/Semantics/openmp-utils.h

@@ -29,6 +29,19 @@ class Symbol;
 
 // Add this namespace to avoid potential conflicts
 namespace omp {
+// There is no consistent way to get the source of an ActionStmt, but there
+// is "source" in Statement<T>. This structure keeps the ActionStmt with the
+// extracted source for further use.
+struct SourcedActionStmt {
+  const parser::ActionStmt *stmt{nullptr};
+  parser::CharBlock source;
+
+  operator bool() const { return stmt != nullptr; }
+};
+
+SourcedActionStmt GetActionStmt(const parser::ExecutionPartConstruct *x);
+SourcedActionStmt GetActionStmt(const parser::Block &block);
+
 std::string ThisVersion(unsigned version);
 std::string TryVersion(unsigned version);