[mypyc] Fix for loop iteration when the loop body reassigns the iterable (#22137)

Fixes mypyc/mypyc#1233

`for` loops over lists, tuples, strings, bytes, dicts and `range()`
(also inside `enumerate()`, `zip()` and `reversed()`) read the iterable,
or the end of the range, on every iteration. When it was a local
variable, the loop read the variable's register, so assigning to the
variable in the loop body changed the iteration: `for i in range(n): n =
0` ran once, `for x in a: a = [9]` stopped after the first item, and a
loop over a tuple read past the end of the new tuple, which could
segfault.

Now, if the loop target or body assigns to the local variable that holds
the iterable or the end of the range, its value is copied to a new
register before the loop. A small AST visitor finds these assignments:
plain, augmented and starred assignments, `:=`, `del`, `for` targets,
`with ... as`, `except ... as` and `match` capture patterns. Other loops
generate the same IR as before. Copying unconditionally would add a copy
to every loop over a local variable and change the IR of 23 existing
test cases.

The comprehension form, `[x for x in t if (t := (9,))]`, is fixed
separately by #22133.

Author: rheard

mypyc/irbuild/for_helpers.py

@@ -7,27 +7,37 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Set as AbstractSet
 from contextlib import contextmanager
 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
 from mypyc.ir.ops import (
@@ -73,7 +83,7 @@
 from mypyc.irbuild.constant_fold import constant_fold_expr
 from mypyc.irbuild.context import FuncInfo
 from mypyc.irbuild.env_class import add_vars_to_env, finalize_env_class, setup_env_class
-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,
@@ -110,6 +120,7 @@ def for_loop_helper(
     is_async: bool,
     line: int,
     after_init: GenFunc | None = None,
+    reassigned: AbstractSet[SymbolNode] = frozenset(),
 ) -> None:
     """Generate IR for a loop.
 
@@ -120,6 +131,8 @@ def for_loop_helper(
         else_insts: a function that generates the else block instructions
         after_init: a function that generates instructions to run after the iterable has
             been evaluated (and iter() called on it) but before the first iteration
+        reassigned: local variables that the loop index or body may assign to
+            (see find_loop_assignments)
     """
     # Body of the loop
     body_block = BasicBlock()
@@ -134,7 +147,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,
     )
     if after_init is not None:
         after_init()
@@ -584,6 +604,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,
@@ -593,10 +702,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
@@ -611,7 +724,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)
@@ -620,7 +733,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)
@@ -642,10 +755,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
@@ -669,7 +782,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 (
@@ -681,7 +794,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 (
@@ -691,7 +804,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)
@@ -701,7 +814,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)
@@ -1362,15 +1475,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:
@@ -1398,15 +1524,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: