[mypyc] Refactor IR building for generator functions

mypyc/irbuild/builder.py

@@ -1169,7 +1169,7 @@ def flatten_classes(self, arg: RefExpr | TupleExpr) -> list[ClassIR] | None:
                     return None
             return res
 
-    def enter(self, fn_info: FuncInfo | str = "") -> None:
+    def enter(self, fn_info: FuncInfo | str = "", *, ret_type: RType = none_rprimitive) -> None:
         if isinstance(fn_info, str):
             fn_info = FuncInfo(name=fn_info)
         self.builder = LowLevelIRBuilder(self.errors, self.options)
@@ -1179,7 +1179,7 @@ def enter(self, fn_info: FuncInfo | str = "") -> None:
         self.runtime_args.append([])
         self.fn_info = fn_info
         self.fn_infos.append(self.fn_info)
-        self.ret_types.append(none_rprimitive)
+        self.ret_types.append(ret_type)
         if fn_info.is_generator:
             self.nonlocal_control.append(GeneratorNonlocalControl())
         else:
@@ -1219,10 +1219,9 @@ def enter_method(
             self_type: If not None, override default type of the implicit 'self'
                 argument (by default, derive type from class_ir)
         """
-        self.enter(fn_info)
+        self.enter(fn_info, ret_type=ret_type)
         self.function_name_stack.append(name)
         self.class_ir_stack.append(class_ir)
-        self.ret_types[-1] = ret_type
         if self_type is None:
             self_type = RInstance(class_ir)
         self.add_argument(SELF_NAME, self_type)

mypyc/irbuild/env_class.py

@@ -191,6 +191,41 @@ def add_args_to_env(
                 builder.add_var_to_env_class(arg.variable, rtype, base, reassign=reassign)
 
 
+def add_vars_to_env(builder: IRBuilder) -> None:
+    """Add relevant local variables and nested functions to the environment class.
+
+    Add all variables and functions that are declared/defined within current
+    function and are referenced in functions nested within this one to this
+    function's environment class so the nested functions can reference
+    them even if they are declared after the nested function's definition.
+    Note that this is done before visiting the body of the function.
+    """
+    env_for_func: FuncInfo | ImplicitClass = builder.fn_info
+    if builder.fn_info.is_generator:
+        env_for_func = builder.fn_info.generator_class
+    elif builder.fn_info.is_nested or builder.fn_info.in_non_ext:
+        env_for_func = builder.fn_info.callable_class
+
+    if builder.fn_info.fitem in builder.free_variables:
+        # Sort the variables to keep things deterministic
+        for var in sorted(builder.free_variables[builder.fn_info.fitem], key=lambda x: x.name):
+            if isinstance(var, Var):
+                rtype = builder.type_to_rtype(var.type)
+                builder.add_var_to_env_class(var, rtype, env_for_func, reassign=False)
+
+    if builder.fn_info.fitem in builder.encapsulating_funcs:
+        for nested_fn in builder.encapsulating_funcs[builder.fn_info.fitem]:
+            if isinstance(nested_fn, FuncDef):
+                # The return type is 'object' instead of an RInstance of the
+                # callable class because differently defined functions with
+                # the same name and signature across conditional blocks
+                # will generate different callable classes, so the callable
+                # class that gets instantiated must be generic.
+                builder.add_var_to_env_class(
+                    nested_fn, object_rprimitive, env_for_func, reassign=False
+                )
+
+
 def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: ImplicitClass) -> None:
     """Enable calling a nested function (with a callable class) recursively.
 

mypyc/irbuild/function.py

@@ -69,22 +69,14 @@
     instantiate_callable_class,
     setup_callable_class,
 )
-from mypyc.irbuild.context import FuncInfo, ImplicitClass
+from mypyc.irbuild.context import FuncInfo
 from mypyc.irbuild.env_class import (
+    add_vars_to_env,
     finalize_env_class,
     load_env_registers,
-    load_outer_envs,
     setup_env_class,
-    setup_func_for_recursive_call,
-)
-from mypyc.irbuild.generator import (
-    add_methods_to_generator_class,
-    add_raise_exception_blocks_to_generator_class,
-    create_switch_for_generator_class,
-    gen_generator_func,
-    populate_switch_for_generator_class,
-    setup_env_for_generator_class,
 )
+from mypyc.irbuild.generator import gen_generator_func, gen_generator_func_body
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.irbuild.util import is_constant
 from mypyc.primitives.dict_ops import dict_get_method_with_none, dict_new_op, dict_set_item_op
@@ -235,123 +227,81 @@ def c() -> None:
         func_name = singledispatch_main_func_name(name)
     else:
         func_name = name
-    builder.enter(
-        FuncInfo(
-            fitem=fitem,
-            name=func_name,
-            class_name=class_name,
-            namespace=gen_func_ns(builder),
-            is_nested=is_nested,
-            contains_nested=contains_nested,
-            is_decorated=is_decorated,
-            in_non_ext=in_non_ext,
-            add_nested_funcs_to_env=add_nested_funcs_to_env,
-        )
+
+    fn_info = FuncInfo(
+        fitem=fitem,
+        name=func_name,
+        class_name=class_name,
+        namespace=gen_func_ns(builder),
+        is_nested=is_nested,
+        contains_nested=contains_nested,
+        is_decorated=is_decorated,
+        in_non_ext=in_non_ext,
+        add_nested_funcs_to_env=add_nested_funcs_to_env,
     )
+    is_generator = fn_info.is_generator
+    builder.enter(fn_info, ret_type=sig.ret_type)
 
     # Functions that contain nested functions need an environment class to store variables that
     # are free in their nested functions. Generator functions need an environment class to
     # store a variable denoting the next instruction to be executed when the __next__ function
     # is called, along with all the variables inside the function itself.
-    if builder.fn_info.contains_nested or builder.fn_info.is_generator:
+    if contains_nested or is_generator:
         setup_env_class(builder)
 
-    if builder.fn_info.is_nested or builder.fn_info.in_non_ext:
+    if is_nested or in_non_ext:
         setup_callable_class(builder)
 
-    if builder.fn_info.is_generator:
-        # Do a first-pass and generate a function that just returns a generator object.
-        gen_generator_func(builder)
-        args, _, blocks, ret_type, fn_info = builder.leave()
-        func_ir, func_reg = gen_func_ir(
-            builder, args, blocks, sig, fn_info, cdef, is_singledispatch
+    if is_generator:
+        # First generate a function that just constructs and returns a generator object.
+        func_ir, func_reg = gen_generator_func(
+            builder,
+            lambda args, blocks, fn_info: gen_func_ir(
+                builder, args, blocks, sig, fn_info, cdef, is_singledispatch
+            ),
         )
 
         # Re-enter the FuncItem and visit the body of the function this time.
-        builder.enter(fn_info)
-        setup_env_for_generator_class(builder)
+        symtable = gen_generator_func_body(builder, fn_info, sig)
 
-        load_outer_envs(builder, builder.fn_info.generator_class)
-        top_level = builder.top_level_fn_info()
-        if (
-            builder.fn_info.is_nested
-            and isinstance(fitem, FuncDef)
-            and top_level
-            and top_level.add_nested_funcs_to_env
-        ):
-            setup_func_for_recursive_call(builder, fitem, builder.fn_info.generator_class)
-        create_switch_for_generator_class(builder)
-        add_raise_exception_blocks_to_generator_class(builder, fitem.line)
+        # Evaluate argument defaults in the surrounding scope, since we
+        # calculate them *once* when the function definition is evaluated.
+        calculate_arg_defaults(builder, fn_info, func_reg, symtable)
     else:
-        load_env_registers(builder)
-        gen_arg_defaults(builder)
+        func_ir, func_reg = gen_func_body(builder, sig, cdef, is_singledispatch)
 
-    if builder.fn_info.contains_nested and not builder.fn_info.is_generator:
-        finalize_env_class(builder)
+    if is_singledispatch:
+        # add the generated main singledispatch function
+        builder.functions.append(func_ir)
+        # create the dispatch function
+        assert isinstance(fitem, FuncDef)
+        return gen_dispatch_func_ir(builder, fitem, fn_info.name, name, sig)
 
-    builder.ret_types[-1] = sig.ret_type
+    return func_ir, func_reg
 
-    # Add all variables and functions that are declared/defined within this
-    # function and are referenced in functions nested within this one to this
-    # function's environment class so the nested functions can reference
-    # them even if they are declared after the nested function's definition.
-    # Note that this is done before visiting the body of this function.
-
-    env_for_func: FuncInfo | ImplicitClass = builder.fn_info
-    if builder.fn_info.is_generator:
-        env_for_func = builder.fn_info.generator_class
-    elif builder.fn_info.is_nested or builder.fn_info.in_non_ext:
-        env_for_func = builder.fn_info.callable_class
-
-    if builder.fn_info.fitem in builder.free_variables:
-        # Sort the variables to keep things deterministic
-        for var in sorted(builder.free_variables[builder.fn_info.fitem], key=lambda x: x.name):
-            if isinstance(var, Var):
-                rtype = builder.type_to_rtype(var.type)
-                builder.add_var_to_env_class(var, rtype, env_for_func, reassign=False)
-
-    if builder.fn_info.fitem in builder.encapsulating_funcs:
-        for nested_fn in builder.encapsulating_funcs[builder.fn_info.fitem]:
-            if isinstance(nested_fn, FuncDef):
-                # The return type is 'object' instead of an RInstance of the
-                # callable class because differently defined functions with
-                # the same name and signature across conditional blocks
-                # will generate different callable classes, so the callable
-                # class that gets instantiated must be generic.
-                builder.add_var_to_env_class(
-                    nested_fn, object_rprimitive, env_for_func, reassign=False
-                )
 
-    builder.accept(fitem.body)
+def gen_func_body(
+    builder: IRBuilder, sig: FuncSignature, cdef: ClassDef | None, is_singledispatch: bool
+) -> tuple[FuncIR, Value | None]:
+    load_env_registers(builder)
+    gen_arg_defaults(builder)
+    if builder.fn_info.contains_nested:
+        finalize_env_class(builder)
+    add_vars_to_env(builder)
+    builder.accept(builder.fn_info.fitem.body)
     builder.maybe_add_implicit_return()
 
-    if builder.fn_info.is_generator:
-        populate_switch_for_generator_class(builder)
-
     # Hang on to the local symbol table for a while, since we use it
     # to calculate argument defaults below.
     symtable = builder.symtables[-1]
 
     args, _, blocks, ret_type, fn_info = builder.leave()
 
-    if fn_info.is_generator:
-        add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
-    else:
-        func_ir, func_reg = gen_func_ir(
-            builder, args, blocks, sig, fn_info, cdef, is_singledispatch
-        )
+    func_ir, func_reg = gen_func_ir(builder, args, blocks, sig, fn_info, cdef, is_singledispatch)
 
     # Evaluate argument defaults in the surrounding scope, since we
     # calculate them *once* when the function definition is evaluated.
     calculate_arg_defaults(builder, fn_info, func_reg, symtable)
-
-    if is_singledispatch:
-        # add the generated main singledispatch function
-        builder.functions.append(func_ir)
-        # create the dispatch function
-        assert isinstance(fitem, FuncDef)
-        return gen_dispatch_func_ir(builder, fitem, fn_info.name, name, sig)
-
     return func_ir, func_reg
 
 

mypyc/irbuild/generator.py

@@ -10,7 +10,9 @@
 
 from __future__ import annotations
 
-from mypy.nodes import ARG_OPT, Var
+from typing import Callable
+
+from mypy.nodes import ARG_OPT, FuncDef, SymbolNode, Var
 from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME, SELF_NAME
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
@@ -31,13 +33,16 @@
     Value,
 )
 from mypyc.ir.rtypes import RInstance, int_rprimitive, object_rprimitive
-from mypyc.irbuild.builder import IRBuilder, gen_arg_defaults
+from mypyc.irbuild.builder import IRBuilder, SymbolTarget, gen_arg_defaults
 from mypyc.irbuild.context import FuncInfo, GeneratorClass
 from mypyc.irbuild.env_class import (
     add_args_to_env,
+    add_vars_to_env,
     finalize_env_class,
     load_env_registers,
     load_outer_env,
+    load_outer_envs,
+    setup_func_for_recursive_call,
 )
 from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl
 from mypyc.primitives.exc_ops import (
@@ -49,13 +54,67 @@
 )
 
 
-def gen_generator_func(builder: IRBuilder) -> None:
+def gen_generator_func(
+    builder: IRBuilder,
+    gen_func_ir: Callable[
+        [list[Register], list[BasicBlock], FuncInfo], tuple[FuncIR, Value | None]
+    ],
+) -> tuple[FuncIR, Value | None]:
+    """Generate IR for generator function that returns generator object."""
     setup_generator_class(builder)
     load_env_registers(builder)
     gen_arg_defaults(builder)
     finalize_env_class(builder)
     builder.add(Return(instantiate_generator_class(builder)))
 
+    args, _, blocks, ret_type, fn_info = builder.leave()
+    func_ir, func_reg = gen_func_ir(args, blocks, fn_info)
+    return func_ir, func_reg
+
+
+def gen_generator_func_body(
+    builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature
+) -> dict[SymbolNode, SymbolTarget]:
+    """Generate IR based on the body of a generator function.
+
+    Add "__next__", "__iter__" and other generator methods to the generator
+    class that implements the function (each function gets a separate class).
+
+    Return the symbol table for the body.
+    """
+    builder.enter(fn_info, ret_type=sig.ret_type)
+    setup_env_for_generator_class(builder)
+
+    load_outer_envs(builder, builder.fn_info.generator_class)
+    top_level = builder.top_level_fn_info()
+    fitem = fn_info.fitem
+    if (
+        builder.fn_info.is_nested
+        and isinstance(fitem, FuncDef)
+        and top_level
+        and top_level.add_nested_funcs_to_env
+    ):
+        setup_func_for_recursive_call(builder, fitem, builder.fn_info.generator_class)
+    create_switch_for_generator_class(builder)
+    add_raise_exception_blocks_to_generator_class(builder, fitem.line)
+
+    add_vars_to_env(builder)
+
+    builder.accept(fitem.body)
+    builder.maybe_add_implicit_return()
+
+    populate_switch_for_generator_class(builder)
+
+    # Hang on to the local symbol table, since the caller will use it
+    # to calculate argument defaults.
+    symtable = builder.symtables[-1]
+
+    args, _, blocks, ret_type, fn_info = builder.leave()
+
+    add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
+
+    return symtable
+
 
 def instantiate_generator_class(builder: IRBuilder) -> Value:
     fitem = builder.fn_info.fitem