[mypyc] Fix for loop iteration when the loop body reassigns the iterable

Author: rheard

mypyc/irbuild/for_helpers.py

@@ -7,26 +7,36 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Set as AbstractSet
 from typing import ClassVar, cast
 
 from mypy.nodes import (
     ARG_POS,
+    AssignmentExpr,
+    AssignmentStmt,
     CallExpr,
+    DelStmt,
     DictionaryComprehension,
     Expression,
+    ForStmt,
     GeneratorExpr,
     ListExpr,
     Lvalue,
     MemberExpr,
     NameExpr,
+    OperatorAssignmentStmt,
     RefExpr,
     SetExpr,
     StarExpr,
+    SymbolNode,
+    TryStmt,
     TupleExpr,
     TypeAlias,
     Var,
+    WithStmt,
 )
+from mypy.patterns import AsPattern, MappingPattern, StarredPattern
+from mypy.traverser import TraverserVisitor
 from mypy.types import LiteralType, TupleType, get_proper_type, get_proper_types
 from mypyc.common import GENERATOR_HELPER_NAME, IS_FREE_THREADED
 from mypyc.ir.ops import (
@@ -70,7 +80,7 @@
 )
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.constant_fold import constant_fold_expr
-from mypyc.irbuild.targets import AssignmentTargetTuple
+from mypyc.irbuild.targets import AssignmentTargetRegister, 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,
@@ -106,6 +116,7 @@ def for_loop_helper(
     else_insts: GenFunc | None,
     is_async: bool,
     line: int,
+    reassigned: AbstractSet[SymbolNode] = frozenset(),
 ) -> None:
     """Generate IR for a loop.
 
@@ -114,6 +125,8 @@ def for_loop_helper(
         expr: the expression to iterate over
         body_insts: a function that generates the body of the loop
         else_insts: a function that generates the else block instructions
+        reassigned: local variables that the loop index or body may assign to
+            (see find_loop_assignments)
     """
     # Body of the loop
     body_block = BasicBlock()
@@ -128,7 +141,14 @@ def for_loop_helper(
     normal_loop_exit = else_block if else_insts is not None else exit_block
 
     for_gen = make_for_loop_generator(
-        builder, index, expr, body_block, normal_loop_exit, line, is_async=is_async
+        builder,
+        index,
+        expr,
+        body_block,
+        normal_loop_exit,
+        line,
+        is_async=is_async,
+        reassigned=reassigned,
     )
 
     builder.push_loop_stack(step_block, exit_block)
@@ -479,6 +499,95 @@ def is_range_ref(expr: RefExpr) -> bool:
     )
 
 
+def find_loop_assignments(s: ForStmt) -> set[SymbolNode]:
+    """Return the variables that the index or the body of a for loop may assign to."""
+    collector = AssignmentCollector()
+    collector.add_lvalue(s.index)
+    s.body.accept(collector)
+    return collector.symbols
+
+
+class AssignmentCollector(TraverserVisitor):
+    """Collect the symbols of the variables assigned to (or deleted) in a subtree."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.symbols: set[SymbolNode] = set()
+
+    def add_lvalue(self, lvalue: Expression | None) -> None:
+        if isinstance(lvalue, NameExpr):
+            if lvalue.node is not None:
+                self.symbols.add(lvalue.node)
+        elif isinstance(lvalue, (TupleExpr, ListExpr)):
+            for item in lvalue.items:
+                self.add_lvalue(item)
+        elif isinstance(lvalue, StarExpr):
+            self.add_lvalue(lvalue.expr)
+
+    def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
+        for lvalue in o.lvalues:
+            self.add_lvalue(lvalue)
+        super().visit_assignment_stmt(o)
+
+    def visit_operator_assignment_stmt(self, o: OperatorAssignmentStmt) -> None:
+        self.add_lvalue(o.lvalue)
+        super().visit_operator_assignment_stmt(o)
+
+    def visit_assignment_expr(self, o: AssignmentExpr) -> None:
+        self.add_lvalue(o.target)
+        super().visit_assignment_expr(o)
+
+    def visit_for_stmt(self, o: ForStmt) -> None:
+        self.add_lvalue(o.index)
+        super().visit_for_stmt(o)
+
+    def visit_with_stmt(self, o: WithStmt) -> None:
+        for target in o.target:
+            self.add_lvalue(target)
+        super().visit_with_stmt(o)
+
+    def visit_try_stmt(self, o: TryStmt) -> None:
+        for var in o.vars:
+            self.add_lvalue(var)
+        super().visit_try_stmt(o)
+
+    def visit_del_stmt(self, o: DelStmt) -> None:
+        self.add_lvalue(o.expr)
+        super().visit_del_stmt(o)
+
+    def visit_as_pattern(self, o: AsPattern) -> None:
+        self.add_lvalue(o.name)
+        super().visit_as_pattern(o)
+
+    def visit_starred_pattern(self, o: StarredPattern) -> None:
+        self.add_lvalue(o.capture)
+        super().visit_starred_pattern(o)
+
+    def visit_mapping_pattern(self, o: MappingPattern) -> None:
+        self.add_lvalue(o.rest)
+        super().visit_mapping_pattern(o)
+
+
+def accept_loop_state(
+    builder: IRBuilder, expr: Expression, reassigned: AbstractSet[SymbolNode]
+) -> Value:
+    """Evaluate the object to iterate over, or the end of a range.
+
+    Specialized loops read this value again on every iteration. If it's the register
+    of a local variable that the loop may assign to, copy it to a new register, since
+    the assignment must not affect the iteration (as in "for x in a: a = []").
+    """
+    value = builder.accept(expr)
+    if isinstance(value, Register):
+        for symbol in reassigned:
+            target = builder.symtables[-1].get(symbol)
+            if isinstance(target, AssignmentTargetRegister) and target.register is value:
+                copy = Register(value.type)
+                builder.assign(copy, value, expr.line)
+                return copy
+    return value
+
+
 def make_for_loop_generator(
     builder: IRBuilder,
     index: Lvalue,
@@ -488,10 +597,14 @@ def make_for_loop_generator(
     line: int,
     is_async: bool = False,
     nested: bool = False,
+    reassigned: AbstractSet[SymbolNode] = frozenset(),
 ) -> ForGenerator:
     """Return helper object for generating a for loop over an iterable.
 
     If "nested" is True, this is a nested iterator such as "e" in "enumerate(e)".
