[flang] [openmp] Correct OpenMP Allocate Directive parse tree structure to associate directives with allocate statements correctly.

[I_Perry]

Summary

This patch implements a parse tree rewrite such that openmp allocate directives which should be associated with fortran allocate statements are done so.

[I_Perry]

Added test to stop unambiguous start of executable allocate from being ignored.

[kiranchandramohan]

Can this be in the Canonicalize OpenMP file?
flang/lib/Semantics/canonicalize-omp.cpp

[I_Perry]

Addressed comment to move changes to canonicalize-omp.cpp

[I_Perry]

Removed included <iterator> in rewrite-parse-tree.cpp

flang/lib/Parser/unparse.cpp

void Unparse(const OmpAtomicWrite &x) {
     Put("\n");
     EndOpenMP();
     Walk(std::get<Statement<AssignmentStmt>>(x.t));
     BeginOpenMP();
     Walk(std::get<std::optional<OmpEndAtomic>>(x.t), "!$OMP END ATOMIC\n");
     EndOpenMP();
   }
   void Unparse(const OpenMPExecutableAllocate &x) {
+    auto &optionalList{
+        std::get<std::optional<std::list<OpenMPDeclarativeAllocate>>>(x.t)};
+    if (optionalList.has_value()) {
+      for (auto &declarativeAllocate : optionalList.value()) {
+        Walk(declarativeAllocate);
+      }
+    }
     BeginOpenMP();
     Word("!$OMP ALLOCATE");
     Walk(" (", std::get<std::optional<OmpObjectList>>(x.t), ")");
     Walk(std::get<OmpClauseList>(x.t));
     Put("\n");
     EndOpenMP();
     Walk(std::get<Statement<AllocateStmt>>(x.t));
   }

flang/lib/Semantics/canonicalize-omp.cpp

 //===-- lib/Semantics/canonicalize-omp.cpp --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "canonicalize-omp.h"
+#include "flang/Common/reference.h"
 #include "flang/Parser/parse-tree-visitor.h"
+#include "flang/Parser/tools.h"
 
 // After Loop Canonicalization, rewrite OpenMP parse tree to make OpenMP
 // Constructs more structured which provide explicit scopes for later
 // structural checks and semantic analysis.
 //   1. move structured DoConstruct and OmpEndLoopDirective into
 //      OpenMPLoopConstruct. Compilation will not proceed in case of errors
 //      after this pass.
 //   2. TBD
for (auto it{block.begin()}; it != block.end(); ++it) {
         messages_.Say(dir.source,
             "The %s directive must follow the DO loop associated with the "
             "loop construct"_err_en_US,
             parser::ToUpperCaseLetters(dir.source.ToString()));
       }
     } // Block list
   }
 
+  bool Pre(parser::ExecutionPart &x) {
+    executionPartIndex = 0;
+    return true;
+  }
+
+  void Post(parser::ExecutionPart &x) {
+    auto it{x.v.begin()};
+    int previousExecutableCandidate{0};
+    for (auto &executableCandidate : movedExecutableOmpAllocateCandidates) {
+      std::advance(it, executableCandidate - previousExecutableCandidate);
+      x.v.erase(it++);
+      previousExecutableCandidate = executableCandidate + 1;
+    }
+    movedExecutableOmpAllocateCandidates.clear();
+  }
+
+  void Post(parser::ExecutionPartConstruct &x) {
+    auto *declarativeAllocate{
+        parser::Unwrap<parser::OpenMPDeclarativeAllocate>(x)};
+    auto *executableAllocate{
+        parser::Unwrap<parser::OpenMPExecutableAllocate>(x)};
+    if (declarativeAllocate) {
+      executableOmpAllocateCandidates.push_back(
+          execPartPair{x, executionPartIndex});
+
+    } else if (executableAllocate &&
+        std::get<std::optional<parser::OmpObjectList>>(executableAllocate->t)
+            .has_value() &&
+        executableOmpAllocateCandidates.size() > 0) {
+      auto &declarativeList{
+          std::get<std::optional<std::list<parser::OpenMPDeclarativeAllocate>>>(
+              executableAllocate->t)};
+      // OpenMPExecutable allocate's list of child declaratives is an optional
+      // so this list needs to be set here
+      declarativeList = std::list<parser::OpenMPDeclarativeAllocate>();
+      for (auto &executableCandidate : executableOmpAllocateCandidates) {
+        auto *declarativeCandidate{
+            parser::Unwrap<parser::OpenMPDeclarativeAllocate>(
+                *(std::get<0>(executableCandidate)))};
+        declarativeList.value().push_back(std::move(*declarativeCandidate));
+        movedExecutableOmpAllocateCandidates.push_back(
+            std::get<1>(executableCandidate));
+      }
+      executableOmpAllocateCandidates.clear();
+    } else {
+      executableOmpAllocateCandidates.clear();
+    }
+    ++executionPartIndex;
+  }
+
 private:
   template <typename T> T *GetConstructIf(parser::ExecutionPartConstruct &x) {
     if (auto *y{std::get_if<parser::ExecutableConstruct>(&x.u)}) {
       if (auto *z{std::get_if<common::Indirection<T>>(&y->u)}) {
         return &z->value();
       }
     }
     return nullptr;
if (++nextIt != block.end()) {
         return; // found do-loop
       }
     }
     messages_.Say(dir.source,
         "A DO loop must follow the %s directive"_err_en_US,
         parser::ToUpperCaseLetters(dir.source.ToString()));
   }
 
+  using execPartPair =
+      std::pair<common::Reference<parser::ExecutionPartConstruct>, int>;
+  std::list<execPartPair> executableOmpAllocateCandidates;
+  std::list<int> movedExecutableOmpAllocateCandidates;
+  int executionPartIndex{0};
+
   parser::Messages &messages_;
 };
 
 bool CanonicalizeOmp(parser::Messages &messages, parser::Program &program) {
   CanonicalizationOfOmp omp{messages};
   Walk(program, omp);
   return !messages.AnyFatalError();
 }
 } // namespace Fortran::semantics

flang/lib/Semantics/rewrite-parse-tree.cpp

 #include "flang/Common/indirection.h"
 #include "flang/Parser/parse-tree-visitor.h"
 #include "flang/Parser/parse-tree.h"
 #include "flang/Parser/tools.h"
 #include "flang/Semantics/scope.h"
 #include "flang/Semantics/semantics.h"
 #include "flang/Semantics/symbol.h"
 #include "flang/Semantics/tools.h"
+#include <iterator>
 #include <list>
 
 namespace Fortran::semantics {
 
 using namespace parser::literals;
 
 /// Convert misidentified statement functions to array element assignments.
 /// Convert misidentified format expressions to namelist group names.