[llvm-branch-commits] [flang] [flang][OpenMP] Remember original range in ExecutionPartIterator (PR #185290)

[Krzysztof Parzyszek via llvm-branch-commits]

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

Storing the original range (instead of just the "remaining part") will allow the iterator component to be reused.

>From 77235406105e7e686f1ca19bbfb4956aeef9257e Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Wed, 18 Feb 2026 10:53:13 -0600
Subject: [PATCH] [flang][OpenMP] Remember original range in
 ExecutionPartIterator

Storing the original range (instead of just the "remaining part")
will allow the iterator component to be reused.
---
 flang/include/flang/Parser/openmp-utils.h | 25 ++++++++++++++++++-----
 flang/lib/Parser/openmp-utils.cpp         | 11 ++++------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/flang/include/flang/Parser/openmp-utils.h b/flang/include/flang/Parser/openmp-utils.h
index f23e52585d567..20754ad28d26d 100644
--- a/flang/include/flang/Parser/openmp-utils.h
+++ b/flang/include/flang/Parser/openmp-utils.h
@@ -296,14 +296,28 @@ struct ExecutionPartIterator {
   using IteratorType = Block::const_iterator;
   using IteratorRange = llvm::iterator_range<IteratorType>;
 
+  // An iterator range with a third iterator indicating a position inside
+  // the range.
+  struct IteratorGauge : public IteratorRange {
+    IteratorGauge(IteratorType b, IteratorType e)
+        : IteratorRange(b, e), at(b) {}
+    IteratorGauge(IteratorRange r) : IteratorRange(r), at(r.begin()) {}
+
+    bool atEnd() const { return at == end(); }
+    IteratorType at;
+  };
+
   struct Construct {
     Construct(IteratorType b, IteratorType e, const ExecutionPartConstruct *c)
-        : range(b, e), owner(c) {}
+        : location(b, e), owner(c) {}
     template <typename R>
     Construct(const R &r, const ExecutionPartConstruct *c)
-        : range(r), owner(c) {}
+        : location(r), owner(c) {}
     Construct(const Construct &c) = default;
-    IteratorRange range;
+    // The original range of the construct with the current position in it.
+    // The location.at is the construct currently being pointed at, or
+    // stepped into.
+    IteratorGauge location;
     const ExecutionPartConstruct *owner;
   };
 
@@ -332,6 +346,7 @@ struct ExecutionPartIterator {
 
   bool valid() const { return !stack_.empty(); }
 
+  const std::vector<Construct> &stack() const { return stack_; }
   decltype(auto) operator*() const { return *at(); }
   bool operator==(const ExecutionPartIterator &other) const {
     if (valid() != other.valid()) {
@@ -339,7 +354,7 @@ struct ExecutionPartIterator {
     }
     // Invalid iterators are considered equal.
     return !valid() ||
-        stack_.back().range.begin() == other.stack_.back().range.begin();
+        stack_.back().location.at == other.stack_.back().location.at;
   }
   bool operator!=(const ExecutionPartIterator &other) const {
     return !(*this == other);
@@ -368,7 +383,7 @@ struct ExecutionPartIterator {
   using iterator_category = std::forward_iterator_tag;
 
 private:
-  IteratorType at() const { return stack_.back().range.begin(); };
+  IteratorType at() const { return stack_.back().location.at; };
 
   // If the iterator is not at a legal location, keep advancing it until
   // it lands at a legal location or becomes invalid.
diff --git a/flang/lib/Parser/openmp-utils.cpp b/flang/lib/Parser/openmp-utils.cpp
index 6d4326af78344..f83a658104396 100644
--- a/flang/lib/Parser/openmp-utils.cpp
+++ b/flang/lib/Parser/openmp-utils.cpp
@@ -243,8 +243,7 @@ void ExecutionPartIterator::step() {
     } else if (auto *loop{GetDoConstruct(*where)}) {
       stack_.emplace_back(std::get<Block>(loop->t), &*where);
     } else {
-      stack_.back().range =
-          IteratorRange(std::next(where), stack_.back().range.end());
+      ++stack_.back().location.at;
     }
     adjust();
   }
@@ -254,8 +253,7 @@ void ExecutionPartIterator::next() {
   // Advance the iterator to the next legal position. If the current
   // position is a DO-loop or a loop construct, step over it.
   if (valid()) {
-    stack_.back().range =
-        IteratorRange(std::next(at()), stack_.back().range.end());
+    ++stack_.back().location.at;
     adjust();
   }
 }
@@ -264,11 +262,10 @@ void ExecutionPartIterator::adjust() {
   // If the iterator is not at a legal location, keep advancing it until
   // it lands at a legal location or becomes invalid.
   while (valid()) {
-    if (stack_.back().range.empty()) {
+    if (stack_.back().location.atEnd()) {
       stack_.pop_back();
       if (valid()) {
-        stack_.back().range =
-            IteratorRange(std::next(at()), stack_.back().range.end());
+        ++stack_.back().location.at;
       }
     } else if (auto *block{GetFortranBlockConstruct(*at())}) {
       stack_.emplace_back(std::get<Block>(block->t), &*at());