[mypyc] Don't assign the range loop variable if the range is empty

mypyc/irbuild/for_helpers.py

@@ -70,7 +70,7 @@
 )
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.constant_fold import constant_fold_expr
-from mypyc.irbuild.targets import AssignmentTarget, AssignmentTargetTuple
+from mypyc.irbuild.targets import AssignmentTargetTuple
 from mypyc.irbuild.vec import vec_append, vec_create, vec_get_item_unsafe, vec_init_item_unsafe
 from mypyc.primitives.dict_ops import (
     dict_check_size_op,
@@ -654,6 +654,7 @@ def __init__(
         self.index = index
         self.body_block = body_block
         self.line = line
+        self.nested = nested
         # Some for loops need a cleanup block that we execute at exit. We
         # create a cleanup block if needed. However, if we are generating a for
         # loop for a nested iterator, such as "e" in "enumerate(e)", the
@@ -946,19 +947,34 @@ def gen_condition(self) -> None:
             # (unless input is immutable type).
             len_reg = builder.read(self.length_reg, line)
         comparison = builder.binary_op(builder.read(self.index_target, line), len_reg, "<", line)
-        builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
+        if self.nested:
+            # In zip() or enumerate(), read the item right after the length check, as
+            # the sequence's iterator would. Other iterators may be advanced and other
+            # loop targets assigned before our begin_body(), and either could shrink
+            # the sequence.
+            read_block = BasicBlock()
+            builder.add_bool_branch(comparison, read_block, self.loop_exit)
+            builder.activate_block(read_block)
+            self.next_reg = self.read_item()
+            builder.goto(self.body_block)
+        else:
+            builder.add_bool_branch(comparison, self.body_block, self.loop_exit)
 
-    def begin_body(self) -> None:
+    def read_item(self) -> Value:
         builder = self.builder
         line = self.line
-        # Read the next list item.
-        value_box = unsafe_index(
+        return unsafe_index(
             builder,
             builder.read(self.expr_target, line),
             builder.read(self.index_target, line),
             line,
         )
-        assert value_box
+
+    def begin_body(self) -> None:
+        builder = self.builder
+        line = self.line
+        # Read the next list item, unless gen_condition() already did.
+        value_box = self.next_reg if self.nested else self.read_item()
         # We coerce to the type of list elements here so that
         # iterating with tuple unpacking generates a tuple based
         # unpack instead of an iterator based one.
@@ -1139,9 +1155,6 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
         index_reg = Register(index_type, line=self.line)
         builder.assign(index_reg, start_reg, self.line)
         self.index_reg = index_reg
-        # Initialize loop index to 0. Assert that the index target is assignable.
-        self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
-        builder.assign(self.index_target, builder.read(self.index_reg, self.line), self.line)
 
     def convert_arg(self, value: Value) -> Value:
         """Convert a range() argument to an int using __index__, like range() does."""
@@ -1175,9 +1188,13 @@ def gen_condition(self) -> None:
     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).
+        # assigned if the range is empty, and 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)
+        line = self.line
+        builder.assign(
+            builder.get_assignment_target(self.index), builder.read(self.index_reg, line), line
+        )
 
     def gen_step(self) -> None:
         builder = self.builder
@@ -1209,10 +1226,9 @@ class ForInfiniteCounter(ForGenerator):
     def init(self) -> None:
         builder = self.builder
         # Create a register to store the state of the loop index and
-        # initialize this register along with the loop index to 0.
+        # initialize this register to 0.
         zero = Integer(0)
         self.index_reg = builder.ensure_register(zero)
-        self.index_target: Register | AssignmentTarget = builder.get_assignment_target(self.index)
 
     def gen_step(self) -> None:
         builder = self.builder
@@ -1226,8 +1242,10 @@ def gen_step(self) -> None:
         builder.assign(self.index_reg, new_val, line)
 
     def begin_body(self) -> None:
-        self.builder.assign(
-            self.index_target, self.builder.read(self.index_reg, self.line), self.line
+        builder = self.builder
+        line = self.line
+        builder.assign(
+            builder.get_assignment_target(self.index), builder.read(self.index_reg, line), line
         )
 
 

mypyc/test-data/irbuild-basic.test

@@ -3512,14 +3512,12 @@ L5:
 def range_in_loop():
     sum :: int
     r0 :: short_int
-    i :: int
     r1 :: bit
-    r2 :: int
+    i, r2 :: int
     r3 :: short_int
 L0:
     sum = 0
     r0 = 8
-    i = r0
 L1:
     r1 = int_lt r0, 24
     if r1 goto L2 else goto L4 :: bool

mypyc/test-data/irbuild-i64.test

@@ -536,13 +536,13 @@ def g(a):
 L0:
     return 1
 def f(x):
-    x, r0, n :: i64
+    x, r0 :: i64
     r1 :: bit
+    n :: i64
     r2 :: None
     r3 :: i64
 L0:
     r0 = 0
-    n = r0
 L1:
     r1 = r0 < x :: signed
     if r1 goto L2 else goto L4 :: bool
@@ -1661,12 +1661,11 @@ def f() -> None:
         y = x
 [out]
 def f():
-    r0, x :: i64
+    r0 :: i64
     r1 :: bit
-    y, r2 :: i64
+    x, y, r2 :: i64
 L0:
     r0 = 0
-    x = r0
 L1:
     r1 = r0 < 4 :: signed
     if r1 goto L2 else goto L4 :: bool
@@ -1688,12 +1687,11 @@ def f() -> None:
         y = x
 [out]
 def f():
-    r0, x :: i64
+    r0 :: i64
     r1 :: bit
-    y, r2 :: i64
+    x, y, r2 :: i64
 L0:
     r0 = 0
-    x = r0
 L1:
     r1 = r0 < 4 :: signed
     if r1 goto L2 else goto L4 :: bool

mypyc/test-data/irbuild-lists.test

@@ -290,16 +290,15 @@ def increment(l):
     l :: list
     r0 :: native_int
     r1, r2 :: short_int
-    i :: int
     r3 :: bit
+    i :: int
     r4, r5, r6 :: object
     r7 :: bit
     r8 :: short_int
 L0:
     r0 = var_object_size l
     r1 = r0 << 1
     r2 = 0
-    i = r2
 L1:
     r3 = int_lt r2, r1
     if r3 goto L2 else goto L4 :: bool

mypyc/test-data/irbuild-set.test

@@ -213,9 +213,8 @@ L5:
 def test4():
     r0 :: set
     r1 :: short_int
-    x :: int
     r2 :: bit
-    r3 :: int
+    x, r3 :: int
     r4 :: object
     r5 :: i32
     r6 :: bit
@@ -224,7 +223,6 @@ def test4():
 L0:
     r0 = PySet_New(0)
     r1 = 2
-    x = r1
 L1:
     r2 = int_lt r1, 12
     if r2 goto L2 else goto L4 :: bool
@@ -244,9 +242,8 @@ L4:
 def test5():
     r0 :: set
     r1 :: short_int
-    x :: int
     r2 :: bit
-    r3 :: int
+    x, r3 :: int
     r4 :: object
     r5 :: i32
     r6 :: bit
@@ -255,7 +252,6 @@ def test5():
 L0:
     r0 = PySet_New(0)
     r1 = 2
-    x = r1
 L1:
     r2 = int_lt r1, 12
     if r2 goto L2 else goto L4 :: bool