[flang-commits] [flang] [flang][OpenMP] Implement loop construct iterator range (PR #170734)

[Krzysztof Parzyszek via flang-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/170734

Since we're trying to preserve compiler directives in loop constructs, not every element of the associated parser::Block needs to be a loop or an OpenMP loop construct. Implement a helper class `LoopRange` to make it easy to iterate over elements of parser::Block that are loops or loop constructs.

>From ca0b4a3ae509f29aaaaa9a67140533e3ceca544b Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Tue, 2 Dec 2025 14:34:45 -0600
Subject: [PATCH] [flang][OpenMP] Implement loop construct iterator range

Since we're trying to preserve compiler directives in loop constructs,
not every element of the associated parser::Block needs to be a loop
or an OpenMP loop construct. Implement a helper class `LoopRange` to
make it easy to iterate over elements of parser::Block that are loops
or loop constructs.
---
 flang/include/flang/Parser/openmp-utils.h | 74 +++++++++++++++++++++++
 flang/lib/Parser/openmp-utils.cpp         | 37 ++++++++++++
 2 files changed, 111 insertions(+)

diff --git a/flang/include/flang/Parser/openmp-utils.h b/flang/include/flang/Parser/openmp-utils.h
index e164f63aa189b..27e4aa6db0f0b 100644
--- a/flang/include/flang/Parser/openmp-utils.h
+++ b/flang/include/flang/Parser/openmp-utils.h
@@ -237,6 +237,80 @@ struct OmpAllocateInfo {
 
 OmpAllocateInfo SplitOmpAllocate(const OmpAllocateDirective &x);
 
+namespace detail {
+template <bool IsConst, typename T>
+struct ConstIf {
+  using type = std::conditional_t<IsConst, std::add_const_t<T>, T>;
+};
+
+template <bool IsConst, typename T>
+using ConstIfT = typename ConstIf<IsConst, T>::type;
+}
+
+template <bool IsConst> struct LoopRange {
+  using QualBlock = detail::ConstIfT<IsConst, Block>;
+  using QualReference = decltype(std::declval<QualBlock>().front());
+  using QualPointer = std::remove_reference_t<QualReference> *;
+
+  LoopRange(QualBlock &x) { Initialize(x); }
+  LoopRange(QualReference x);
+
+  LoopRange(detail::ConstIfT<IsConst, OpenMPLoopConstruct> &x)
+      : LoopRange(std::get<Block>(x.t)) {}
+  LoopRange(detail::ConstIfT<IsConst, DoConstruct> &x)
+      : LoopRange(std::get<Block>(x.t)) {}
+
+  size_t size() const { return items.size(); }
+
+  struct iterator;
+
+  iterator begin();
+  iterator end();
+
+private:
+  void Initialize(QualBlock &body);
+
+  std::vector<QualPointer> items;
+};
+
+template <typename T> LoopRange(T &x) -> LoopRange<std::is_const_v<T>>;
+
+template <bool IsConst> struct LoopRange<IsConst>::iterator {
+  QualReference operator*() { return **at; }
+
+  bool operator==(const iterator &other) const { return at == other.at; }
+  bool operator!=(const iterator &other) const { return at != other.at; }
+
+  iterator &operator++() {
+    ++at;
+    return *this;
+  }
+  iterator &operator--() {
+    --at;
+    return *this;
+  }
+  iterator &operator++(int);
+  iterator &operator--(int);
+
+private:
+  friend struct LoopRange;
+  typename decltype(LoopRange::items)::iterator at;
+};
+
+template <bool IsConst> inline auto LoopRange<IsConst>::begin() -> iterator {
+  iterator x;
+  x.at = items.begin();
+  return x;
+}
+
+template <bool IsConst> inline auto LoopRange<IsConst>::end() -> iterator {
+  iterator x;
+  x.at = items.end();
+  return x;
+}
+
+using ConstLoopRange = LoopRange<true>;
+
 } // namespace Fortran::parser::omp
 
 #endif // FORTRAN_PARSER_OPENMP_UTILS_H
diff --git a/flang/lib/Parser/openmp-utils.cpp b/flang/lib/Parser/openmp-utils.cpp
index 4c38917e87d29..099c7faf328c3 100644
--- a/flang/lib/Parser/openmp-utils.cpp
+++ b/flang/lib/Parser/openmp-utils.cpp
@@ -205,4 +205,41 @@ OmpAllocateInfo SplitOmpAllocate(const OmpAllocateDirective &x) {
   return info;
 }
 
+template <bool IsConst>
+LoopRange<IsConst>::LoopRange(QualReference x) {
+  if (auto *doLoop{Unwrap<DoConstruct>(x)}) {
+    Initialize(std::get<Block>(doLoop->t));
+  } else if (auto *omp{Unwrap<OpenMPLoopConstruct>(x)}) {
+    Initialize(std::get<Block>(omp->t));
+  }
+}
+
+template <bool IsConst>
+void LoopRange<IsConst>::Initialize(QualBlock &body) {
+  using QualIterator = decltype(std::declval<QualBlock>().begin());
+  auto makeRange{[](auto &container) {
+    return llvm::make_range(container.begin(), container.end());
+  }};
+
+  std::vector<llvm::iterator_range<QualIterator>> nest{makeRange(body)};
+  do {
+    auto at{nest.back().begin()};
+    auto end{nest.back().end()};
+    nest.pop_back();
+    while (at != end) {
+      if (auto *block{Unwrap<BlockConstruct>(*at)}) {
+        nest.push_back(llvm::make_range(std::next(at), end));
+        nest.push_back(makeRange(std::get<Block>(block->t)));
+        break;
+      } else if (Unwrap<DoConstruct>(*at) || Unwrap<OpenMPLoopConstruct>(*at)) {
+        items.push_back(&*at);
+      }
+      ++at;
+    }
+  } while (!nest.empty());
+}
+
+template struct LoopRange<false>;
+template struct LoopRange<true>;
+
 } // namespace Fortran::parser::omp