Skip to content

Commit 978b711

Browse files
VaggelisDVaggelisD
authored
[mypyc] Fix range loop variable off-by-one after loop exit (#21098)
Fixes mypyc/mypyc#1191 Previously, ForRange.gen_step() updated both the internal index register and the user-visible loop variable after incrementing. This meant the loop variable was set to the incremented value before the condition check could reject it, causing an off-by-one overshoot on loop exit. - Before: ``` ┌─────────────────────────────────────┐ │ L1: condition │ │ if index_reg < end → L2 else L4 │ └──────────────┬──────────────────┬───┘ │ true │ false ▼ │ ┌─────────────────────────────────┐ │ │ L2: body │ │ │ ... use index_target ... │ │ └──────────────┬──────────────────┘ │ ▼ │ ┌─────────────────────────────────┐ │ │ L3: step │ │ │ index_reg = index_reg + step │ │ │ index_target = index_reg ◄── BUG │ │ goto L1 │ │ └─────────────────────────────────┘ │ ▼ ┌───────────────┐ │ L4: exit │ │ index_target │ │ = overshot! │ └───────────────┘ ``` <br /> - After: ``` ┌─────────────────────────────────────┐ │ L1: condition │ │ if index_reg < end → L2 else L4 │ └──────────────┬──────────────────┬───┘ │ true │ false ▼ │ ┌─────────────────────────────────┐ │ │ L2: body │ │ │ index_target = index_reg ◄── FIX │ │ ... use index_target ... │ │ └──────────────┬──────────────────┘ │ ▼ │ ┌─────────────────────────────────┐ │ │ L3: step │ │ │ index_reg = index_reg + step │ │ │ goto L1 │ │ └─────────────────────────────────┘ │ ▼ ┌───────────────┐ │ L4: exit │ │ index_target │ │ = correct ✓ │ └───────────────┘ ```
1 parent 67ada30 commit 978b711

File tree

9 files changed

+54
-27
lines changed

9 files changed

+54
-27
lines changed

‎mypyc/irbuild/for_helpers.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1141,6 +1141,13 @@ def gen_condition(self) -> None:
11411141
)
11421142
builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
11431143

1144+
def begin_body(self) -> None:
1145+
# Update the user-visible loop variable at the start of the body,
1146+
# after the condition check passes. This ensures the variable isn't
1147+
# "overshot" when the loop exits (matching CPython semantics).
1148+
builder = self.builder
1149+
builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line)
1150+
11441151
def gen_step(self) -> None:
11451152
builder = self.builder
11461153
line = self.line
@@ -1163,7 +1170,6 @@ def gen_step(self) -> None:
11631170
builder.read(self.index_reg, line), Integer(self.step), "+", line
11641171
)
11651172
builder.assign(self.index_reg, new_val, line)
1166-
builder.assign(self.index_target, new_val, line)
11671173

11681174

11691175
class ForInfiniteCounter(ForGenerator):

‎mypyc/test-data/irbuild-basic.test‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3458,12 +3458,12 @@ L1:
34583458
r1 = int_lt r0, 24
34593459
if r1 goto L2 else goto L4 :: bool
34603460
L2:
3461+
i = r0
34613462
r2 = CPyTagged_Add(sum, i)
34623463
sum = r2
34633464
L3:
34643465
r3 = r0 + 4
34653466
r0 = r3
3466-
i = r3
34673467
goto L1
34683468
L4:
34693469
return 1

