diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py index a622315e4808..c81c1360ff8b 100644 --- a/mypyc/irbuild/builder.py +++ b/mypyc/irbuild/builder.py @@ -1500,6 +1500,16 @@ def leave(self) -> tuple[list[Register], list[RuntimeArg], list[BasicBlock], RTy self.fn_info = self.fn_infos[-1] return builder.args, runtime_args, builder.blocks, ret_type, fn_info + @property + def code_fn_info(self) -> FuncInfo: + """Return info about the function whose code is being generated. + + This is fn_info, except in a comprehension scope, which is generated inline in the + enclosing function. Use this for state related to the generated code, such as the + generator class of a coroutine, and fn_info for name binding. + """ + return next(info for info in reversed(self.fn_infos) if not info.is_comprehension_scope) + @contextmanager def enter_scope(self, fn_info: FuncInfo) -> Iterator[None]: """Push a lightweight scope for comprehensions. @@ -1517,8 +1527,9 @@ def enter_scope(self, fn_info: FuncInfo) -> Iterator[None]: self.runtime_args.append([]) self.fn_info = fn_info self.fn_infos.append(self.fn_info) - self.ret_types.append(none_rprimitive) - self.nonlocal_control.append(BaseNonlocalControl()) + # Yields (from await) and nonlocal control flow belong to the enclosing function. + self.ret_types.append(self.ret_types[-1]) + self.nonlocal_control.append(self.nonlocal_control[-1]) try: yield finally: diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py index 42d2facaac9f..2d91d1f32591 100644 --- a/mypyc/irbuild/env_class.py +++ b/mypyc/irbuild/env_class.py @@ -106,15 +106,16 @@ def instantiate_env_class(builder: IRBuilder) -> Value: ) ) else: + if builder.fn_info.is_comprehension_scope and builder.code_fn_info.is_generator: + # The comprehension may contain an await, and only registers can be kept + # alive across one. + curr_env_reg = builder.ensure_register(curr_env_reg) # Top-level functions and comprehension scopes store env reg directly. builder.fn_info._curr_env_reg = curr_env_reg # Comprehension scopes link to parent env if it exists. parent = builder.fn_infos[-2] if builder.fn_info.is_nested else None if parent is not None and parent._env_class is not None: - if parent.is_generator: - parent_env_reg = parent.generator_class.curr_env_reg - else: - parent_env_reg = parent._curr_env_reg + parent_env_reg = get_curr_env_reg(parent) else: parent_env_reg = None if parent_env_reg is not None: @@ -125,6 +126,16 @@ def instantiate_env_class(builder: IRBuilder) -> Value: return curr_env_reg +def get_curr_env_reg(fn_info: FuncInfo) -> Value | None: + """Return the register holding the current environment of a function, if any.""" + if fn_info.is_generator: + return fn_info.generator_class.curr_env_reg + elif fn_info.is_nested and not fn_info.is_comprehension_scope: + # Nested functions store their environment on the callable class. + return fn_info.callable_class._curr_env_reg + return fn_info._curr_env_reg + + def load_env_registers(builder: IRBuilder, prefix: str = "") -> None: """Load the registers for the current FuncItem being visited. diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py index a59e6630aa78..ef5cbed7f66e 100644 --- a/mypyc/irbuild/expression.py +++ b/mypyc/irbuild/expression.py @@ -115,7 +115,9 @@ from mypyc.irbuild.constant_fold import constant_fold_expr from mypyc.irbuild.for_helpers import ( comprehension_helper, + comprehension_scope, raise_error_if_contains_unreachable_names, + translate_generator_to_list, translate_list_comprehension, translate_set_comprehension, translate_vec_comprehension, @@ -499,8 +501,7 @@ def implicit_super_arg(builder: IRBuilder) -> Var: """ # A comprehension can have a scope of its own, but it has no arguments, and # super() uses the function that contains the comprehension. - fn_info = next(info for info in reversed(builder.fn_infos) if not info.is_comprehension_scope) - return fn_info.fitem.arguments[0].variable + return builder.code_fn_info.fitem.arguments[0].variable def is_instance_method_self(builder: IRBuilder, var: Var) -> bool: @@ -1527,14 +1528,10 @@ def _visit_display( # Comprehensions # -# mypyc always inlines comprehensions (the loop body is emitted directly into -# the enclosing function's IR, no implicit function call like CPython). -# -# However, when a comprehension body contains a lambda, we need a lightweight -# scope boundary so the closure/env-class machinery can see the comprehension -# as a separate scope. The comprehension is still inlined (same basic blocks -# and registers), but we push a new FuncInfo and set up an env class so the -# lambda can capture loop variables through the standard env-class chain. +# mypyc always inlines list, set and dict comprehensions (the loop body is +# emitted directly into the enclosing function's IR, no implicit function call +# like CPython). When a comprehension body contains a lambda, it gets a +# lightweight scope (see for_helpers.comprehension_scope). def transform_list_comprehension(builder: IRBuilder, o: ListComprehension) -> Value: @@ -1542,11 +1539,8 @@ def transform_list_comprehension(builder: IRBuilder, o: ListComprehension) -> Va builder.error("Unpacking in comprehensions is not supported yet", o.line) return builder.none() gen = o.generator - if gen in builder.comprehension_to_fitem: - return _translate_comprehension_with_scope( - builder, gen, lambda: translate_list_comprehension(builder, gen) - ) - return translate_list_comprehension(builder, gen) + with comprehension_scope(builder, gen): + return translate_list_comprehension(builder, gen) def transform_set_comprehension(builder: IRBuilder, o: SetComprehension) -> Value: @@ -1554,20 +1548,16 @@ def transform_set_comprehension(builder: IRBuilder, o: SetComprehension) -> Valu builder.error("Unpacking in comprehensions is not supported yet", o.line) return builder.none() gen = o.generator - if gen in builder.comprehension_to_fitem: - return _translate_comprehension_with_scope( - builder, gen, lambda: translate_set_comprehension(builder, gen) - ) - return translate_set_comprehension(builder, gen) + with comprehension_scope(builder, gen): + return translate_set_comprehension(builder, gen) def transform_dictionary_comprehension(builder: IRBuilder, o: DictionaryComprehension) -> Value: if raise_error_if_contains_unreachable_names(builder, o): return builder.none() - if o in builder.comprehension_to_fitem: - return _translate_comprehension_with_scope(builder, o, lambda: _dict_comp_body(builder, o)) - return _dict_comp_body(builder, o) + with comprehension_scope(builder, o): + return _dict_comp_body(builder, o) def _dict_comp_body(builder: IRBuilder, o: DictionaryComprehension) -> Value: @@ -1587,31 +1577,6 @@ def gen_inner_stmts() -> None: return builder.read(d, o.line) -def _translate_comprehension_with_scope( - builder: IRBuilder, - node: GeneratorExpr | DictionaryComprehension, - gen_body: Callable[[], Value], -) -> Value: - """Wrap a comprehension body with a lightweight scope for closure capture.""" - from mypyc.irbuild.context import FuncInfo - from mypyc.irbuild.env_class import add_vars_to_env, finalize_env_class, setup_env_class - - comprehension_fdef = builder.comprehension_to_fitem[node] - fn_info = FuncInfo( - fitem=comprehension_fdef, - name=comprehension_fdef.name, - is_nested=True, - contains_nested=True, - is_comprehension_scope=True, - ) - - with builder.enter_scope(fn_info): - setup_env_class(builder) - finalize_env_class(builder) - add_vars_to_env(builder) - return gen_body() - - # Misc @@ -1631,17 +1596,7 @@ def transform_generator_expr(builder: IRBuilder, o: GeneratorExpr) -> Value: builder.error("Unpacking in comprehensions is not supported yet", o.line) return builder.none() builder.warning("Treating generator comprehension as list", o.line) - if o in builder.comprehension_to_fitem: - return builder.primitive_op( - iter_op, - [ - _translate_comprehension_with_scope( - builder, o, lambda: translate_list_comprehension(builder, o) - ) - ], - o.line, - ) - return builder.primitive_op(iter_op, [translate_list_comprehension(builder, o)], o.line) + return builder.primitive_op(iter_op, [translate_generator_to_list(builder, o)], o.line) def transform_assignment_expr(builder: IRBuilder, o: AssignmentExpr) -> Value: diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 072736f5fda3..7484f0c6331e 100644 --- a/mypyc/irbuild/for_helpers.py +++ b/mypyc/irbuild/for_helpers.py @@ -7,7 +7,8 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Iterator +from contextlib import contextmanager from typing import ClassVar, cast from mypy.nodes import ( @@ -70,6 +71,8 @@ ) from mypyc.irbuild.builder import IRBuilder 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.vec import vec_append, vec_create, vec_get_item_unsafe, vec_init_item_unsafe from mypyc.primitives.dict_ops import ( @@ -223,6 +226,20 @@ def for_loop_helper_with_index( builder.activate_block(exit_block) +def can_preallocate_from_generator(builder: IRBuilder, gen: GeneratorExpr) -> bool: + """Can a sequence be built from a generator expression with a preallocated length? + + This requires a single loop without conditions over a sequence of known length. + """ + if not (len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0): + return False + rtype = builder.node_type(gen.sequences[0]) + if isinstance(rtype, RTuple): + # for_loop_helper_with_index crashes for empty tuples + return bool(rtype.types) + return is_sequence_rprimitive(rtype) or isinstance(rtype, RVec) + + def sequence_from_generator_preallocate_helper( builder: IRBuilder, gen: GeneratorExpr, @@ -254,23 +271,16 @@ def sequence_from_generator_preallocate_helper( after_first_init: Optional function that generates instructions to run after the sequence has been evaluated (see comprehension_helper). """ - if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0: + if can_preallocate_from_generator(builder, gen): line = gen.line sequence_expr = gen.sequences[0] rtype = builder.node_type(sequence_expr) - if not (is_sequence_rprimitive(rtype) or isinstance(rtype, (RTuple, RVec))): - return None if isinstance(rtype, RTuple): # If input is RTuple, box it to tuple_rprimitive for generic iteration # TODO: this can be optimized a bit better with an unrolled ForRTuple helper proper_type = get_proper_type(builder.types[sequence_expr]) assert isinstance(proper_type, TupleType), proper_type - - # the for_loop_helper_with_index crashes for empty tuples, bail out - if not proper_type.items: - return None - proper_types = get_proper_types(proper_type.items) get_item_ops: list[LoadLiteral | TupleGet] @@ -369,7 +379,8 @@ def translate_generator_to_list( CPython evaluates them after the first iterable (when the generator object is created) but before the generator runs, so we do the same. """ - return translate_list_comprehension(builder, gen, after_first_init=before_loop) + with comprehension_scope(builder, gen): + return translate_list_comprehension(builder, gen, after_first_init=before_loop) def generator_loop_helper( @@ -385,10 +396,42 @@ def generator_loop_helper( translate_generator_to_list for the meaning of before_loop. """ loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async)) - comprehension_helper( - builder, loop_params, gen_inner_stmts, gen.line, after_first_init=before_loop + with comprehension_scope(builder, gen): + comprehension_helper( + builder, loop_params, gen_inner_stmts, gen.line, after_first_init=before_loop + ) + + +@contextmanager +def comprehension_scope( + builder: IRBuilder, node: GeneratorExpr | DictionaryComprehension +) -> Iterator[None]: + """Enter a lightweight scope for a comprehension, if it needs one. + + A comprehension gets a scope when it contains a lambda, so that the closure/env-class + machinery sees the comprehension as a separate scope and the lambda can capture loop + variables. The comprehension is still generated inline (same basic blocks and + registers). + """ + if node not in builder.comprehension_to_fitem: + yield + return + + comprehension_fdef = builder.comprehension_to_fitem[node] + fn_info = FuncInfo( + fitem=comprehension_fdef, + name=comprehension_fdef.name, + is_nested=True, + contains_nested=True, + is_comprehension_scope=True, ) + with builder.enter_scope(fn_info): + setup_env_class(builder) + finalize_env_class(builder) + add_vars_to_env(builder) + yield + def raise_error_if_contains_unreachable_names( builder: IRBuilder, gen: GeneratorExpr | DictionaryComprehension diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 86d4d40171c1..a5b8bf4f56b1 100644 --- a/mypyc/irbuild/generator.py +++ b/mypyc/irbuild/generator.py @@ -217,7 +217,7 @@ def add_raise_exception_blocks_to_generator_class(builder: IRBuilder, line: int) helper method for generator functions, and raises an exception if those flags are set. """ - cls = builder.fn_info.generator_class + cls = builder.code_fn_info.generator_class assert cls.exc_regs is not None exc_type, exc_val, exc_tb = cls.exc_regs diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py index 27f461acb785..520582f0a996 100644 --- a/mypyc/irbuild/specialize.py +++ b/mypyc/irbuild/specialize.py @@ -93,7 +93,8 @@ from mypyc.irbuild.builder import IRBuilder, find_walrus_targets, get_call_target_fullname from mypyc.irbuild.constant_fold import constant_fold_expr from mypyc.irbuild.for_helpers import ( - comprehension_helper, + can_preallocate_from_generator, + comprehension_scope, generator_loop_helper, get_expr_length_value, sequence_from_generator_preallocate_helper, @@ -441,17 +442,20 @@ def translate_list_from_generator_call( len(expr.args) == 1 and expr.arg_kinds[0] == ARG_POS and isinstance(expr.args[0], GeneratorExpr) + and can_preallocate_from_generator(builder, expr.args[0]) ): + gen = expr.args[0] def set_item(x: Value, y: Value, z: Value, line: int) -> None: builder.call_c(new_list_set_item_op, [x, y, z], line) - return sequence_from_generator_preallocate_helper( - builder, - expr.args[0], - empty_op_llbuilder=builder.builder.new_list_op_with_length, - set_item_op=set_item, - ) + with comprehension_scope(builder, gen): + return sequence_from_generator_preallocate_helper( + builder, + gen, + empty_op_llbuilder=builder.builder.new_list_op_with_length, + set_item_op=set_item, + ) return None @@ -480,17 +484,20 @@ def translate_tuple_from_generator_call( len(expr.args) == 1 and expr.arg_kinds[0] == ARG_POS and isinstance(expr.args[0], GeneratorExpr) + and can_preallocate_from_generator(builder, expr.args[0]) ): + gen = expr.args[0] def set_item(x: Value, y: Value, z: Value, line: int) -> None: builder.call_c(new_tuple_set_item_op, [x, y, z], line) - return sequence_from_generator_preallocate_helper( - builder, - expr.args[0], - empty_op_llbuilder=builder.builder.new_tuple_with_length, - set_item_op=set_item, - ) + with comprehension_scope(builder, gen): + return sequence_from_generator_preallocate_helper( + builder, + gen, + empty_op_llbuilder=builder.builder.new_tuple_with_length, + set_item_op=set_item, + ) return None @@ -508,7 +515,9 @@ def translate_set_from_generator_call( and expr.arg_kinds[0] == ARG_POS and isinstance(expr.args[0], GeneratorExpr) ): - return translate_set_comprehension(builder, expr.args[0]) + gen = expr.args[0] + with comprehension_scope(builder, gen): + return translate_set_comprehension(builder, gen) return None @@ -626,7 +635,6 @@ def any_all_helper( init_val = initial_value() retval = Register(bool_rprimitive, line=init_val.line) builder.assign(retval, init_val, init_val.line) - loop_params = list(zip(gen.indices, gen.sequences, gen.condlists, gen.is_async)) true_block, false_block, exit_block = BasicBlock(), BasicBlock(), BasicBlock() def gen_inner_stmts() -> None: @@ -638,7 +646,7 @@ def gen_inner_stmts() -> None: builder.goto(exit_block) builder.activate_block(false_block) - comprehension_helper(builder, loop_params, gen_inner_stmts, gen.line) + generator_loop_helper(builder, gen, gen_inner_stmts) builder.goto_and_activate(exit_block) return retval diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py index aac19cc8a75e..8f97fc4d0303 100644 --- a/mypyc/irbuild/statement.py +++ b/mypyc/irbuild/statement.py @@ -726,7 +726,7 @@ def transform_try_except( builder.builder.push_error_handler(double_except_block) builder.activate_block(except_entry) old_exc = builder.call_c(error_catch_op, [], line) - if builder.fn_info.is_generator: + if builder.code_fn_info.is_generator: # CFG cleanup may remove the CallC's block while resume cleanup paths remain. old_exc = builder.ensure_register(old_exc) # Compile the except blocks with the nonlocal control flow overridden to clear exc_info @@ -1312,7 +1312,7 @@ def transform_del_item(builder: IRBuilder, target: AssignmentTarget, line: int) def emit_yield(builder: IRBuilder, val: Value, line: int) -> Value: retval = builder.coerce(val, builder.ret_types[-1], line) - cls = builder.fn_info.generator_class + cls = builder.code_fn_info.generator_class # Create a new block for the instructions immediately following the yield expression, and # set the next label so that the next time '__next__' is called on the generator object, # the function continues at the new block. diff --git a/mypyc/test-data/run-generator-exprs.test b/mypyc/test-data/run-generator-exprs.test index 17d0f97e164b..b99a32699e2a 100644 --- a/mypyc/test-data/run-generator-exprs.test +++ b/mypyc/test-data/run-generator-exprs.test @@ -83,3 +83,117 @@ def test_exception_in_later_argument() -> None: with assertRaises(ValueError, "fail"): max((elt(x) for x in it()), key=fail()) check(["iterable", "fail"]) + +def test_lambda_in_generator() -> None: + assert sum(((lambda: elt(x))() for x in it()), arg("start", 10)) == 13 + check(["iterable", "start", "elt 1", "elt 2"]) + assert next(((lambda: elt(x))() for x in it()), arg("default", -1)) == 1 + check(["iterable", "default", "elt 1"]) + n = -1 + assert max(((lambda: elt(x))() for x in it()), key=arg("key", lambda v: v * n)) == 1 + check(["iterable", "key", "elt 1", "elt 2"]) + +[case testLambdaInSpecializedGeneratorExpr] +from typing import Callable + +def test_lambdas_capture_loop_variable() -> None: + fs = tuple((lambda: x) for x in range(3)) + assert [f() for f in fs] == [2, 2, 2] + gs = list((lambda: x) for x in range(3)) + assert [g() for g in gs] == [2, 2, 2] + assert sorted(f() for f in [(lambda: x) for x in range(3)]) == [2, 2, 2] + +def test_lambdas_called_while_iterating() -> None: + assert sum((lambda: x)() for x in range(4)) == 6 + assert sum(((lambda: x)() for x in range(4)), 10) == 16 + assert any((lambda: x)() == 2 for x in range(3)) + assert not all((lambda: x)() < 2 for x in range(3)) + assert next((lambda: x)() for x in range(5) if x > 2) == 3 + assert next(((lambda: x)() for x in range(5) if x > 10), -1) == -1 + assert set((lambda: x)() % 2 for x in range(5)) == {0, 1} + assert tuple((lambda: x)() for x in [5, 6]) == (5, 6) + assert list((lambda: x)() for x in "ab") == ["a", "b"] + assert max(((lambda: x)() for x in [3, 1, 2]), key=lambda v: -v) == 1 + assert ",".join((lambda: s)() for s in ["a", "b"]) == "a,b" + +def test_lambda_captures_outer_variable() -> None: + n = 10 + assert sum((lambda: x + n)() for x in range(3)) == 33 + fs: list[Callable[[], int]] = [] + fs.extend((lambda: x * n) for x in range(3)) + assert [f() for f in fs] == [20, 20, 20] + +[case testLambdaInGeneratorExprInNestedFunction] +from typing import Callable, Iterator + +def test_nested_function() -> None: + n = 7 + def inner() -> list[int]: + a = [f() for f in tuple((lambda: n + x) for x in range(3))] + b = [f() for f in [(lambda: n + x) for x in range(3)]] + c = [f() for f in tuple((lambda y=x: y + n) for x in range(3))] + return a + b + c + [sum((lambda: n + x)() for x in range(3))] + assert inner() == [9, 9, 9, 9, 9, 9, 7, 8, 9, 24] + +def test_nested_generator() -> None: + n = 7 + def inner() -> Iterator[int]: + yield sum((lambda: n + x)() for x in range(3)) + for f in tuple((lambda: n + x) for x in range(3)): + yield f() + assert list(inner()) == [24, 9, 9, 9] + +def test_default_bound_lambdas() -> None: + data = [1, 2] + assert [f() for f in list((lambda y=x: y) for x in data)] == [1, 2] + assert [f() for f in tuple((lambda y=x: y) for x in data)] == [1, 2] + +[case testLambdaInGeneratorExprWithAwait] +import asyncio +from typing import Callable + +async def items() -> list[int]: + await asyncio.sleep(0) + return [1, 2] + +async def value(x: int) -> int: + await asyncio.sleep(0) + return x + +async def first_iterable() -> int: + return sum((lambda: x)() for x in await items()) + +async def consumer_arg() -> int: + return sum(((lambda: x)() for x in range(3)), await value(10)) + +async def list_comprehension() -> list[int]: + return [(lambda: x)() for x in await items()] + +async def await_in_body() -> list[int]: + return [(lambda: x)() + await value(10) for x in await items() if await value(x) > 1] + +async def capture_across_await() -> list[int]: + fs: list[Callable[[], int]] = [(lambda: x + y) for x in await items() for y in [await value(10)]] + return [f() for f in fs] + +async def nested() -> list[int]: + n = 3 + async def inner() -> list[int]: + return [(lambda: x + n)() + await value(10) for x in await items()] + return await inner() + +def test_await() -> None: + assert asyncio.run(first_iterable()) == 3 + assert asyncio.run(consumer_arg()) == 13 + assert asyncio.run(list_comprehension()) == [1, 2] + assert asyncio.run(await_in_body()) == [12] + assert asyncio.run(capture_across_await()) == [12, 12] + assert asyncio.run(nested()) == [14, 15] + +[file asyncio/__init__.pyi] +from typing import Any, Awaitable, TypeVar +T = TypeVar("T") +async def sleep(t: float) -> None: ... +def run(x: Awaitable[T]) -> T: ... + +[typing fixtures/typing-full.pyi]