[mypyc] Fix lambdas in generator expressions consumed by builtins

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:

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.
 

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,47 +1528,36 @@ 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:
     if isinstance(o.generator.left_expr, StarExpr):
         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:
     if isinstance(o.generator.left_expr, StarExpr):
         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:

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

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