‎mypyc/test-data/irbuild-i64.test‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -547,11 +547,11 @@ L1:
547547
r1 = r0 < x :: signed
548548
if r1 goto L2 else goto L4 :: bool
549549
L2:
550+
n = r0
550551
r2 = g(n)
551552
L3:
552553
r3 = r0 + 1
553554
r0 = r3
554-
n = r3
555555
goto L1
556556
L4:
557557
return 1
@@ -1613,11 +1613,11 @@ L1:
16131613
r1 = r0 < 4 :: signed
16141614
if r1 goto L2 else goto L4 :: bool
16151615
L2:
1616+
x = r0
16161617
y = x
16171618
L3:
16181619
r2 = r0 + 1
16191620
r0 = r2
1620-
x = r2
16211621
goto L1
16221622
L4:
16231623
return 1
@@ -1640,11 +1640,11 @@ L1:
16401640
r1 = r0 < 4 :: signed
16411641
if r1 goto L2 else goto L4 :: bool
16421642
L2:
1643+
x = r0
16431644
y = x
16441645
L3:
16451646
r2 = r0 + 1
16461647
r0 = r2
1647-
x = r2
16481648
goto L1
16491649
L4:
16501650
return 1

‎mypyc/test-data/irbuild-lists.test‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -286,14 +286,14 @@ L1:
286286
r3 = int_lt r2, r1
287287
if r3 goto L2 else goto L4 :: bool
288288
L2:
289+
i = r2
289290
r4 = CPyList_GetItem(l, i)
290291
r5 = object 1
291292
r6 = PyNumber_InPlaceAdd(r4, r5)
292293
r7 = CPyList_SetItem(l, i, r6)
293294
L3:
294295
r8 = r2 + 2
295296
r2 = r8
296-
i = r8
297297
goto L1
298298
L4:
299299
return l

‎mypyc/test-data/irbuild-set.test‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -229,14 +229,14 @@ L1:
229229
r2 = int_lt r1, 12
230230
if r2 goto L2 else goto L4 :: bool
231231
L2:
232+
x = r1
232233
r3 = f(x)
233234
r4 = box(int, r3)
234235
r5 = PySet_Add(r0, r4)
235236
r6 = r5 >= 0 :: signed
236237
L3:
237238
r7 = r1 + 4
238239
r1 = r7
239-
x = r7
240240
goto L1
241241
L4:
242242
d = r0
@@ -260,14 +260,14 @@ L1:
260260
r2 = int_lt r1, 12
261261
if r2 goto L2 else goto L4 :: bool
262262
L2:
263+
x = r1
263264
r3 = f(x)
264265
r4 = box(int, r3)
265266
r5 = PySet_Add(r0, r4)
266267
r6 = r5 >= 0 :: signed
267268
L3:
268269
r7 = r1 + 4
269270
r1 = r7
270-
x = r7
271271
goto L1
272272
L4:
273273
e = r0

‎mypyc/test-data/irbuild-statements.test‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,12 @@ L1:
1919
r1 = int_lt r0, 10
2020
if r1 goto L2 else goto L4 :: bool
2121
L2:
22+
i = r0
2223
r2 = CPyTagged_Add(x, i)
2324
x = r2
2425
L3:
2526
r3 = r0 + 2
2627
r0 = r3
27-
i = r3
2828
goto L1
2929
L4:
3030
return 1
@@ -45,10 +45,10 @@ L1:
4545
r1 = int_lt r0, a
4646
if r1 goto L2 else goto L4 :: bool
4747
L2:
48+
i = r0
4849
L3:
4950
r2 = CPyTagged_Add(r0, 2)
5051
r0 = r2
51-
i = r2
5252
goto L1
5353
L4:
5454
return 1
@@ -70,10 +70,10 @@ L1:
7070
r1 = int_gt r0, 0
7171
if r1 goto L2 else goto L4 :: bool
7272
L2:
73+
i = r0
7374
L3:
7475
r2 = r0 + -2
7576
r0 = r2
76-
i = r2
7777
goto L1
7878
L4:
7979
return 1
@@ -113,11 +113,11 @@ L1:
113113
r1 = int_lt r0, 10
114114
if r1 goto L2 else goto L4 :: bool
115115
L2:
116+
n = r0
116117
goto L4
117118
L3:
118119
r2 = r0 + 2
119120
r0 = r2
120-
n = r2
121121
goto L1
122122
L4:
123123
return 1
@@ -183,10 +183,10 @@ L1:
183183
r1 = int_lt r0, 10
184184
if r1 goto L2 else goto L4 :: bool
185185
L2:
186+
n = r0
186187
L3:
187188
r2 = r0 + 2
188189
r0 = r2
189-
n = r2
190190
goto L1
191191
L4:
192192
return 1
@@ -1016,13 +1016,13 @@ L4:
10161016
r8 = list_get_item_unsafe b, r1
10171017
r9 = unbox(int, r8)
10181018
y = r9
1019+
z = r2
10191020
x = 0
10201021
L5:
10211022
r10 = r1 + 1
10221023
r1 = r10
10231024
r11 = r2 + 2
10241025
r2 = r11
1025-
z = r11
10261026
goto L1
10271027
L6:
10281028
r12 = CPy_NoErrOccurred()

‎mypyc/test-data/irbuild-vec-i64.test‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -337,13 +337,13 @@ L1:
337337
r3 = r2 < 5 :: signed
338338
if r3 goto L2 else goto L4 :: bool
339339
L2:
340+
x = r2
340341
r4 = x + 1
341342
r5 = VecI64Api.append(r1, r4)
342343
r1 = r5
343344
L3:
344345
r6 = r2 + 1
345346
r2 = r6
346-
x = r6
347347
goto L1
348348
L4:
349349
return r1
@@ -464,9 +464,9 @@ def f():
464464
r2 :: short_int
465465
r3, ___tmp_6 :: i64
466466
r4 :: bit
467-
r5 :: vec[i64]
468-
r6 :: short_int
469-
r7 :: i64
467+
r5 :: i64
468+
r6 :: vec[i64]
469+
r7 :: short_int
470470
L0:
471471
r0 = VecI64Api.alloc(0, 0)
472472
r1 = r0
@@ -477,13 +477,13 @@ L1:
477477
r4 = int_lt r2, 14
478478
if r4 goto L2 else goto L4 :: bool
479479
L2:
480-
r5 = VecI64Api.append(r1, ___tmp_6)
481-
r1 = r5
480+
r5 = r2 >> 1
481+
___tmp_6 = r5
482+
r6 = VecI64Api.append(r1, ___tmp_6)
483+
r1 = r6
482484
L3:
483-
r6 = r2 + 2
484-
r2 = r6
485-
r7 = r6 >> 1
486-
___tmp_6 = r7
485+
r7 = r2 + 2
486+
r2 = r7
487487
goto L1
488488
L4:
489489
return r1

‎mypyc/test-data/irbuild-vec-t.test‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,7 @@ L1:
275275
r5 = r4 < 5 :: signed
276276
if r5 goto L2 else goto L4 :: bool
277277
L2:
278+
x = r4
278279
r6 = 'x'
279280
r7 = load_address PyUnicode_Type
280281
r8 = r7
@@ -283,7 +284,6 @@ L2:
283284
L3:
284285
r10 = r4 + 1
285286
r4 = r10
286-
x = r10
287287
goto L1
288288
L4:
289289
return r3

‎mypyc/test-data/run-loops.test‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -277,8 +277,29 @@ nested_enumerate()
277277
nested_range()
278278
nested_list()
279279
gen = nested_yield()
280-
for k in range(12):
281-
assert next(gen) == k % 4
280+
expected = [0, 1, 2, 2, 0, 1, 2, 2, 0, 1, 2, 2]
281+
assert list(gen) == expected
282+
[out]
283+
284+
[case testRangeLoopVariableAfterLoop]
285+
for i in range(3):
286+
pass
287+
assert i == 2, f"expected 2, got {i}"
288+
289+
for j in range(0, 10, 3):
290+
pass
291+
assert j == 9, f"expected 9, got {j}"
292+
293+
for k in range(5, 0, -1):
294+
pass
295+
assert k == 1, f"expected 1, got {k}"
296+
297+
for m in range(0, 10, 2):
298+
pass
299+
assert m == 8, f"expected 8, got {m}"
300+
301+
[file driver.py]
302+
from native import *
282303
[out]
283304

284305
[case testForIterable]

0 commit comments

Comments
 (0)