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>>(associate.value().t));
35                 },
36                 [&](common::Indirection<BlockConstruct> &blockConstruct) {
37           CanonicalizeIfMatch(block, stack, i,
38               std::get<Statement<EndBlockStmt>>(blockConstruct.value().t));
39                 },
40                 [&](common::Indirection<ChangeTeamConstruct> &changeTeam) {
41           CanonicalizeIfMatch(block, stack, i,
42               std::get<Statement<EndChangeTeamStmt>>(changeTeam.value().t));
43                 },
44                 [&](common::Indirection<CriticalConstruct> &critical) {
45           CanonicalizeIfMatch(block, stack, i,
46               std::get<Statement<EndCriticalStmt>>(critical.value().t));
47                 },
48                 [&](common::Indirection<DoConstruct> &doConstruct) {
49           CanonicalizeIfMatch(block, stack, i,
50               std::get<Statement<EndDoStmt>>(doConstruct.value().t));
51                 },
52                 [&](common::Indirection<IfConstruct> &ifConstruct) {
53           CanonicalizeIfMatch(block, stack, i,
54               std::get<Statement<EndIfStmt>>(ifConstruct.value().t));
55                 },
56                 [&](common::Indirection<CaseConstruct> &caseConstruct) {
57           CanonicalizeIfMatch(block, stack, i,
58               std::get<Statement<EndSelectStmt>>(caseConstruct.value().t));
59                 },
60                 [&](common::Indirection<SelectRankConstruct> &selectRank) {
61           CanonicalizeIfMatch(block, stack, i,
62               std::get<Statement<EndSelectStmt>>(selectRank.value().t));
63                 },
64                 [&](common::Indirection<SelectTypeConstruct> &selectType) {
65           CanonicalizeIfMatch(block, stack, i,
66               std::get<Statement<EndSelectStmt>>(selectType.value().t));
67                 },
68                 [&](common::Indirection<ForallConstruct> &forall) {
69           CanonicalizeIfMatch(block, stack, i,
70               std::get<Statement<EndForallStmt>>(forall.value().t));
71                 },
72                 [&](common::Indirection<WhereConstruct> &where) {
73                   CanonicalizeIfMatch(block, stack, i,
74                       std::get<Statement<EndWhereStmt>>(where.value().t));
75                 },
76                 [&](Statement<common::Indirection<LabelDoStmt>> &labelDoStmt) {
77           auto &label{std::get<Label>(labelDoStmt.statement.value().t)};
78           stack.push_back(LabelInfo{i, label});
79                 },
80                 [&](Statement<common::Indirection<EndDoStmt>> &endDoStmt) {
81           CanonicalizeIfMatch(block, stack, i, endDoStmt);
82                 },
83                 [&](Statement<ActionStmt> &actionStmt) {
84           CanonicalizeIfMatch(block, stack, i, actionStmt);
85                 },
86             },
87             executableConstruct->u);
88       }
89     }
90   }
91 
92 private:
93   template <typename T>
94   void CanonicalizeIfMatch(Block &originalBlock, std::vector<LabelInfo> &stack,
95       Block::iterator &i, Statement<T> &statement) {
96     if (!stack.empty() && statement.label &&
97         stack.back().label == *statement.label) {
98       auto currentLabel{stack.back().label};
99       if constexpr (std::is_same_v<T, common::Indirection<EndDoStmt>>) {
100         std::get<ExecutableConstruct>(i->u).u = Statement<ActionStmt>{
101             std::optional<Label>{currentLabel}, ContinueStmt{}};
102       }
103       auto next{++i};
104       do {
105         Block block;
106         auto doLoop{stack.back().iter};
107         auto originalSource{
108             std::get<Statement<common::Indirection<LabelDoStmt>>>(
109                 std::get<ExecutableConstruct>(doLoop->u).u)
110                 .source};
111         block.splice(block.begin(), originalBlock, ++stack.back().iter, next);
112         auto &labelDo{std::get<Statement<common::Indirection<LabelDoStmt>>>(
113             std::get<ExecutableConstruct>(doLoop->u).u)};
114         auto &loopControl{
115             std::get<std::optional<LoopControl>>(labelDo.statement.value().t)};
116         auto &name{std::get<std::optional<Name>>(labelDo.statement.value().t)};
117         Statement<NonLabelDoStmt> nonLabelDoStmt{std::move(labelDo.label),
118             NonLabelDoStmt{
119                 std::make_tuple(common::Clone(name), std::move(loopControl))}};
120         nonLabelDoStmt.source = originalSource;
121         std::get<ExecutableConstruct>(doLoop->u).u =
122             common::Indirection<DoConstruct>{
123                 std::make_tuple(std::move(nonLabelDoStmt), std::move(block),
124                     Statement<EndDoStmt>{
125                         std::optional<Label>{}, EndDoStmt{std::move(name)}})};
126         stack.pop_back();
127       } while (!stack.empty() && stack.back().label == currentLabel);
128       i = --next;
129     }
130   }
131 };
132 
133 bool CanonicalizeDo(Program &program) {
134   CanonicalizationOfDoLoops canonicalizationOfDoLoops;
135   Walk(program, canonicalizationOfDoLoops);
136   return true;
137 }
138 
139 } // namespace Fortran::parser
140