diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index a3c7349993af..fdb274a4a664 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/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 " 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 ". - 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()". - 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,7 +1370,13 @@ 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 @@ -1265,7 +1384,14 @@ def init(self, index1: Lvalue, index2: Lvalue, expr: Expression) -> None: 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,7 +1419,12 @@ 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. @@ -1301,7 +1432,14 @@ def init(self, indexes: list[Lvalue], exprs: list[Expression]) -> None: 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) diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index aac19cc8a75e..19d78804abf2 100644 --- a/mypyc/irbuild/statement.py +++ b/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), ) diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test index 68e0e2d79f9c..e5cf6e07bec1 100644 --- a/mypyc/test-data/irbuild-statements.test +++ b/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: diff --git a/mypyc/test-data/run-loops.test b/mypyc/test-data/run-loops.test index 5f2f0ae52101..9c03e5058742 100644 --- a/mypyc/test-data/run-loops.test +++ b/mypyc/test-data/run-loops.test @@ -281,6 +281,240 @@ expected = [0, 1, 2, 2, 0, 1, 2, 2, 0, 1, 2, 2] assert list(gen) == expected [out] +[case testForLoopReassignIterable] +from typing import Any, Iterator + +from mypy_extensions import i64 + +class Ctx: + def __enter__(self) -> tuple[int, ...]: + return (9,) + + def __exit__(self, *args: Any) -> None: + pass + +def test_sequence() -> None: + t: tuple[int, ...] = (1, 2, 3) + items = [] + for x in t: + t = (9,) + items.append(x) + assert items == [1, 2, 3] + + a = [1, 2, 3] + items = [] + for x in a: + a = [9] + items.append(x) + assert items == [1, 2, 3] + + s = "abc" + chars = [] + for c in s: + s = "z" + chars.append(c) + assert chars == ["a", "b", "c"] + + b = b"abc" + items = [] + for x in b: + b = b"z" + items.append(x) + assert items == [97, 98, 99] + +def test_loop_target_is_iterable() -> None: + s = "abc" + chars = [] + for s in s: + chars.append(s) + assert chars == ["a", "b", "c"] + +def test_reversed_enumerate_zip() -> None: + a = [1, 2, 3] + items = [] + for x in reversed(a): + a = [9] + items.append(x) + assert items == [3, 2, 1] + + a = [1, 2, 3] + pairs = [] + for i, x in enumerate(a): + a = [9] + pairs.append((i, x)) + assert pairs == [(0, 1), (1, 2), (2, 3)] + + a = [1, 2, 3] + b = ["a", "b", "c"] + zipped = [] + for x, y in zip(a, b): + a = [9] + b = ["z"] + zipped.append((x, y)) + assert zipped == [(1, "a"), (2, "b"), (3, "c")] + +def test_dict() -> None: + d = {"a": 1, "b": 2} + keys = [] + for k in d: + d = {"z": 0} + keys.append(k) + assert keys == ["a", "b"] + + d = {"a": 1, "b": 2} + keys = [] + for k in d.keys(): + d = {"z": 0} + keys.append(k) + assert keys == ["a", "b"] + + d = {"a": 1, "b": 2} + values = [] + for v in d.values(): + d = {"z": 0} + values.append(v) + assert values == [1, 2] + + d = {"a": 1, "b": 2} + items = [] + for k, v in d.items(): + d = {"z": 0} + items.append((k, v)) + assert items == [("a", 1), ("b", 2)] + + # The dict that is iterated over is still checked for size changes + d = {"a": 1, "b": 2} + try: + for k in d: + old = d + d = {"x": 1, "y": 2} + old["c"] = 3 + except RuntimeError as e: + assert str(e) == "dictionary changed size during iteration" + else: + assert False + +def range_end_arg(n: int) -> list[int]: + items = [] + for i in range(n): + n = 0 + items.append(i) + return items + +def test_range() -> None: + assert range_end_arg(3) == [0, 1, 2] + + n = 5 + items = [] + for i in range(n): + n -= 1 + items.append(i) + assert items == [0, 1, 2, 3, 4] + + lo = 0 + items = [] + for i in range(3, lo, -1): + lo = 5 + items.append(i) + assert items == [3, 2, 1] + + n = 3 + pairs = [] + for i, j in zip(range(n), range(10)): + n = 0 + pairs.append((i, j)) + assert pairs == [(0, 0), (1, 1), (2, 2)] + + m: i64 = 3 + items64 = [] + for k in range(m): + m = 0 + items64.append(k) + assert items64 == [0, 1, 2] + +def test_other_assignments() -> None: + t: tuple[int, ...] = (1, 2, 3) + items = [] + for x in t: + if x == 1: + del t + items.append(x) + assert items == [1, 2, 3] + + t = (1, 2, 3) + items = [] + for x in t: + if t := (9,): + items.append(x) + assert items == [1, 2, 3] + + t = (1, 2, 3) + items = [] + for x in t: + [(t := (9,)) for _ in range(1)] + items.append(x) + assert items == [1, 2, 3] + + t = (1, 2, 3) + items = [] + for x in t: + for t in [(9,)]: + pass + items.append(x) + assert items == [1, 2, 3] + + t = (1, 2, 3) + items = [] + for x in t: + with Ctx() as t: + pass + items.append(x) + assert items == [1, 2, 3] + + a = [1, 2, 3] + items = [] + for x in a: + first, *a = [9, 9] + items.append(x) + assert items == [1, 2, 3] + +def test_match() -> None: + t: tuple[int, ...] = (1, 2, 3) + items = [] + for x in t: + match (9,): + case t: + pass + items.append(x) + assert items == [1, 2, 3] + + a = [1, 2, 3] + items = [] + for x in a: + match [9]: + case [*a]: + pass + items.append(x) + assert items == [1, 2, 3] + + d = {"a": 1, "b": 2} + keys = [] + for k in d: + match {"z": 0}: + case {**d}: + pass + keys.append(k) + assert keys == ["a", "b"] + +def gen() -> Iterator[int]: + t: tuple[int, ...] = (1, 2, 3) + for x in t: + t = (9,) + yield x + +def test_generator() -> None: + assert list(gen()) == [1, 2, 3] + [case testRangeLoopVariableAfterLoop] for i in range(3): pass