[flang-commits] [flang] [flang] Rewrite an IF body ending in CYCLE as an IF/ELSE (PR #224718)

[Kareem Ergawy via flang-commits]

================
@@ -843,13 +850,117 @@ class PFTBuilder {
             std::string cycleName = getConstructName(*cycleStmt);
             if (cycleName.empty() || cycleName == doName)
               // This candidate will match doStmt's EndDoStmt.
-              ifCandidateStack.push_back({it, {}, true});
+              ifCandidateStack.push_back({it, {}, /*isCycleStmt=*/true});
           }
         }
+      } else if (doStmt && eval.isA<parser::IfConstruct>() &&
+                 eval.evaluationList->size() > 3) {
+        // An IF body whose last statement is a CYCLE.
+        lower::pft::EvaluationList &bodyList = *eval.evaluationList;
+        auto branchIt = std::prev(std::prev(bodyList.end()));
+        const auto *cycleStmt = branchIt->getIf<parser::CycleStmt>();
+        // A CYCLE ending an ELSE or ELSE IF branch sits in the same position
+        // but must not match.
+        bool hasElseBranch =
+            llvm::any_of(bodyList, [](const lower::pft::Evaluation &bodyEval) {
+              return bodyEval.isIntermediateConstructStmt();
+            });
+        if (cycleStmt && !hasElseBranch && !branchIt->label) {
+          std::string cycleName = getConstructName(*cycleStmt);
+          if (cycleName.empty() || cycleName == doName)
+            // This candidate will match doStmt's EndDoStmt.
+            ifCandidateStack.push_back({it,
+                                        {},
+                                        /*isCycleStmt=*/true,
+                                        /*hasLeadingStmts=*/true});
+        }
       }
     }
   }
 
+  /// Rewrite an IfConstruct whose body ends in a CycleStmt, with statements
+  /// ahead of it, as an IF/ELSE. The pre-branch-analysis code:
+  ///
+  ///       <<IfConstruct>>
+  ///         1 IfThenStmt: if(cond) then
+  ///         2 Statement: ...
+  ///         3 CycleStmt: cycle
+  ///         4 EndIfStmt
+  ///       <<End IfConstruct>>
+  ///       5 Statement: ...
+  ///       6 EndDoStmt
+  ///
+  /// becomes:
+  ///
+  ///       <<IfConstruct>>
+  ///         1 IfThenStmt: if(cond) then
+  ///         2 Statement: ...
+  ///         * ElseStmt
+  ///         5 Statement: ...
+  ///         4 EndIfStmt
+  ///       <<End IfConstruct>>
+  ///       6 EndDoStmt
+  ///
+  /// When the branch is the whole IF body there is nothing to keep, so
+  /// rewriteIfGotos deletes the branch and negates the condition. Here
+  /// statement 2 must still run when the condition holds, so the condition is
+  /// left as written and the statements that followed the construct go into an
+  /// ELSE branch instead.
+  ///
+  /// The synthesized ElseStmt has no source position and no index of its own.
+  template <typename FirstStmtFn>
+  void
+  rewriteIfCycleAsIfElse(lower::pft::EvaluationList &evaluationList,
+                         lower::pft::EvaluationList &ifBodyList,
+                         lower::pft::EvaluationList::iterator ifConstructIt,
+                         lower::pft::EvaluationList::iterator successorIt,
+                         lower::pft::EvaluationList::iterator it,
+                         FirstStmtFn firstStmt) {
+    lower::pft::EvaluationList::iterator branchStmtIt =
+        std::prev(std::prev(ifBodyList.end()));
+    assert(branchStmtIt->isA<parser::CycleStmt>() &&
+           "expected cycle statement");
+
+    // The last statement reachable in the THEN branch, which currently falls
+    // through to the CycleStmt. Descend through nested constructs to find it.
+    lower::pft::Evaluation *thenTail = &*std::prev(branchStmtIt);
+    while (thenTail->evaluationList && !thenTail->evaluationList->empty())
+      thenTail = &thenTail->evaluationList->back();
+    assert(thenTail->lexicalSuccessor == &*branchStmtIt &&
+           "expected fallthrough to the cycle statement");
----------------
ergawy wrote:

The following input triggers the assertion:
```fortran
subroutine s(n, v)
  integer :: n, i, v(n)
  do i = 1, n
    if (v(i) == 1) then
      v(i) = 7
100   format(I5)
      cycle
    end if
    v(i) = 2
  end do
end subroutine
```

Claude suggested the following change:
```c++
    auto inLexicalChain = [](const lower::pft::Evaluation &e) {
      return e.isActionStmt() || e.isConstructStmt() || e.isEndStmt() ||
             e.isExecutableDirective();
    };

    auto fallthroughPredecessor = [&](lower::pft::EvaluationList::iterator at) {
      do {
        --at;
      } while (!at->isConstruct() && !inLexicalChain(*at));
      lower::pft::Evaluation *tail = &*at;
      while (tail->evaluationList && !tail->evaluationList->empty())
        tail = &tail->evaluationList->back();
      return tail;
    };
```

I believe, we need to skip non-executable statements.

(Same loop is repeated below and in my changes linked above so that's a stronger reason to have a well-tested util to find the lexical predecessor).

https://github.com/llvm/llvm-project/pull/224718