[flang][OpenMP] Use new modifiers in ALLOCATE clause

flang/examples/FeatureList/FeatureList.cpp

@@ -516,10 +516,10 @@ struct NodeVisitor {
   READ_FEATURE(OmpReductionInitializerClause)
   READ_FEATURE(OmpReductionIdentifier)
   READ_FEATURE(OmpAllocateClause)
-  READ_FEATURE(OmpAllocateClause::AllocateModifier)
-  READ_FEATURE(OmpAllocateClause::AllocateModifier::Allocator)
-  READ_FEATURE(OmpAllocateClause::AllocateModifier::ComplexModifier)
-  READ_FEATURE(OmpAllocateClause::AllocateModifier::Align)
+  READ_FEATURE(OmpAllocateClause::Modifier)
+  READ_FEATURE(OmpAllocatorSimpleModifier)
+  READ_FEATURE(OmpAllocatorComplexModifier)
+  READ_FEATURE(OmpAlignModifier)
   READ_FEATURE(OmpScheduleClause)
   READ_FEATURE(OmpScheduleClause::Kind)
   READ_FEATURE(OmpScheduleClause::Modifier)

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

@@ -586,10 +586,10 @@ class ParseTreeDumper {
   NODE(parser, OmpReductionInitializerClause)
   NODE(parser, OmpReductionIdentifier)
   NODE(parser, OmpAllocateClause)
-  NODE(OmpAllocateClause, AllocateModifier)
-  NODE(OmpAllocateClause::AllocateModifier, Allocator)
-  NODE(OmpAllocateClause::AllocateModifier, ComplexModifier)
-  NODE(OmpAllocateClause::AllocateModifier, Align)
+  NODE(OmpAllocateClause, Modifier)
+  NODE(parser, OmpAlignModifier)
+  NODE(parser, OmpAllocatorComplexModifier)
+  NODE(parser, OmpAllocatorSimpleModifier)
   NODE(parser, OmpScheduleClause)
   NODE(OmpScheduleClause, Modifier)
   NODE_ENUM(OmpScheduleClause, Kind)

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

@@ -3457,6 +3457,30 @@ inline namespace modifier {
 //   ENUM_CLASS(Value, Keyword1, Keyword2);
 // };
 
+// Ref: [5.1:184-185], [5.2:178-179]
+//
+// align-modifier ->
+//    ALIGN(alignment)                              // since 5.1
+struct OmpAlignModifier {
+  WRAPPER_CLASS_BOILERPLATE(OmpAlignModifier, ScalarIntExpr);
+};
+
+// Ref: [5.0:158-159], [5.1:184-185], [5.2:178-179]
+//
+// allocator-simple-modifier ->
+//    allocator                                     // since 5.0
+struct OmpAllocatorSimpleModifier {
+  WRAPPER_CLASS_BOILERPLATE(OmpAllocatorSimpleModifier, ScalarIntExpr);
+};
+
+// Ref: [5.1:184-185], [5.2:178-179]
+//
+// allocator-complex-modifier ->
+//    ALLOCATOR(allocator)                          // since 5.1
+struct OmpAllocatorComplexModifier {
+  WRAPPER_CLASS_BOILERPLATE(OmpAllocatorComplexModifier, ScalarIntExpr);
+};
+
 // Ref: [5.2:252-254]
 //
 // chunk-modifier ->
@@ -3646,24 +3670,20 @@ struct OmpAlignedClause {
   std::tuple<OmpObjectList, std::optional<ScalarIntConstantExpr>> t;
 };
 
-// OMP 5.0 2.11.4 allocate-clause -> ALLOCATE ([allocator:] variable-name-list)
-// OMP 5.2 2.13.4 allocate-clause -> ALLOCATE ([allocate-modifier [,
-//  	                               allocate-modifier] :]
-//                                   variable-name-list)
-//                allocate-modifier -> allocator | align
+// Ref: [5.0:158-159], [5.1:184-185], [5.2:178-179]
+//
+// allocate-clause ->
+//    ALLOCATE(
+//        [allocator-simple-modifier:] list) |      // since 5.0
+//    ALLOCATE([modifier...:] list)                 // since 5.1
+// modifier ->
+//    allocator-simple-modifier |
+//    allocator-complex-modifier | align-modifier   // since 5.1
 struct OmpAllocateClause {
-  struct AllocateModifier {
-    WRAPPER_CLASS(Allocator, ScalarIntExpr);
-    WRAPPER_CLASS(Align, ScalarIntExpr);
-    struct ComplexModifier {
-      TUPLE_CLASS_BOILERPLATE(ComplexModifier);
-      std::tuple<Allocator, Align> t;
-    };
-    UNION_CLASS_BOILERPLATE(AllocateModifier);
-    std::variant<Allocator, ComplexModifier, Align> u;
-  };
+  MODIFIER_BOILERPLATE(OmpAlignModifier, OmpAllocatorSimpleModifier,
+      OmpAllocatorComplexModifier);
   TUPLE_CLASS_BOILERPLATE(OmpAllocateClause);
-  std::tuple<std::optional<AllocateModifier>, OmpObjectList> t;
+  std::tuple<MODIFIERS(), OmpObjectList> t;
 };
 
 // OMP 5.0 2.4 atomic-default-mem-order-clause ->

flang/include/flang/Semantics/openmp-modifiers.h

@@ -67,6 +67,9 @@ template <typename SpecificTy> const OmpModifierDescriptor &OmpGetDescriptor();
 #define DECLARE_DESCRIPTOR(name) \
   template <> const OmpModifierDescriptor &OmpGetDescriptor<name>()
 
+DECLARE_DESCRIPTOR(parser::OmpAlignModifier);
+DECLARE_DESCRIPTOR(parser::OmpAllocatorComplexModifier);
+DECLARE_DESCRIPTOR(parser::OmpAllocatorSimpleModifier);
 DECLARE_DESCRIPTOR(parser::OmpChunkModifier);
 DECLARE_DESCRIPTOR(parser::OmpDependenceType);
 DECLARE_DESCRIPTOR(parser::OmpExpectation);
@@ -216,10 +219,26 @@ OmpGetRepeatableModifier(const std::optional<std::list<UnionTy>> &modifiers) {
       OmpSpecificModifierIterator(items, items->end()));
 }
 
+// Attempt to prevent creating a range based on an expiring modifier list.
 template <typename SpecificTy, typename UnionTy>
 llvm::iterator_range<OmpSpecificModifierIterator<SpecificTy>>
 OmpGetRepeatableModifier(std::optional<std::list<UnionTy>> &&) = delete;
 
+template <typename SpecificTy, typename UnionTy>
+Fortran::parser::CharBlock OmpGetModifierSource(
+    const std::optional<std::list<UnionTy>> &modifiers,
+    const SpecificTy *specific) {
+  if (!modifiers || !specific) {
+    return Fortran::parser::CharBlock{};
+  }
+  for (auto &m : *modifiers) {
+    if (std::get_if<SpecificTy>(&m.u) == specific) {
+      return m.source;
+    }
+  }
+  llvm_unreachable("`specific` must be a member of `modifiers`");
+}
+
 namespace detail {
 template <typename T> constexpr const T *make_nullptr() {
   return static_cast<const T *>(nullptr);

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

@@ -81,12 +81,8 @@ genAllocateClause(lower::AbstractConverter &converter,
   // Check if allocate clause has allocator specified. If so, add it
   // to list of allocators, otherwise, add default allocator to
   // list of allocators.
-  using SimpleModifier = Allocate::AllocatorSimpleModifier;
   using ComplexModifier = Allocate::AllocatorComplexModifier;
-  if (auto &mod = std::get<std::optional<SimpleModifier>>(clause.t)) {
-    mlir::Value operand = fir::getBase(converter.genExprValue(*mod, stmtCtx));
-    allocatorOperands.append(objects.size(), operand);
-  } else if (auto &mod = std::get<std::optional<ComplexModifier>>(clause.t)) {
+  if (auto &mod = std::get<std::optional<ComplexModifier>>(clause.t)) {
     mlir::Value operand = fir::getBase(converter.genExprValue(mod->v, stmtCtx));
     allocatorOperands.append(objects.size(), operand);
   } else {

flang/lib/Lower/OpenMP/Clauses.cpp

@@ -415,47 +415,29 @@ Aligned make(const parser::OmpClause::Aligned &inp,
 Allocate make(const parser::OmpClause::Allocate &inp,
               semantics::SemanticsContext &semaCtx) {
   // inp.v -> parser::OmpAllocateClause
-  using wrapped = parser::OmpAllocateClause;
-  auto &t0 = std::get<std::optional<wrapped::AllocateModifier>>(inp.v.t);
+  auto &mods = semantics::OmpGetModifiers(inp.v);
+  auto *m0 = semantics::OmpGetUniqueModifier<parser::OmpAlignModifier>(mods);
+  auto *m1 =
+      semantics::OmpGetUniqueModifier<parser::OmpAllocatorComplexModifier>(
+          mods);
+  auto *m2 =
+      semantics::OmpGetUniqueModifier<parser::OmpAllocatorSimpleModifier>(mods);
   auto &t1 = std::get<parser::OmpObjectList>(inp.v.t);
 
-  if (!t0) {
-    return Allocate{{/*AllocatorSimpleModifier=*/std::nullopt,
-                     /*AllocatorComplexModifier=*/std::nullopt,
-                     /*AlignModifier=*/std::nullopt,
-                     /*List=*/makeObjects(t1, semaCtx)}};
-  }
+  auto makeAllocator = [&](auto *mod) -> std::optional<Allocator> {
+    if (mod)
+      return Allocator{makeExpr(mod->v, semaCtx)};
+    return std::nullopt;
+  };
 
-  using Tuple = decltype(Allocate::t);
+  auto makeAlign = [&](const parser::ScalarIntExpr &expr) {
+    return Align{makeExpr(expr, semaCtx)};
+  };
 
-  return Allocate{Fortran::common::visit(
-      common::visitors{
-          // simple-modifier
-          [&](const wrapped::AllocateModifier::Allocator &v) -> Tuple {
-            return {/*AllocatorSimpleModifier=*/makeExpr(v.v, semaCtx),
-                    /*AllocatorComplexModifier=*/std::nullopt,
-                    /*AlignModifier=*/std::nullopt,
-                    /*List=*/makeObjects(t1, semaCtx)};
-          },
-          // complex-modifier + align-modifier
-          [&](const wrapped::AllocateModifier::ComplexModifier &v) -> Tuple {
-            auto &s0 = std::get<wrapped::AllocateModifier::Allocator>(v.t);
-            auto &s1 = std::get<wrapped::AllocateModifier::Align>(v.t);
-            return {
-                /*AllocatorSimpleModifier=*/std::nullopt,
-                /*AllocatorComplexModifier=*/Allocator{makeExpr(s0.v, semaCtx)},
-                /*AlignModifier=*/Align{makeExpr(s1.v, semaCtx)},
-                /*List=*/makeObjects(t1, semaCtx)};
-          },
-          // align-modifier
-          [&](const wrapped::AllocateModifier::Align &v) -> Tuple {
-            return {/*AllocatorSimpleModifier=*/std::nullopt,
-                    /*AllocatorComplexModifier=*/std::nullopt,
-                    /*AlignModifier=*/Align{makeExpr(v.v, semaCtx)},
-                    /*List=*/makeObjects(t1, semaCtx)};
-          },
-      },
-      t0->u)};
+  auto maybeAllocator = m1 ? makeAllocator(m1) : makeAllocator(m2);
+  return Allocate{{/*AllocatorComplexModifier=*/std::move(maybeAllocator),
+                   /*AlignModifier=*/maybeApplyToV(makeAlign, m0),
+                   /*List=*/makeObjects(t1, semaCtx)}};
 }
 
 Allocator make(const parser::OmpClause::Allocator &inp,

flang/lib/Lower/OpenMP/Clauses.h

@@ -153,10 +153,11 @@ std::optional<ResultTy> maybeApply(FuncTy &&func,
   return func(*arg);
 }
 
-template <
+template <           //
     typename FuncTy, //
     typename ArgTy,  //
-    typename ResultTy = std::invoke_result_t<FuncTy, typename ArgTy::Value>>
+    typename ResultTy =
+        std::invoke_result_t<FuncTy, decltype(std::declval<ArgTy>().v)>>
 std::optional<ResultTy> maybeApplyToV(FuncTy &&func, const ArgTy *arg) {
   if (!arg)
     return std::nullopt;

flang/lib/Parser/openmp-parsers.cpp

@@ -91,6 +91,14 @@ static TypeDeclarationStmt makeIterSpecDecl(std::list<ObjectName> &&names) {
 
 // --- Parsers for clause modifiers -----------------------------------
 
+TYPE_PARSER(construct<OmpAlignModifier>( //
+    "ALIGN" >> parenthesized(scalarIntExpr)))
+
+TYPE_PARSER(construct<OmpAllocatorComplexModifier>(
+    "ALLOCATOR" >> parenthesized(scalarIntExpr)))
+
+TYPE_PARSER(construct<OmpAllocatorSimpleModifier>(scalarIntExpr))
+
 TYPE_PARSER(construct<OmpChunkModifier>( //
     "SIMD" >> pure(OmpChunkModifier::Value::Simd)))
 
@@ -183,6 +191,16 @@ TYPE_PARSER(construct<OmpVariableCategory>(
     "SCALAR" >> pure(OmpVariableCategory::Value::Scalar)))
 
 // This could be auto-generated.
+TYPE_PARSER(sourced(construct<OmpAllocateClause::Modifier>(sourced(
+    construct<OmpAllocateClause::Modifier>(Parser<OmpAlignModifier>{}) ||
+    construct<OmpAllocateClause::Modifier>(
+        Parser<OmpAllocatorComplexModifier>{}) ||
+    construct<OmpAllocateClause::Modifier>(
+        Parser<OmpAllocatorSimpleModifier>{})))))
+
+TYPE_PARSER(sourced(
+    construct<OmpDefaultmapClause::Modifier>(Parser<OmpVariableCategory>{})))
+
 TYPE_PARSER(sourced(construct<OmpFromClause::Modifier>(
     sourced(construct<OmpFromClause::Modifier>(Parser<OmpExpectation>{}) ||
         construct<OmpFromClause::Modifier>(Parser<OmpMapper>{}) ||
@@ -211,9 +229,6 @@ TYPE_PARSER(sourced(construct<OmpToClause::Modifier>(
         construct<OmpToClause::Modifier>(Parser<OmpMapper>{}) ||
         construct<OmpToClause::Modifier>(Parser<OmpIterator>{})))))
 
-TYPE_PARSER(sourced(
-    construct<OmpDefaultmapClause::Modifier>(Parser<OmpVariableCategory>{})))
-
 // --- Parsers for clauses --------------------------------------------
 
 /// `MOBClause` is a clause that has a
@@ -334,29 +349,7 @@ TYPE_PARSER(construct<OmpInReductionClause>(
 //                                   variable-name-list)
 //                allocate-modifier -> allocator | align
 TYPE_PARSER(construct<OmpAllocateClause>(
-    maybe(
-        first(
-            construct<OmpAllocateClause::AllocateModifier>("ALLOCATOR" >>
-                construct<OmpAllocateClause::AllocateModifier::ComplexModifier>(
-                    parenthesized(construct<
-                        OmpAllocateClause::AllocateModifier::Allocator>(
-                        scalarIntExpr)) /
-                        ",",
-                    "ALIGN" >> parenthesized(construct<
-                                   OmpAllocateClause::AllocateModifier::Align>(
-                                   scalarIntExpr)))),
-            construct<OmpAllocateClause::AllocateModifier>("ALLOCATOR" >>
-                parenthesized(
-                    construct<OmpAllocateClause::AllocateModifier::Allocator>(
-                        scalarIntExpr))),
-            construct<OmpAllocateClause::AllocateModifier>("ALIGN" >>
-                parenthesized(
-                    construct<OmpAllocateClause::AllocateModifier::Align>(
-                        scalarIntExpr))),
-            construct<OmpAllocateClause::AllocateModifier>(
-                construct<OmpAllocateClause::AllocateModifier::Allocator>(
-                    scalarIntExpr))) /
-        ":"),
+    maybe(nonemptyList(Parser<OmpAllocateClause::Modifier>{}) / ":"),
     Parser<OmpObjectList>{}))
 
 // iteration-offset -> +/- non-negative-constant-expr

flang/lib/Parser/unparse.cpp

@@ -2148,35 +2148,21 @@ class UnparseVisitor {
     Walk(std::get<OmpObjectList>(x.t));
   }
   void Unparse(const OmpAllocateClause &x) {
-    Walk(
-        std::get<std::optional<OmpAllocateClause::AllocateModifier>>(x.t), ":");
+    using Modifier = OmpAllocateClause::Modifier;
+    Walk(std::get<std::optional<std::list<Modifier>>>(x.t), ": ");
     Walk(std::get<OmpObjectList>(x.t));
   }
-  void Unparse(const OmpAllocateClause::AllocateModifier &x) {
-    common::visit(
-        common::visitors{
-            [&](const OmpAllocateClause::AllocateModifier::Allocator &y) {
-              Walk(y);
-            },
-            [&](const OmpAllocateClause::AllocateModifier::ComplexModifier &y) {
-              Word("ALLOCATOR(");
-              Walk(std::get<OmpAllocateClause::AllocateModifier::Allocator>(
-                  y.t));
-              Put(")");
-              Put(",");
-              Walk(std::get<OmpAllocateClause::AllocateModifier::Align>(y.t));
-            },
-            [&](const OmpAllocateClause::AllocateModifier::Align &y) {
-              Walk(y);
-            },
-        },
-        x.u);
-  }
-  void Unparse(const OmpAllocateClause::AllocateModifier::Align &x) {
+  void Unparse(const OmpAlignModifier &x) {
     Word("ALIGN(");
     Walk(x.v);
     Put(")");
   }
+  void Unparse(const OmpAllocatorSimpleModifier &x) { Walk(x.v); }
+  void Unparse(const OmpAllocatorComplexModifier &x) {
+    Word("ALLOCATOR(");
+    Walk(x.v);
+    Put(")");
+  }
   void Unparse(const OmpOrderClause &x) {
     using Modifier = OmpOrderClause::Modifier;
     Walk(std::get<std::optional<std::list<Modifier>>>(x.t), ":");