[flang][OpenMP] Add support for getting parent context and clauses and modify the relevant sema checks to use the new functions.

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

GitOrigin-RevId: 3fc9ca87f1ab7c8e342cb942729df30f1ba80418
diff --git a/lib/Semantics/check-directive-structure.h b/lib/Semantics/check-directive-structure.h
index 3ff4462..6ef6023 100644
--- a/lib/Semantics/check-directive-structure.h
+++ b/lib/Semantics/check-directive-structure.h
@@ -150,6 +150,11 @@
     return dirContext_.back();
   }
 
+  DirectiveContext &GetContextParent() {
+    CHECK(dirContext_.size() >= 2);
+    return dirContext_[dirContext_.size() - 2];
+  }
+
   void SetContextClause(const PC &clause) {
     GetContext().clauseSource = clause.source;
     GetContext().clause = &clause;
@@ -198,6 +203,7 @@
     GetContext().actualClauses.push_back(type);
   }
 
+  // Check if the given clause is present in the current context
   const PC *FindClause(C type) {
     auto it{GetContext().clauseInfo.find(type)};
     if (it != GetContext().clauseInfo.end()) {
@@ -206,6 +212,15 @@
     return nullptr;
   }
 
+  // Check if the given clause is present in the parent context
+  const PC *FindClauseParent(C type) {
+    auto it{GetContextParent().clauseInfo.find(type)};
+    if (it != GetContextParent().clauseInfo.end()) {
+      return it->second;
+    }
+    return nullptr;
+  }
+
   std::pair<typename ClauseMapTy::iterator, typename ClauseMapTy::iterator>
   FindClauses(C type) {
     auto it{GetContext().clauseInfo.equal_range(type)};
diff --git a/lib/Semantics/check-omp-structure.cpp b/lib/Semantics/check-omp-structure.cpp
index e8b632f..b23ae51 100644
--- a/lib/Semantics/check-omp-structure.cpp
+++ b/lib/Semantics/check-omp-structure.cpp
@@ -131,7 +131,7 @@
     const parser::CharBlock &source, const OmpDirectiveSet &set) {
   // set contains all the invalid closely nested directives
   // for the given directive (`source` here)
-  if (CurrentDirectiveIsNested() && set.test(GetContext().directive)) {
+  if (CurrentDirectiveIsNested() && set.test(GetContextParent().directive)) {
     context_.Say(source,
         "A worksharing region may not be closely nested inside a "
         "worksharing, explicit task, taskloop, critical, ordered, atomic, or "
@@ -158,9 +158,9 @@
     CheckMatching<parser::OmpLoopDirective>(beginDir, endDir);
   }
 
-  if (beginDir.v != llvm::omp::Directive::OMPD_do) {
-    PushContextAndClauseSets(beginDir.source, beginDir.v);
-  } else {
+  PushContextAndClauseSets(beginDir.source, beginDir.v);
+
+  if (beginDir.v == llvm::omp::Directive::OMPD_do) {
     // 2.7.1 do-clause -> private-clause |
     //                    firstprivate-clause |
     //                    lastprivate-clause |
@@ -181,7 +181,6 @@
             llvm::omp::Directive::OMPD_ordered,
             llvm::omp::Directive::OMPD_atomic,
             llvm::omp::Directive::OMPD_master});
-    PushContextAndClauseSets(beginDir.source, llvm::omp::Directive::OMPD_do);
   }
   SetLoopInfo(x);
 
@@ -318,15 +317,17 @@
 
   CheckMatching<parser::OmpBlockDirective>(beginDir, endDir);
 
+  PushContextAndClauseSets(beginDir.source, beginDir.v);
+
   // TODO: This check needs to be extended while implementing nesting of regions
   // checks.
   if (beginDir.v == llvm::omp::Directive::OMPD_single) {
     HasInvalidWorksharingNesting(
         beginDir.source, {llvm::omp::Directive::OMPD_do});
   }
-  CheckIfDoOrderedClause(beginDir);
+  if (CurrentDirectiveIsNested())
+    CheckIfDoOrderedClause(beginDir);
 
-  PushContextAndClauseSets(beginDir.source, beginDir.v);
   CheckNoBranching(block, beginDir.v, beginDir.source);
 
   switch (beginDir.v) {
@@ -342,7 +343,7 @@
 void OmpStructureChecker::CheckIfDoOrderedClause(
     const parser::OmpBlockDirective &blkDirective) {
   if (blkDirective.v == llvm::omp::OMPD_ordered) {
-    if (!FindClause(llvm::omp::Clause::OMPC_ordered)) {
+    if (!FindClauseParent(llvm::omp::Clause::OMPC_ordered)) {
       context_.Say(blkDirective.source,
           "The ORDERED clause must be present on the loop"
           " construct if any ORDERED region ever binds"