From f2ac513812e197758e29354c190004eb7bc2f8b6 Mon Sep 17 00:00:00 2001
From: Peixin-Qiao <qiaopeixin@huawei.com>
Date: Fri, 11 Mar 2022 15:20:23 +0800
Subject: [PATCH] [flang] Fix processing ModuleLikeUnit evaluationList

Push the ModuleLikeUnit evalutionList when entering module unit. Pop it
when exiting module unit if there is no module procedure. Otherwise, pop
it when entering the first module procedure.

Reviewed By: V Donaldson

Differential Revision: https://reviews.llvm.org/D120460
---
 flang/lib/Lower/PFTBuilder.cpp      | 23 +++++++++-----------
 flang/test/Lower/pre-fir-tree06.f90 | 42 +++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/flang/lib/Lower/PFTBuilder.cpp b/flang/lib/Lower/PFTBuilder.cpp
index b204d4a..4d80ff5 100644
--- a/flang/lib/Lower/PFTBuilder.cpp
+++ b/flang/lib/Lower/PFTBuilder.cpp
@@ -160,8 +160,6 @@ public:
       exitFunction();
     } else if constexpr (lower::pft::isConstruct<A> ||
                          lower::pft::isDirective<A>) {
-      if constexpr (lower::pft::isDeclConstruct<A>)
-        return;
       exitConstructOrDirective();
     }
   }
@@ -245,11 +243,6 @@ private:
     if (evaluationListStack.empty())
       return;
     auto evaluationList = evaluationListStack.back();
-    if (evaluationList->empty() &&
-        pftParentStack.back().getIf<lower::pft::ModuleLikeUnit>()) {
-      popEvaluationList();
-      return;
-    }
     if (evaluationList->empty() || !evaluationList->back().isEndStmt()) {
       const auto &endStmt =
           pftParentStack.back().get<lower::pft::FunctionLikeUnit>().endStmt;
@@ -279,10 +272,20 @@ private:
     lastLexicalEvaluation = nullptr;
   }
 
+  /// Pop the ModuleLikeUnit evaluationList when entering the first module
+  /// procedure.
+  void cleanModuleEvaluationList() {
+    if (evaluationListStack.empty())
+      return;
+    if (pftParentStack.back().isA<lower::pft::ModuleLikeUnit>())
+      popEvaluationList();
+  }
+
   /// Initialize a new function-like unit and make it the builder's focus.
   template <typename A>
   bool enterFunction(const A &func,
                      const semantics::SemanticsContext &semanticsContext) {
+    cleanModuleEvaluationList();
     endFunctionBody(); // enclosing host subprogram body, if any
     Fortran::lower::pft::FunctionLikeUnit &unit =
         addFunction(lower::pft::FunctionLikeUnit{func, pftParentStack.back(),
@@ -316,12 +319,6 @@ private:
     pushEvaluationList(eval.evaluationList.get());
     pftParentStack.emplace_back(eval);
     constructAndDirectiveStack.emplace_back(&eval);
-    if constexpr (lower::pft::isDeclConstruct<A>) {
-      popEvaluationList();
-      pftParentStack.pop_back();
-      constructAndDirectiveStack.pop_back();
-      popEvaluationList();
-    }
     return true;
   }
 
diff --git a/flang/test/Lower/pre-fir-tree06.f90 b/flang/test/Lower/pre-fir-tree06.f90
index b79a55a..847da6e 100644
--- a/flang/test/Lower/pre-fir-tree06.f90
+++ b/flang/test/Lower/pre-fir-tree06.f90
@@ -10,3 +10,45 @@ module m
 end
 ! CHECK: End ModuleLike
 
+! CHECK: ModuleLike
+module m2
+  integer, save :: i
+  ! CHECK-NEXT: OpenMPDeclarativeConstruct
+  !$omp threadprivate(i)
+contains
+  subroutine sub()
+    i = 1;
+  end
+  subroutine sub2()
+    i = 2;
+  end
+end
+! CHECK: End ModuleLike
+
+! CHECK: Program main
+program main
+  real :: y
+  ! CHECK-NEXT: OpenMPDeclarativeConstruct
+  !$omp threadprivate(y)
+end
+! CHECK: End Program main
+
+! CHECK: Subroutine sub1
+subroutine sub1()
+  real, save :: p
+  ! CHECK-NEXT: OpenMPDeclarativeConstruct
+  !$omp threadprivate(p)
+end
+! CHECK: End Subroutine sub1
+
+! CHECK: Subroutine sub2
+subroutine sub2()
+  real, save :: q
+  ! CHECK-NEXT: OpenMPDeclarativeConstruct
+  !$omp threadprivate(q)
+contains
+  subroutine sub()
+  end
+end
+! CHECK: End Subroutine sub2
+
-- 
2.7.4