[flang-commits] [flang] [flang][openacc] Avoid crash when collapse loop nest has extra directive (PR #73166)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/73166

None

>From 3e2a5b631d97d586259cd4d5ca5a852b282bbfa7 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 22 Nov 2023 12:14:32 -0800
Subject: [PATCH] [flang][openacc] Avoid crash when collapse loop nest has
 extra directive

---
 flang/lib/Semantics/resolve-directives.cpp | 17 ++++++++++++++---
 flang/test/Semantics/OpenACC/acc-loop.f90  |  8 ++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index bbb105e3516da18..45329b5cca9c03b 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1183,13 +1183,23 @@ void AccAttributeVisitor::CheckAssociatedLoopIndex(
   }
 
   const auto getNextDoConstruct =
-      [this](const parser::Block &block) -> const parser::DoConstruct * {
+      [this](const parser::Block &block,
+          std::int64_t &level) -> const parser::DoConstruct * {
     for (const auto &entry : block) {
       if (const auto *doConstruct = GetDoConstructIf(entry)) {
         return doConstruct;
       } else if (parser::Unwrap<parser::CompilerDirective>(entry)) {
         // It is allowed to have a compiler directive associated with the loop.
         continue;
+      } else if (const auto &accLoop{
+                     parser::Unwrap<parser::OpenACCLoopConstruct>(entry)}) {
+        if (level == 0)
+          break;
+        const auto &beginDir{
+            std::get<parser::AccBeginLoopDirective>(accLoop->t)};
+        context_.Say(beginDir.source,
+            "LOOP directive not expected in COLLAPSE loop nest"_err_en_US);
+        level = 0;
       } else {
         break;
       }
@@ -1198,11 +1208,12 @@ void AccAttributeVisitor::CheckAssociatedLoopIndex(
   };
 
   const auto &outer{std::get<std::optional<parser::DoConstruct>>(x.t)};
-  for (const parser::DoConstruct *loop{&*outer}; loop && level > 0; --level) {
+  for (const parser::DoConstruct *loop{&*outer}; loop && level > 0;) {
     // Go through all nested loops to ensure index variable exists.
     GetLoopIndex(*loop);
     const auto &block{std::get<parser::Block>(loop->t)};
-    loop = getNextDoConstruct(block);
+    --level;
+    loop = getNextDoConstruct(block, level);
   }
   CHECK(level == 0);
 }
diff --git a/flang/test/Semantics/OpenACC/acc-loop.f90 b/flang/test/Semantics/OpenACC/acc-loop.f90
index 1c0a515441f432b..ee1e6e0cce09740 100644
--- a/flang/test/Semantics/OpenACC/acc-loop.f90
+++ b/flang/test/Semantics/OpenACC/acc-loop.f90
@@ -268,4 +268,12 @@ program openacc_loop_validity
   end do
   !$acc end loop
 
+  !$acc loop collapse(2)
+  do i = 1, 10
+    !ERROR: LOOP directive not expected in COLLAPSE loop nest
+    !$acc loop
+    do j = 1, 10
+    end do
+  end do
+
 end program openacc_loop_validity