[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

@@ -160,8 +160,6 @@ class PFTBuilder {
       exitFunction();
     } else if constexpr (lower::pft::isConstruct<A> ||
                          lower::pft::isDirective<A>) {
-      if constexpr (lower::pft::isDeclConstruct<A>)
-        return;
       exitConstructOrDirective();
     }
   }
@@ -245,11 +243,6 @@ class PFTBuilder {
     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 @@ class PFTBuilder {
     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 @@ class PFTBuilder {
     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;
   }
 

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
+