[flang][OpenMP] Implement loop construct iterator range

[kparzysz]

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.

[llvmbot]

@llvm/pr-subscribers-flang-parser

Author: Krzysztof Parzyszek (kparzysz)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/170734.diff

2 Files Affected:

• (modified) flang/include/flang/Parser/openmp-utils.h (+74)
• (modified) flang/lib/Parser/openmp-utils.cpp (+37)

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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[kparzysz]

PR stack:

1. [flang][OpenMP] Implement loop construct iterator range #170734 (this PR)
2. [flang][OpenMP] Generalize checks of loop construct structure #170735

[github-actions]

🐧 Linux x64 Test Results

• 4100 tests passed
• 201 tests skipped

✅ The build succeeded and all tests passed.