1 files changed, 39 insertions, 0 deletions
diff --git a/flang/lib/Parser/openmp-utils.cpp b/flang/lib/Parser/openmp-utils.cpp
index 95ad3f6..b9d3763c 100644
--- a/flang/lib/Parser/openmp-utils.cpp
+++ b/flang/lib/Parser/openmp-utils.cpp
@@ -22,6 +22,25 @@
 
 namespace Fortran::parser::omp {
 
+const OpenMPDeclarativeConstruct *GetOmp(const DeclarationConstruct &x) {
+  if (auto *y = std::get_if<SpecificationConstruct>(&x.u)) {
+    if (auto *z{std::get_if<common::Indirection<OpenMPDeclarativeConstruct>>(
+            &y->u)}) {
+      return &z->value();
+    }
+  }
+  return nullptr;
+}
+
+const OpenMPConstruct *GetOmp(const ExecutionPartConstruct &x) {
+  if (auto *y{std::get_if<ExecutableConstruct>(&x.u)}) {
+    if (auto *z{std::get_if<common::Indirection<OpenMPConstruct>>(&y->u)}) {
+      return &z->value();
+    }
+  }
+  return nullptr;
+}
+
 const OmpObjectList *GetOmpObjectList(const OmpClause &clause) {
   // Clauses with OmpObjectList as its data member
   using MemberObjectListClauses = std::tuple<OmpClause::Copyin,
@@ -86,4 +105,24 @@ const OmpInitializerExpression *GetInitializerExpr(const OmpClause &init) {
   return nullptr;
 }
 
+static void SplitOmpAllocateHelper(
+    OmpAllocateInfo &n, const OmpAllocateDirective &x) {
+  n.dirs.push_back(&x);
+  const Block &body{std::get<Block>(x.t)};
+  if (!body.empty()) {
+    if (auto *omp{GetOmp(body.front())}) {
+      if (auto *ad{std::get_if<OmpAllocateDirective>(&omp->u)}) {
+        return SplitOmpAllocateHelper(n, *ad);
+      }
+    }
+    n.body = &body.front();
+  }
+}
+
+OmpAllocateInfo SplitOmpAllocate(const OmpAllocateDirective &x) {
+  OmpAllocateInfo info;
+  SplitOmpAllocateHelper(info, x);
+  return info;
+}
+
 } // namespace Fortran::parser::omp