1#     Copyright 2021, Kay Hayen, mailto:kay.hayen@gmail.com
2#
3#     Part of "Nuitka", an optimizing Python compiler that is compatible and
4#     integrates with CPython, but also works on its own.
5#
6#     Licensed under the Apache License, Version 2.0 (the "License");
7#     you may not use this file except in compliance with the License.
8#     You may obtain a copy of the License at
9#
10#        http://www.apache.org/licenses/LICENSE-2.0
11#
12#     Unless required by applicable law or agreed to in writing, software
13#     distributed under the License is distributed on an "AS IS" BASIS,
14#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15#     See the License for the specific language governing permissions and
16#     limitations under the License.
17#
18""" Reformulation of for loop statements.
19
20Consult the developer manual for information. TODO: Add ability to sync
21source code comments with developer manual sections.
22
23"""
24
25from nuitka.nodes.AssignNodes import (
26    StatementAssignmentVariable,
27    StatementReleaseVariable,
28)
29from nuitka.nodes.BuiltinIteratorNodes import (
30    ExpressionAsyncIter,
31    ExpressionAsyncNext,
32    ExpressionBuiltinIter1,
33)
34from nuitka.nodes.BuiltinNextNodes import ExpressionBuiltinNext1
35from nuitka.nodes.ComparisonNodes import ExpressionComparisonIs
36from nuitka.nodes.ConditionalNodes import makeStatementConditional
37from nuitka.nodes.ConstantRefNodes import makeConstantRefNode
38from nuitka.nodes.LoopNodes import StatementLoop, StatementLoopBreak
39from nuitka.nodes.StatementNodes import StatementsSequence
40from nuitka.nodes.VariableRefNodes import ExpressionTempVariableRef
41from nuitka.nodes.YieldNodes import ExpressionYieldFromWaitable
42
43from .ReformulationAssignmentStatements import buildAssignmentStatements
44from .ReformulationTryExceptStatements import makeTryExceptSingleHandlerNode
45from .ReformulationTryFinallyStatements import makeTryFinallyStatement
46from .TreeHelpers import (
47    buildNode,
48    buildStatementsNode,
49    makeStatementsSequence,
50    makeStatementsSequenceFromStatements,
51    popBuildContext,
52    pushBuildContext,
53)
54
55
56def _buildForLoopNode(provider, node, sync, source_ref):
57    # The for loop is re-formulated according to developer manual. An iterator
58    # is created, and looped until it gives StopIteration. The else block is
59    # taken if a for loop exits normally, i.e. because of iterator
60    # exhaustion. We do this by introducing an indicator variable.
61
62    # We handle async and sync both here, leading to cases, pylint: disable=too-many-locals
63
64    source = buildNode(provider, node.iter, source_ref)
65
66    # Temporary variables, we need one for the iterator, and one for the current
67    # value.
68    temp_scope = provider.allocateTempScope("for_loop")
69
70    tmp_iter_variable = provider.allocateTempVariable(
71        temp_scope=temp_scope, name="for_iterator"
72    )
73    tmp_value_variable = provider.allocateTempVariable(
74        temp_scope=temp_scope, name="iter_value"
75    )
76
77    else_block = buildStatementsNode(
78        provider=provider,
79        nodes=node.orelse if node.orelse else None,
80        source_ref=source_ref,
81    )
82
83    if else_block is not None:
84        # Indicator variable, will end up with C bool type, and need not be released.
85        tmp_break_indicator = provider.allocateTempVariable(
86            temp_scope=temp_scope, name="break_indicator", temp_type="bool"
87        )
88
89        statements = [
90            StatementAssignmentVariable(
91                variable=tmp_break_indicator,
92                source=makeConstantRefNode(constant=True, source_ref=source_ref),
93                source_ref=source_ref,
94            )
95        ]
96    else:
97        statements = []
98
99    statements.append(StatementLoopBreak(source_ref=source_ref))
100
101    handler_body = makeStatementsSequence(
102        statements=statements, allow_none=False, source_ref=source_ref
103    )
104
105    if sync:
106        next_node = ExpressionBuiltinNext1(
107            value=ExpressionTempVariableRef(
108                variable=tmp_iter_variable, source_ref=source_ref
109            ),
110            source_ref=source_ref,
111        )
112    else:
113        next_node = ExpressionYieldFromWaitable(
114            expression=ExpressionAsyncNext(
115                value=ExpressionTempVariableRef(
116                    variable=tmp_iter_variable, source_ref=source_ref
117                ),
118                source_ref=source_ref,
119            ),
120            source_ref=source_ref,
121        )
122
123    statements = (
124        makeTryExceptSingleHandlerNode(
125            tried=StatementAssignmentVariable(
126                variable=tmp_value_variable, source=next_node, source_ref=source_ref
127            ),
128            exception_name="StopIteration" if sync else "StopAsyncIteration",
129            handler_body=handler_body,
130            source_ref=source_ref,
131        ),
132        buildAssignmentStatements(
133            provider=provider,
134            node=node.target,
135            source=ExpressionTempVariableRef(
136                variable=tmp_value_variable, source_ref=source_ref
137            ),
138            source_ref=source_ref,
139        ),
140    )
141
142    pushBuildContext("loop_body")
143    statements += (
144        buildStatementsNode(provider=provider, nodes=node.body, source_ref=source_ref),
145    )
146    popBuildContext()
147
148    loop_body = makeStatementsSequence(
149        statements=statements, allow_none=True, source_ref=source_ref
150    )
151
152    cleanup_statements = [
153        StatementReleaseVariable(variable=tmp_value_variable, source_ref=source_ref),
154        StatementReleaseVariable(variable=tmp_iter_variable, source_ref=source_ref),
155    ]
156
157    if else_block is not None:
158        statements = [
159            StatementAssignmentVariable(
160                variable=tmp_break_indicator,
161                source=makeConstantRefNode(constant=False, source_ref=source_ref),
162                source_ref=source_ref,
163            )
164        ]
165    else:
166        statements = []
167
168    if sync:
169        iter_source = ExpressionBuiltinIter1(
170            value=source, source_ref=source.getSourceReference()
171        )
172    else:
173        iter_source = ExpressionYieldFromWaitable(
174            expression=ExpressionAsyncIter(
175                value=source, source_ref=source.getSourceReference()
176            ),
177            source_ref=source.getSourceReference(),
178        )
179
180    statements += (
181        # First create the iterator and store it.
182        StatementAssignmentVariable(
183            variable=tmp_iter_variable, source=iter_source, source_ref=source_ref
184        ),
185        makeTryFinallyStatement(
186            provider=provider,
187            tried=StatementLoop(loop_body=loop_body, source_ref=source_ref),
188            final=StatementsSequence(
189                statements=cleanup_statements, source_ref=source_ref
190            ),
191            source_ref=source_ref,
192        ),
193    )
194
195    if else_block is not None:
196        statements.append(
197            makeStatementConditional(
198                condition=ExpressionComparisonIs(
199                    left=ExpressionTempVariableRef(
200                        variable=tmp_break_indicator, source_ref=source_ref
201                    ),
202                    right=makeConstantRefNode(constant=True, source_ref=source_ref),
203                    source_ref=source_ref,
204                ),
205                yes_branch=else_block,
206                no_branch=None,
207                source_ref=source_ref,
208            )
209        )
210
211    return makeStatementsSequenceFromStatements(*statements)
212
213
214def buildForLoopNode(provider, node, source_ref):
215    return _buildForLoopNode(provider, node, True, source_ref)
216
217
218def buildAsyncForLoopNode(provider, node, source_ref):
219    return _buildForLoopNode(provider, node, False, source_ref)
220