[mypyc] Fix range loop variable off-by-one after loop exit

mypyc/irbuild/for_helpers.py

@@ -1141,6 +1141,13 @@ def gen_condition(self) -> None:
         )
         builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
 
+    def begin_body(self) -> None:
+        # Update the user-visible loop variable at the start of the body,
+        # after the condition check passes. This ensures the variable isn't
+        # "overshot" when the loop exits (matching CPython semantics).
+        builder = self.builder
+        builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line)
+
     def gen_step(self) -> None:
         builder = self.builder
         line = self.line
@@ -1163,7 +1170,6 @@ def gen_step(self) -> None:
                 builder.read(self.index_reg, line), Integer(self.step), "+", line
             )
         builder.assign(self.index_reg, new_val, line)
-        builder.assign(self.index_target, new_val, line)
 
 
 class ForInfiniteCounter(ForGenerator):

mypyc/test-data/irbuild-basic.test

@@ -3458,12 +3458,12 @@ L1:
     r1 = int_lt r0, 24
     if r1 goto L2 else goto L4 :: bool
 L2:
+    i = r0
     r2 = CPyTagged_Add(sum, i)
     sum = r2
 L3:
     r3 = r0 + 4
     r0 = r3
-    i = r3
     goto L1
 L4:
     return 1

mypyc/test-data/irbuild-i64.test

@@ -547,11 +547,11 @@ L1:
     r1 = r0 < x :: signed
     if r1 goto L2 else goto L4 :: bool
 L2:
+    n = r0
     r2 = g(n)
 L3:
     r3 = r0 + 1
     r0 = r3
-    n = r3
     goto L1
 L4:
     return 1
@@ -1613,11 +1613,11 @@ L1:
     r1 = r0 < 4 :: signed
     if r1 goto L2 else goto L4 :: bool
 L2:
+    x = r0
     y = x
 L3:
     r2 = r0 + 1
     r0 = r2
-    x = r2
     goto L1
 L4:
     return 1
@@ -1640,11 +1640,11 @@ L1:
     r1 = r0 < 4 :: signed
     if r1 goto L2 else goto L4 :: bool
 L2:
+    x = r0
     y = x
 L3:
     r2 = r0 + 1
     r0 = r2
-    x = r2
     goto L1
 L4:
     return 1

mypyc/test-data/irbuild-lists.test

@@ -286,14 +286,14 @@ L1:
     r3 = int_lt r2, r1
     if r3 goto L2 else goto L4 :: bool
 L2:
+    i = r2
     r4 = CPyList_GetItem(l, i)
     r5 = object 1
     r6 = PyNumber_InPlaceAdd(r4, r5)
     r7 = CPyList_SetItem(l, i, r6)
 L3:
     r8 = r2 + 2
     r2 = r8
-    i = r8
     goto L1
 L4:
     return l

mypyc/test-data/irbuild-set.test

@@ -229,14 +229,14 @@ L1:
     r2 = int_lt r1, 12
     if r2 goto L2 else goto L4 :: bool
 L2:
+    x = r1
     r3 = f(x)
     r4 = box(int, r3)
     r5 = PySet_Add(r0, r4)
     r6 = r5 >= 0 :: signed
 L3:
     r7 = r1 + 4
     r1 = r7
-    x = r7
     goto L1
 L4:
     d = r0
@@ -260,14 +260,14 @@ L1:
     r2 = int_lt r1, 12
     if r2 goto L2 else goto L4 :: bool
 L2:
+    x = r1
     r3 = f(x)
     r4 = box(int, r3)
     r5 = PySet_Add(r0, r4)
     r6 = r5 >= 0 :: signed
 L3:
     r7 = r1 + 4
     r1 = r7
-    x = r7
     goto L1
 L4:
     e = r0

mypyc/test-data/irbuild-statements.test

@@ -19,12 +19,12 @@ L1:
     r1 = int_lt r0, 10
     if r1 goto L2 else goto L4 :: bool
 L2:
+    i = r0
     r2 = CPyTagged_Add(x, i)
     x = r2
 L3:
     r3 = r0 + 2
     r0 = r3
-    i = r3
     goto L1
 L4:
     return 1
@@ -45,10 +45,10 @@ L1:
     r1 = int_lt r0, a
     if r1 goto L2 else goto L4 :: bool
 L2:
+    i = r0
 L3:
     r2 = CPyTagged_Add(r0, 2)
     r0 = r2
-    i = r2
     goto L1
 L4:
     return 1
@@ -70,10 +70,10 @@ L1:
     r1 = int_gt r0, 0
     if r1 goto L2 else goto L4 :: bool
 L2:
+    i = r0
 L3:
     r2 = r0 + -2
     r0 = r2
-    i = r2
     goto L1
 L4:
     return 1
@@ -113,11 +113,11 @@ L1:
     r1 = int_lt r0, 10
     if r1 goto L2 else goto L4 :: bool
 L2:
+    n = r0
     goto L4
 L3:
     r2 = r0 + 2
     r0 = r2
-    n = r2
     goto L1
 L4:
     return 1
@@ -183,10 +183,10 @@ L1:
     r1 = int_lt r0, 10
     if r1 goto L2 else goto L4 :: bool
 L2:
+    n = r0
 L3:
     r2 = r0 + 2
     r0 = r2
-    n = r2
     goto L1
 L4:
     return 1
@@ -1016,13 +1016,13 @@ L4:
     r8 = list_get_item_unsafe b, r1
     r9 = unbox(int, r8)
     y = r9
+    z = r2
     x = 0
 L5:
     r10 = r1 + 1
     r1 = r10
     r11 = r2 + 2
     r2 = r11
-    z = r11
     goto L1
 L6:
     r12 = CPy_NoErrOccurred()

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

@@ -337,13 +337,13 @@ L1:
     r3 = r2 < 5 :: signed
     if r3 goto L2 else goto L4 :: bool
 L2:
+    x = r2
     r4 = x + 1
     r5 = VecI64Api.append(r1, r4)
     r1 = r5
 L3:
     r6 = r2 + 1
     r2 = r6
-    x = r6
     goto L1
 L4:
     return r1
@@ -464,9 +464,9 @@ def f():
     r2 :: short_int
     r3, ___tmp_6 :: i64
     r4 :: bit
-    r5 :: vec[i64]
-    r6 :: short_int
-    r7 :: i64
+    r5 :: i64
+    r6 :: vec[i64]
+    r7 :: short_int
 L0:
     r0 = VecI64Api.alloc(0, 0)
     r1 = r0
@@ -477,13 +477,13 @@ L1:
     r4 = int_lt r2, 14
     if r4 goto L2 else goto L4 :: bool
 L2:
-    r5 = VecI64Api.append(r1, ___tmp_6)
-    r1 = r5
+    r5 = r2 >> 1
+    ___tmp_6 = r5
+    r6 = VecI64Api.append(r1, ___tmp_6)
+    r1 = r6
 L3:
-    r6 = r2 + 2
-    r2 = r6
-    r7 = r6 >> 1
-    ___tmp_6 = r7
+    r7 = r2 + 2
+    r2 = r7
     goto L1
 L4:
     return r1

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

@@ -275,6 +275,7 @@ L1:
     r5 = r4 < 5 :: signed
     if r5 goto L2 else goto L4 :: bool
 L2:
+    x = r4
     r6 = 'x'
     r7 = load_address PyUnicode_Type
     r8 = r7
@@ -283,7 +284,6 @@ L2:
 L3:
     r10 = r4 + 1
     r4 = r10
-    x = r10
     goto L1
 L4:
     return r3

mypyc/test-data/run-loops.test

@@ -277,8 +277,29 @@ nested_enumerate()
 nested_range()
 nested_list()
 gen = nested_yield()
-for k in range(12):
-    assert next(gen) == k % 4
+expected = [0, 1, 2, 2, 0, 1, 2, 2, 0, 1, 2, 2]
+assert list(gen) == expected
+[out]
+
+[case testRangeLoopVariableAfterLoop]
+for i in range(3):
+    pass
+assert i == 2, f"expected 2, got {i}"
+
+for j in range(0, 10, 3):
+    pass
+assert j == 9, f"expected 9, got {j}"
+
+for k in range(5, 0, -1):
+    pass
+assert k == 1, f"expected 1, got {k}"
+
+for m in range(0, 10, 2):
+    pass
+assert m == 8, f"expected 8, got {m}"
+
+[file driver.py]
+from native import *
 [out]
 
 [case testForIterable]