[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:

I had to the same logic recently for deleting trailing `cycle`. I think this is worth refactoring into its own utility with the assertion in the util. See: https://github.com/llvm/llvm-project/blob/main/flang/lib/Lower/PFTBuilder.cpp#L723-L729.

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