+
+    "reassigned" contains the local variables that the loop may assign to
+    (see accept_loop_state).
     """
 
     # Do an async loop if needed. async is always generic
@@ -506,7 +619,7 @@ def make_for_loop_generator(
     rtyp = builder.node_type(expr)
     if is_sequence_rprimitive(rtyp) or isinstance(rtyp, RVec):
         # Special case "for x in <seq>" for concrete sequence types.
-        expr_reg = builder.accept(expr)
+        expr_reg = accept_loop_state(builder, expr, reassigned)
         target_type = builder.get_sequence_type(expr)
 
         for_list = ForSequence(builder, index, body_block, loop_exit, line, nested)
@@ -515,7 +628,7 @@ def make_for_loop_generator(
 
     if is_dict_rprimitive(rtyp):
         # Special case "for k in <dict>".
-        expr_reg = builder.accept(expr)
+        expr_reg = accept_loop_state(builder, expr, reassigned)
         target_type = builder.get_dict_key_type(expr)
 
         for_dict = ForDictionaryKeys(builder, index, body_block, loop_exit, line, nested)
@@ -537,10 +650,10 @@ def make_for_loop_generator(
             # direction of comparison to do.
             if len(expr.args) == 1:
                 start_reg: Value = Integer(0)
-                end_reg = builder.accept(expr.args[0])
+                end_reg = accept_loop_state(builder, expr.args[0], reassigned)
             else:
                 start_reg = builder.accept(expr.args[0])
-                end_reg = builder.accept(expr.args[1])
+                end_reg = accept_loop_state(builder, expr.args[1], reassigned)
             if len(expr.args) == 3:
                 step = builder.extract_int(expr.args[2])
                 assert step is not None
@@ -564,7 +677,7 @@ def make_for_loop_generator(
             lvalue1 = index.items[0]
             lvalue2 = index.items[1]
             for_enumerate = ForEnumerate(builder, index, body_block, loop_exit, line, nested)
-            for_enumerate.init(lvalue1, lvalue2, expr.args[0])
+            for_enumerate.init(lvalue1, lvalue2, expr.args[0], reassigned)
             return for_enumerate
 
         elif (
@@ -576,7 +689,7 @@ def make_for_loop_generator(
         ):
             # Special case "for x, y in zip(a, b)".
             for_zip = ForZip(builder, index, body_block, loop_exit, line, nested)
-            for_zip.init(index.items, expr.args)
+            for_zip.init(index.items, expr.args, reassigned)
             return for_zip
 
         if (
@@ -586,7 +699,7 @@ def make_for_loop_generator(
             and is_sequence_rprimitive(builder.node_type(expr.args[0]))
         ):
             # Special case "for x in reversed(<list>)".
-            expr_reg = builder.accept(expr.args[0])
+            expr_reg = accept_loop_state(builder, expr.args[0], reassigned)
             target_type = builder.get_sequence_type(expr)
 
             for_list = ForSequence(builder, index, body_block, loop_exit, line, nested)
@@ -596,7 +709,7 @@ def make_for_loop_generator(
         # Special cases for dictionary iterator methods, like dict.items().
         rtype = builder.node_type(expr.callee.expr)
         if is_dict_rprimitive(rtype) and expr.callee.name in ("keys", "values", "items"):
-            expr_reg = builder.accept(expr.callee.expr)
+            expr_reg = accept_loop_state(builder, expr.callee.expr, reassigned)
             for_dict_type: type[ForGenerator] | None = None
             if expr.callee.name == "keys":
                 target_type = builder.get_dict_key_type(expr.callee.expr)
@@ -1257,15 +1370,28 @@ def need_cleanup(self) -> bool:
         # redundant cleanup block, but that's okay.
         return True
 
-    def init(self, index1: Lvalue, index2: Lvalue, expr: Expression) -> None:
+    def init(
+        self,
+        index1: Lvalue,
+        index2: Lvalue,
+        expr: Expression,
+        reassigned: AbstractSet[SymbolNode] = frozenset(),
+    ) -> None:
         # Count from 0 to infinity (for the index lvalue).
         self.index_gen = ForInfiniteCounter(
             self.builder, index1, self.body_block, self.loop_exit, self.line, nested=True
         )
         self.index_gen.init()
         # Iterate over the actual iterable.
         self.main_gen = make_for_loop_generator(
-            self.builder, index2, expr, self.body_block, self.loop_exit, self.line, nested=True
+            self.builder,
+            index2,
+            expr,
+            self.body_block,
+            self.loop_exit,
+            self.line,
+            nested=True,
+            reassigned=reassigned,
         )
 
     def gen_condition(self) -> None:
@@ -1293,15 +1419,27 @@ def need_cleanup(self) -> bool:
         # redundant cleanup block, but that's okay.
         return True
 
-    def init(self, indexes: list[Lvalue], exprs: list[Expression]) -> None:
+    def init(
+        self,
+        indexes: list[Lvalue],
+        exprs: list[Expression],
+        reassigned: AbstractSet[SymbolNode] = frozenset(),
+    ) -> None:
         assert len(indexes) == len(exprs)
         # Condition check will require multiple basic blocks, since there will be
         # multiple conditions to check.
         self.cond_blocks = [BasicBlock() for _ in range(len(indexes) - 1)] + [self.body_block]
         self.gens: list[ForGenerator] = []
         for index, expr, next_block in zip(indexes, exprs, self.cond_blocks):
             gen = make_for_loop_generator(
-                self.builder, index, expr, next_block, self.loop_exit, self.line, nested=True
+                self.builder,
+                index,
+                expr,
+                next_block,
+                self.loop_exit,
+                self.line,
+                nested=True,
+                reassigned=reassigned,
             )
             self.gens.append(gen)
 

mypyc/irbuild/statement.py

@@ -98,7 +98,7 @@
     find_walrus_targets,
     int_borrow_friendly_op,
 )
-from mypyc.irbuild.for_helpers import for_loop_helper
+from mypyc.irbuild.for_helpers import find_loop_assignments, for_loop_helper
 from mypyc.irbuild.generator import add_raise_exception_blocks_to_generator_class
 from mypyc.irbuild.nonlocalcontrol import (
     ExceptNonlocalControl,
@@ -661,7 +661,14 @@ def else_block() -> None:
         builder.accept(s.else_body)
 
     for_loop_helper(
-        builder, s.index, s.expr, body, else_block if s.else_body else None, s.is_async, s.line
+        builder,
+        s.index,
+        s.expr,
+        body,
+        else_block if s.else_body else None,
+        s.is_async,
+        s.line,
+        reassigned=find_loop_assignments(s),
     )
 
 

mypyc/test-data/irbuild-statements.test

@@ -1172,6 +1172,37 @@ L7:
 L8:
     return 1
 
+[case testForLoopReassignIterable]
+def f(a: tuple[int, ...], b: tuple[int, ...]) -> None:
+    for x in a:
+        a = b
+[out]
+def f(a, b):
+    a, b, r0 :: tuple
+    r1, r2 :: native_int
+    r3 :: bit
+    r4 :: object
+    r5, x :: int
+    r6 :: native_int
+L0:
+    r0 = a
+    r1 = var_object_size r0
+    r2 = 0
+L1:
+    r3 = r2 < r1 :: signed
+    if r3 goto L2 else goto L4 :: bool
+L2:
+    r4 = CPySequenceTuple_GetItemUnsafe(r0, r2)
+    r5 = unbox(int, r4)
+    x = r5
+    a = b
+L3:
+    r6 = r2 + 1
+    r2 = r6
+    goto L1
+L4:
+    return 1
+
 [case testConditionalFunctionDefinition]
 if int():
     def foo() -> int: