1 //===-- lib/Semantics/canonicalize-do.cpp ---------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "canonicalize-do.h"
10 #include "flang/Parser/parse-tree-visitor.h"
11 
12 namespace Fortran::parser {
13 
14 class CanonicalizationOfDoLoops {
15   struct LabelInfo {
16     Block::iterator iter;
17     Label label;
18   };
19 
20 public:
21   template <typename T> bool Pre(T &) { return true; }
22   template <typename T> void Post(T &) {}
23   void Post(Block &block) {
24     std::vector<LabelInfo> stack;
25     for (auto i{block.begin()}, end{block.end()}; i != end; ++i) {
26       if (auto *executableConstruct{std::get_if<ExecutableConstruct>(&i->u)}) {
27         std::visit(
28             common::visitors{
29                 [](auto &) {},
30                 // Labels on end-stmt of constructs are accepted by f18 as an
31                 // extension.
32                 [&](common::Indirection<AssociateConstruct> &associate) {
33                   CanonicalizeIfMatch(block, stack, i,
34                       std::get<Statement<EndAssociateStmt>>(
35                           associate.value().t));
36                 },
37                 [&](common::Indirection<BlockConstruct> &blockConstruct) {
38                   CanonicalizeIfMatch(block, stack, i,
39                       std::get<Statement<EndBlockStmt>>(
40                           blockConstruct.value().t));
41                 },
42                 [&](common::Indirection<ChangeTeamConstruct> &changeTeam) {
43                   CanonicalizeIfMatch(block, stack, i,
44                       std::get<Statement<EndChangeTeamStmt>>(
45                           changeTeam.value().t));
46                 },
47                 [&](common::Indirection<CriticalConstruct> &critical) {
48                   CanonicalizeIfMatch(block, stack, i,
49                       std::get<Statement<EndCriticalStmt>>(critical.value().t));
50                 },
51                 [&](common::Indirection<DoConstruct> &doConstruct) {
52                   CanonicalizeIfMatch(block, stack, i,
53                       std::get<Statement<EndDoStmt>>(doConstruct.value().t));
54                 },
55                 [&](common::Indirection<IfConstruct> &ifConstruct) {
56                   CanonicalizeIfMatch(block, stack, i,
57                       std::get<Statement<EndIfStmt>>(ifConstruct.value().t));
58                 },
59                 [&](common::Indirection<CaseConstruct> &caseConstruct) {
60                   CanonicalizeIfMatch(block, stack, i,
61                       std::get<Statement<EndSelectStmt>>(
62                           caseConstruct.value().t));
63                 },
64                 [&](common::Indirection<SelectRankConstruct> &selectRank) {
65                   CanonicalizeIfMatch(block, stack, i,
66                       std::get<Statement<EndSelectStmt>>(selectRank.value().t));
67                 },
68                 [&](common::Indirection<SelectTypeConstruct> &selectType) {
69                   CanonicalizeIfMatch(block, stack, i,
70                       std::get<Statement<EndSelectStmt>>(selectType.value().t));
71                 },
72                 [&](common::Indirection<ForallConstruct> &forall) {
73                   CanonicalizeIfMatch(block, stack, i,
74                       std::get<Statement<EndForallStmt>>(forall.value().t));
75                 },
76                 [&](common::Indirection<WhereConstruct> &where) {
77                   CanonicalizeIfMatch(block, stack, i,
78                       std::get<Statement<EndWhereStmt>>(where.value().t));
79                 },
80                 [&](Statement<common::Indirection<LabelDoStmt>> &labelDoStmt) {
81                   auto &label{std::get<Label>(labelDoStmt.statement.value().t)};
82                   stack.push_back(LabelInfo{i, label});
83                 },
84                 [&](Statement<common::Indirection<EndDoStmt>> &endDoStmt) {
85                   CanonicalizeIfMatch(block, stack, i, endDoStmt);
86                 },
87                 [&](Statement<ActionStmt> &actionStmt) {
88                   CanonicalizeIfMatch(block, stack, i, actionStmt);
89                 },
90             },
91             executableConstruct->u);
92       }
93     }
94   }
95 
96 private:
97   template <typename T>
98   void CanonicalizeIfMatch(Block &originalBlock, std::vector<LabelInfo> &stack,
99       Block::iterator &i, Statement<T> &statement) {
100     if (!stack.empty() && statement.label &&
101         stack.back().label == *statement.label) {
102       auto currentLabel{stack.back().label};
103       if constexpr (std::is_same_v<T, common::Indirection<EndDoStmt>>) {
104         std::get<ExecutableConstruct>(i->u).u = Statement<ActionStmt>{
105             std::optional<Label>{currentLabel}, ContinueStmt{}};
106       }
107       auto next{++i};
108       do {
109         Block block;
110         auto doLoop{stack.back().iter};
111         auto originalSource{
112             std::get<Statement<common::Indirection<LabelDoStmt>>>(
113                 std::get<ExecutableConstruct>(doLoop->u).u)
114                 .source};
115         block.splice(block.begin(), originalBlock, ++stack.back().iter, next);
116         auto &labelDo{std::get<Statement<common::Indirection<LabelDoStmt>>>(
117             std::get<ExecutableConstruct>(doLoop->u).u)};
118         auto &loopControl{
119             std::get<std::optional<LoopControl>>(labelDo.statement.value().t)};
120         auto &name{std::get<std::optional<Name>>(labelDo.statement.value().t)};
121         Statement<NonLabelDoStmt> nonLabelDoStmt{std::move(labelDo.label),
122             NonLabelDoStmt{
123                 std::make_tuple(common::Clone(name), std::move(loopControl))}};
124         nonLabelDoStmt.source = originalSource;
125         std::get<ExecutableConstruct>(doLoop->u).u =
126             common::Indirection<DoConstruct>{
127                 std::make_tuple(std::move(nonLabelDoStmt), std::move(block),
128                     Statement<EndDoStmt>{
129                         std::optional<Label>{}, EndDoStmt{std::move(name)}})};
130         stack.pop_back();
131       } while (!stack.empty() && stack.back().label == currentLabel);
132       i = --next;
133     }
134   }
135 };
136 
137 bool CanonicalizeDo(Program &program) {
138   CanonicalizationOfDoLoops canonicalizationOfDoLoops;
139   Walk(program, canonicalizationOfDoLoops);
140   return true;
141 }
142 
143 } // namespace Fortran::parser
144