[flang] OpenMP allocate directive parse tree fix

Addresses the same issue as the following abandoned revision: D104391.

Rewrite leading declarative allocations so they are nested within their respective executable allocate directive

Original:
```
ExecutionPartConstruct -> OpenMPDeclarativeAllocate
ExecutionPartConstruct -> OpenMPDeclarativeAllocate
ExecutionPartConstruct -> OpenMPExecutableAllocate
```

After rewriting:
```
ExecutionPartConstruct -> OpenMPExecutableAllocate
| ExecutionPartConstruct -> OpenMPDeclarativeAllocate
| ExecutionPartConstruct -> OpenMPDeclarativeAllocate
```

Reviewed By: kiranchandramohan

Differential Revision: https://reviews.llvm.org/D148409

5 files changed, 144 insertions, 2 deletions

diff --git a/flang/lib/Parser/unparse.cpp b/flang/lib/Parser/unparse.cpp
index 6916052cf78d..3b34c4ec89ce 100644
--- a/flang/lib/Parser/unparse.cpp
+++ b/flang/lib/Parser/unparse.cpp
@@ -2325,6 +2325,14 @@ public:
     EndOpenMP();
   }
   void Unparse(const OpenMPExecutableAllocate &x) {
+    const auto &fields =
+        std::get<std::optional<std::list<parser::OpenMPDeclarativeAllocate>>>(
+            x.t);
+    if (fields) {
+      for (const auto &decl : *fields) {
+        Walk(decl);
+      }
+    }
     BeginOpenMP();
     Word("!$OMP ALLOCATE");
     Walk(" (", std::get<std::optional<OmpObjectList>>(x.t), ")");
diff --git a/flang/lib/Semantics/canonicalize-omp.cpp b/flang/lib/Semantics/canonicalize-omp.cpp
index dc8a2a4d93af..013fb408214e 100644
--- a/flang/lib/Semantics/canonicalize-omp.cpp
+++ b/flang/lib/Semantics/canonicalize-omp.cpp
@@ -15,7 +15,9 @@
 //   1. move structured DoConstruct and OmpEndLoopDirective into
 //      OpenMPLoopConstruct. Compilation will not proceed in case of errors
 //      after this pass.
-//   2. TBD
+//   2. Associate declarative OMP allocation directives with their
+//      respective executable allocation directive
+//   3. TBD
 namespace Fortran::semantics {
 
 using namespace parser::literals;
@@ -46,6 +48,8 @@ public:
     } // Block list
   }
 
+  void Post(parser::ExecutionPart &body) { RewriteOmpAllocations(body); }
+
 private:
   template <typename T> T *GetConstructIf(parser::ExecutionPartConstruct &x) {
     if (auto *y{std::get_if<parser::ExecutableConstruct>(&x.u)}) {
@@ -56,6 +60,15 @@ private:
     return nullptr;
   }
 
+  template <typename T> T *GetOmpIf(parser::ExecutionPartConstruct &x) {
+    if (auto *construct{GetConstructIf<parser::OpenMPConstruct>(x)}) {
+      if (auto *omp{std::get_if<T>(&construct->u)}) {
+        return omp;
+      }
+    }
+    return nullptr;
+  }
+
   void RewriteOpenMPLoopConstruct(parser::OpenMPLoopConstruct &x,
       parser::Block &block, parser::Block::iterator it) {
     // Check the sequence of DoConstruct and OmpEndLoopDirective
@@ -106,6 +119,36 @@ private:
         parser::ToUpperCaseLetters(dir.source.ToString()));
   }
 
+  void RewriteOmpAllocations(parser::ExecutionPart &body) {
+    // Rewrite leading declarative allocations so they are nested
+    // within their respective executable allocate directive
+    //
+    // Original:
+    //   ExecutionPartConstruct -> OpenMPDeclarativeAllocate
+    //   ExecutionPartConstruct -> OpenMPDeclarativeAllocate
+    //   ExecutionPartConstruct -> OpenMPExecutableAllocate
+    //
+    // After rewriting:
+    //   ExecutionPartConstruct -> OpenMPExecutableAllocate
+    //     ExecutionPartConstruct -> OpenMPDeclarativeAllocate
+    //     ExecutionPartConstruct -> OpenMPDeclarativeAllocate
+    for (auto it = body.v.rbegin(); it != body.v.rend();) {
+      if (auto *exec = GetOmpIf<parser::OpenMPExecutableAllocate>(*(it++))) {
+        parser::OpenMPDeclarativeAllocate *decl;
+        std::list<parser::OpenMPDeclarativeAllocate> subAllocates;
+        while (it != body.v.rend() &&
+            (decl = GetOmpIf<parser::OpenMPDeclarativeAllocate>(*it))) {
+          subAllocates.push_front(std::move(*decl));
+          it = decltype(it)(body.v.erase(std::next(it).base()));
+        }
+        if (!subAllocates.empty()) {
+          std::get<std::optional<std::list<parser::OpenMPDeclarativeAllocate>>>(
+              exec->t) = {std::move(subAllocates)};
+        }
+      }
+    }
+  }
+
   parser::Messages &messages_;
 };
 
diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index 9eb4d9836491..1052c459632e 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1691,7 +1691,8 @@ void OmpAttributeVisitor::ResolveOmpObject(
                   }
                 }
                 if (ompFlag == Symbol::Flag::OmpDeclarativeAllocateDirective &&
-                    IsAllocatable(*symbol)) {
+                    IsAllocatable(*symbol) &&
+                    !IsNestedInDirective(llvm::omp::Directive::OMPD_allocate)) {
                   context_.Say(designator.source,
                       "List items specified in the ALLOCATE directive must not "
                       "have the ALLOCATABLE attribute unless the directive is "
diff --git a/flang/test/Parser/OpenMP/allocate-tree-spec-part.f90 b/flang/test/Parser/OpenMP/allocate-tree-spec-part.f90
new file mode 100644
index 000000000000..45a693d2cb04
--- /dev/null
+++ b/flang/test/Parser/OpenMP/allocate-tree-spec-part.f90
@@ -0,0 +1,47 @@
+! RUN: %flang_fc1 -fopenmp -fdebug-dump-parse-tree %s | FileCheck %s
+! Ensures associated declarative OMP allocations in the specification
+! part are kept there
+
+program allocate_tree
+    use omp_lib
+    integer, allocatable :: w, xarray(:), zarray(:, :)
+    integer :: f
+!$omp allocate(f) allocator(omp_default_mem_alloc)
+    f = 2
+!$omp allocate(w) allocator(omp_const_mem_alloc)
+!$omp allocate(xarray) allocator(omp_large_cap_mem_alloc)
+!$omp allocate(zarray) allocator(omp_default_mem_alloc)
+!$omp allocate
+    allocate (w, xarray(4), zarray(5, f))
+end program allocate_tree
+
+!CHECK: | | DeclarationConstruct -> SpecificationConstruct -> OpenMPDeclarativeConstruct -> OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | Verbatim
+!CHECK-NEXT: | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'f'
+!CHECK-NEXT: | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | ExecutionPart -> Block
+!CHECK-NEXT: | | ExecutionPartConstruct -> ExecutableConstruct -> ActionStmt -> AssignmentStmt = 'f=2_4'
+!CHECK-NEXT: | | | Variable = 'f'
+!CHECK-NEXT: | | | | Designator -> DataRef -> Name = 'f'
+!CHECK-NEXT: | | | Expr = '2_4'
+!CHECK-NEXT: | | | | LiteralConstant -> IntLiteralConstant = '2'
+!CHECK-NEXT: | | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPExecutableAllocate
+!CHECK-NEXT: | | | Verbatim
+!CHECK-NEXT: | | | OmpClauseList ->
+!CHECK-NEXT: | | | OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | | Verbatim
+!CHECK-NEXT: | | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'w'
+!CHECK-NEXT: | | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | | | OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | | Verbatim
+!CHECK-NEXT: | | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'xarray'
+!CHECK-NEXT: | | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | | | OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | | Verbatim
+!CHECK-NEXT: | | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'zarray'
+!CHECK-NEXT: | | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | | | AllocateStmt
diff --git a/flang/test/Parser/OpenMP/allocate-tree.f90 b/flang/test/Parser/OpenMP/allocate-tree.f90
new file mode 100644
index 000000000000..f04e431e74ae
--- /dev/null
+++ b/flang/test/Parser/OpenMP/allocate-tree.f90
@@ -0,0 +1,43 @@
+! RUN: %flang_fc1 -fopenmp -fdebug-dump-parse-tree %s | FileCheck %s
+! RUN: %flang_fc1 -fopenmp -fdebug-unparse %s | FileCheck %s --check-prefix="UNPARSE"
+! Ensures associated declarative OMP allocations are nested in their
+! corresponding executable allocate directive
+
+program allocate_tree
+    use omp_lib
+    integer, allocatable :: w, xarray(:), zarray(:, :)
+    integer :: z, t
+    t = 2
+    z = 3
+!$omp allocate(w) allocator(omp_const_mem_alloc)
+!$omp allocate(xarray) allocator(omp_large_cap_mem_alloc)
+!$omp allocate(zarray) allocator(omp_default_mem_alloc)
+!$omp allocate
+    allocate(w, xarray(4), zarray(t, z))
+end program allocate_tree
+
+!CHECK: | | ExecutionPartConstruct -> ExecutableConstruct -> OpenMPConstruct -> OpenMPExecutableAllocate
+!CHECK-NEXT: | | | Verbatim
+!CHECK-NEXT: | | | OmpClauseList ->
+!CHECK-NEXT: | | | OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | | Verbatim
+!CHECK-NEXT: | | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'w'
+!CHECK-NEXT: | | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | | | OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | | Verbatim
+!CHECK-NEXT: | | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'xarray'
+!CHECK-NEXT: | | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | | | OpenMPDeclarativeAllocate
+!CHECK-NEXT: | | | | Verbatim
+!CHECK-NEXT: | | | | OmpObjectList -> OmpObject -> Designator -> DataRef -> Name = 'zarray'
+!CHECK-NEXT: | | | | OmpClauseList -> OmpClause -> Allocator -> Scalar -> Integer -> Expr =
+!CHECK-NEXT: | | | | | Designator -> DataRef -> Name =
+!CHECK-NEXT: | | | AllocateStmt
+
+!UNPARSE: !$OMP ALLOCATE (w) ALLOCATOR(1_4)
+!UNPARSE-NEXT: !$OMP ALLOCATE (xarray) ALLOCATOR(1_4)
+!UNPARSE-NEXT: !$OMP ALLOCATE (zarray) ALLOCATOR(1_4)
+!UNPARSE-NEXT: !$OMP ALLOCATE
+!UNPARSE-NEXT: ALLOCATE(w, xarray(4_4), zarray(t,z))