From fb2d0054acc3bf0c7756b301ff1596b0da326ee7 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 14:51:11 +0100
Subject: [PATCH 01/17] [mypyc] Minor refactoring

---
 mypyc/irbuild/function.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index 70e494f063b8..c39a9df9bb68 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -248,18 +248,19 @@ def c() -> None:
             add_nested_funcs_to_env=add_nested_funcs_to_env,
         )
     )
-
+    is_generator = builder.fn_info.is_generator
+    
     # 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:
+    if 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()
@@ -286,7 +287,7 @@ def c() -> None:
         load_env_registers(builder)
         gen_arg_defaults(builder)
 
-    if builder.fn_info.contains_nested and not builder.fn_info.is_generator:
+    if contains_nested and not is_generator:
         finalize_env_class(builder)
 
     builder.ret_types[-1] = sig.ret_type
@@ -298,9 +299,9 @@ def c() -> None:
     # 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:
+    if is_generator:
         env_for_func = builder.fn_info.generator_class
-    elif builder.fn_info.is_nested or builder.fn_info.in_non_ext:
+    elif is_nested or in_non_ext:
         env_for_func = builder.fn_info.callable_class
 
     if builder.fn_info.fitem in builder.free_variables:
@@ -325,7 +326,7 @@ def c() -> None:
     builder.accept(fitem.body)
     builder.maybe_add_implicit_return()
 
-    if builder.fn_info.is_generator:
+    if is_generator:
         populate_switch_for_generator_class(builder)
 
     # Hang on to the local symbol table for a while, since we use it
@@ -334,7 +335,7 @@ def c() -> None:
 
     args, _, blocks, ret_type, fn_info = builder.leave()
 
-    if fn_info.is_generator:
+    if is_generator:
         add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
     else:
         func_ir, func_reg = gen_func_ir(

From 9fb14eadc2978544dc0848456f66c31e2faefb07 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 14:56:08 +0100
Subject: [PATCH 02/17] Minor refactoring

---
 mypyc/irbuild/function.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index c39a9df9bb68..29c2a26311f6 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -249,7 +249,7 @@ def c() -> None:
         )
     )
     is_generator = builder.fn_info.is_generator
-    
+
     # 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
@@ -286,9 +286,8 @@ def c() -> None:
     else:
         load_env_registers(builder)
         gen_arg_defaults(builder)
-
-    if contains_nested and not is_generator:
-        finalize_env_class(builder)
+        if contains_nested:
+            finalize_env_class(builder)
 
     builder.ret_types[-1] = sig.ret_type
 

From 90af72271c7459fe61f89984be3a09815e0c18e6 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 15:04:50 +0100
Subject: [PATCH 03/17] Refactor more

---
 mypyc/irbuild/function.py | 64 +++++++++++++++++++++------------------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index 29c2a26311f6..bff4bc5471f5 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -291,36 +291,7 @@ def c() -> None:
 
     builder.ret_types[-1] = sig.ret_type
 
-    # 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 is_generator:
-        env_for_func = builder.fn_info.generator_class
-    elif is_nested or 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
-                )
+    add_vars_to_env(builder)
 
     builder.accept(fitem.body)
     builder.maybe_add_implicit_return()
@@ -355,6 +326,39 @@ def c() -> None:
     return func_ir, func_reg
 
 
+def add_vars_to_env(builder: IRBuilder) -> None:
+    # 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
+                )
+
+
 def has_nested_func_self_reference(builder: IRBuilder, fitem: FuncItem) -> bool:
     """Does a nested function contain a self-reference in its body?
 

From edbe7ca43dad6ef727613b385c02cf9a6a0a66a7 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 15:12:52 +0100
Subject: [PATCH 04/17] Refactor more

---
 mypyc/irbuild/function.py  | 12 +++---------
 mypyc/irbuild/generator.py | 16 +++++++++++++++-
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index bff4bc5471f5..ddde07d3fe92 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -262,15 +262,9 @@ def c() -> None:
 
     if 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
-        )
-
-        # Re-enter the FuncItem and visit the body of the function this time.
-        builder.enter(fn_info)
-        setup_env_for_generator_class(builder)
+        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))
 
         load_outer_envs(builder, builder.fn_info.generator_class)
         top_level = builder.top_level_fn_info()
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index bc61c4493d55..63962b298835 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -10,6 +10,8 @@
 
 from __future__ import annotations
 
+from typing import Callable
+
 from mypy.nodes import ARG_OPT, Var
 from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME, SELF_NAME
 from mypyc.ir.class_ir import ClassIR
@@ -49,13 +51,25 @@
 )
 
 
-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]:
     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)
+
+    # Re-enter the FuncItem and visit the body of the function this time.
+    builder.enter(fn_info)
+    setup_env_for_generator_class(builder)
+
+    return func_ir, func_reg
+
 
 def instantiate_generator_class(builder: IRBuilder) -> Value:
     fitem = builder.fn_info.fitem

From 476ebe1e858dc9f20dd5949ef251456d57776fa1 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 15:36:56 +0100
Subject: [PATCH 05/17] Refactor

---
 mypyc/irbuild/builder.py   |  4 ++--
 mypyc/irbuild/function.py  | 32 +++++++++++++++++---------------
 mypyc/irbuild/generator.py |  4 ----
 3 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
index d9d3c5ed9cd0..edb873237bf2 100644
--- a/mypyc/irbuild/builder.py
+++ b/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:
diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index ddde07d3fe92..f4c711690868 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -235,20 +235,20 @@ 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 = builder.fn_info.is_generator
+    is_generator = fn_info.is_generator
+    builder.enter(fn_info, ret_type=sig.ret_type if not is_generator else object_rprimitive)
 
     # 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
@@ -266,6 +266,10 @@ def c() -> None:
             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, 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()
         if (
@@ -283,8 +287,6 @@ def c() -> None:
         if contains_nested:
             finalize_env_class(builder)
 
-    builder.ret_types[-1] = sig.ret_type
-
     add_vars_to_env(builder)
 
     builder.accept(fitem.body)
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index 63962b298835..a950a9fdf1e1 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -64,10 +64,6 @@ def gen_generator_func(
     args, _, blocks, ret_type, fn_info = builder.leave()
     func_ir, func_reg = gen_func_ir(args, blocks, fn_info)
 
-    # Re-enter the FuncItem and visit the body of the function this time.
-    builder.enter(fn_info)
-    setup_env_for_generator_class(builder)
-
     return func_ir, func_reg
 
 

From 756e3627401e53d7c24e7821d9ac573ec6af07f1 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 15:40:08 +0100
Subject: [PATCH 06/17] Refactor

---
 mypyc/irbuild/function.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index f4c711690868..717f1d1da2ac 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -248,7 +248,7 @@ def c() -> None:
         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 if not is_generator else object_rprimitive)
+    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
@@ -281,19 +281,23 @@ def c() -> None:
             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)
     else:
         load_env_registers(builder)
         gen_arg_defaults(builder)
         if contains_nested:
             finalize_env_class(builder)
 
-    add_vars_to_env(builder)
+        add_vars_to_env(builder)
 
-    builder.accept(fitem.body)
-    builder.maybe_add_implicit_return()
-
-    if is_generator:
-        populate_switch_for_generator_class(builder)
+        builder.accept(fitem.body)
+        builder.maybe_add_implicit_return()
 
     # Hang on to the local symbol table for a while, since we use it
     # to calculate argument defaults below.

From d736b9ff5a2194856453672bd3dbd19593502e11 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 15:45:02 +0100
Subject: [PATCH 07/17] Refactor more

---
 mypyc/irbuild/function.py | 49 +++++++++++++++++++++------------------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index 717f1d1da2ac..af2db32c78ad 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -267,35 +267,13 @@ def c() -> None:
             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, 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()
-        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)
+        gen_generator_body_func(builder, fn_info, sig)
     else:
         load_env_registers(builder)
         gen_arg_defaults(builder)
         if contains_nested:
             finalize_env_class(builder)
-
         add_vars_to_env(builder)
-
         builder.accept(fitem.body)
         builder.maybe_add_implicit_return()
 
@@ -326,6 +304,31 @@ def c() -> None:
     return func_ir, func_reg
 
 
+def gen_generator_body_func(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature) -> None:
+    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)
+
+
 def add_vars_to_env(builder: IRBuilder) -> None:
     # Add all variables and functions that are declared/defined within this
     # function and are referenced in functions nested within this one to this

From a65ca35cdd0d2b8b06b7e71d25b7e2cfab81ea33 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 15:52:15 +0100
Subject: [PATCH 08/17] Refactor

---
 mypyc/irbuild/env_class.py | 33 ++++++++++++++++++
 mypyc/irbuild/function.py  | 68 ++------------------------------------
 mypyc/irbuild/generator.py | 36 ++++++++++++++++++--
 3 files changed, 70 insertions(+), 67 deletions(-)

diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py
index aa223fe20176..8378b35adaa0 100644
--- a/mypyc/irbuild/env_class.py
+++ b/mypyc/irbuild/env_class.py
@@ -191,6 +191,39 @@ 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 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
+                )
+
+
 def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: ImplicitClass) -> None:
     """Enable calling a nested function (with a callable class) recursively.
 
diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index af2db32c78ad..3f94cd087641 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -73,17 +73,13 @@
 from mypyc.irbuild.env_class import (
     finalize_env_class,
     load_env_registers,
-    load_outer_envs,
     setup_env_class,
-    setup_func_for_recursive_call,
+    add_vars_to_env,
 )
 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,
+    gen_generator_body_func,
 )
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.irbuild.util import is_constant
@@ -261,7 +257,7 @@ def c() -> None:
         setup_callable_class(builder)
 
     if is_generator:
-        # Do a first-pass and generate a function that just returns a generator object.
+        # 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))
@@ -304,64 +300,6 @@ def c() -> None:
     return func_ir, func_reg
 
 
-def gen_generator_body_func(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature) -> None:
-    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)
-
-
-def add_vars_to_env(builder: IRBuilder) -> None:
-    # 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
-                )
-
-
 def has_nested_func_self_reference(builder: IRBuilder, fitem: FuncItem) -> bool:
     """Does a nested function contain a self-reference in its body?
 
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index a950a9fdf1e1..a29f6f15f4f2 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -12,7 +12,7 @@
 
 from typing import Callable
 
-from mypy.nodes import ARG_OPT, Var
+from mypy.nodes import ARG_OPT, Var, FuncDef
 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
@@ -37,9 +37,12 @@
 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 (
@@ -55,6 +58,7 @@ 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)
@@ -63,10 +67,38 @@ def gen_generator_func(
 
     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_body_func(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature) -> None:
+    """Generate IR for the body of the generator function.
+
+    This will be used by the generated "__next__".
+    """
+    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)
+
+
 def instantiate_generator_class(builder: IRBuilder) -> Value:
     fitem = builder.fn_info.fitem
     generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line))

From 52c8989ae414e0a6bb0f3447227df8f8e9515826 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 16:24:52 +0100
Subject: [PATCH 09/17] Lint

---
 mypyc/irbuild/function.py  | 11 +++++++----
 mypyc/irbuild/generator.py |  6 ++++--
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index 3f94cd087641..ca7be123568c 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -69,17 +69,17 @@
     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,
     setup_env_class,
-    add_vars_to_env,
 )
 from mypyc.irbuild.generator import (
     add_methods_to_generator_class,
-    gen_generator_func,
     gen_generator_body_func,
+    gen_generator_func,
 )
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.irbuild.util import is_constant
@@ -260,7 +260,10 @@ def c() -> None:
         # 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))
+            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.
         gen_generator_body_func(builder, fn_info, sig)
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index a29f6f15f4f2..d3c5104fc71d 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -12,7 +12,7 @@
 
 from typing import Callable
 
-from mypy.nodes import ARG_OPT, Var, FuncDef
+from mypy.nodes import ARG_OPT, FuncDef, 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
@@ -56,7 +56,9 @@
 
 def gen_generator_func(
     builder: IRBuilder,
-    gen_func_ir: Callable[[list[Register], list[BasicBlock], FuncInfo], tuple[FuncIR, Value | None]],
+    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)

From 368db2dc000a507009ba869ebd6902a02c502160 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 16:35:09 +0100
Subject: [PATCH 10/17] Refactor

---
 mypyc/irbuild/function.py  | 22 +++++++++++++---------
 mypyc/irbuild/generator.py |  2 +-
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index ca7be123568c..385767cffd40 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -78,7 +78,7 @@
 )
 from mypyc.irbuild.generator import (
     add_methods_to_generator_class,
-    gen_generator_body_func,
+    gen_generator_func_body,
     gen_generator_func,
 )
 from mypyc.irbuild.targets import AssignmentTarget
@@ -266,15 +266,9 @@ def c() -> None:
         )
 
         # Re-enter the FuncItem and visit the body of the function this time.
-        gen_generator_body_func(builder, fn_info, sig)
+        gen_generator_func_body(builder, fn_info, sig)
     else:
-        load_env_registers(builder)
-        gen_arg_defaults(builder)
-        if contains_nested:
-            finalize_env_class(builder)
-        add_vars_to_env(builder)
-        builder.accept(fitem.body)
-        builder.maybe_add_implicit_return()
+        gen_func_body(builder)
 
     # Hang on to the local symbol table for a while, since we use it
     # to calculate argument defaults below.
@@ -303,6 +297,16 @@ def c() -> None:
     return func_ir, func_reg
 
 
+def gen_func_body(builder: IRBuilder) -> 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()
+
+
 def has_nested_func_self_reference(builder: IRBuilder, fitem: FuncItem) -> bool:
     """Does a nested function contain a self-reference in its body?
 
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index d3c5104fc71d..16e8f949dd09 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -72,7 +72,7 @@ def gen_generator_func(
     return func_ir, func_reg
 
 
-def gen_generator_body_func(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature) -> None:
+def gen_generator_func_body(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature) -> None:
     """Generate IR for the body of the generator function.
 
     This will be used by the generated "__next__".

From 8f3d0bff071ed275778641a7b2f08c834a202c76 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 16:37:36 +0100
Subject: [PATCH 11/17] Refactor (partial)

---
 mypyc/irbuild/function.py | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index 385767cffd40..334a7884fca3 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -267,25 +267,34 @@ def c() -> None:
 
         # Re-enter the FuncItem and visit the body of the function this time.
         gen_generator_func_body(builder, fn_info, sig)
-    else:
-        gen_func_body(builder)
 
-    # Hang on to the local symbol table for a while, since we use it
-    # to calculate argument defaults below.
-    symtable = builder.symtables[-1]
+        # 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()
+        args, _, blocks, ret_type, fn_info = builder.leave()
 
-    if is_generator:
         add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
+
+        # 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:
+        gen_func_body(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()
+
         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)
+        # 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

From 300f2aa95271c1f0f0a4aa5846fe9838963e2144 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 16:42:16 +0100
Subject: [PATCH 12/17] Refactor more

---
 mypyc/irbuild/function.py | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index 334a7884fca3..c17233f5cd21 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -78,8 +78,8 @@
 )
 from mypyc.irbuild.generator import (
     add_methods_to_generator_class,
-    gen_generator_func_body,
     gen_generator_func,
+    gen_generator_func_body,
 )
 from mypyc.irbuild.targets import AssignmentTarget
 from mypyc.irbuild.util import is_constant
@@ -280,21 +280,7 @@ def c() -> None:
         # calculate them *once* when the function definition is evaluated.
         calculate_arg_defaults(builder, fn_info, func_reg, symtable)
     else:
-        gen_func_body(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()
-
-        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)
+        func_ir, func_reg = gen_func_body(builder, sig, cdef, is_singledispatch)
 
     if is_singledispatch:
         # add the generated main singledispatch function
@@ -306,7 +292,9 @@ def c() -> None:
     return func_ir, func_reg
 
 
-def gen_func_body(builder: IRBuilder) -> None:
+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:
@@ -315,6 +303,19 @@ def gen_func_body(builder: IRBuilder) -> None:
     builder.accept(builder.fn_info.fitem.body)
     builder.maybe_add_implicit_return()
 
+    # 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()
+
+    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)
+    return func_ir, func_reg
+
 
 def has_nested_func_self_reference(builder: IRBuilder, fitem: FuncItem) -> bool:
     """Does a nested function contain a self-reference in its body?

From 3f1bcc8e4d0f788a993294bacf2fb21cfd420e95 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 16:47:56 +0100
Subject: [PATCH 13/17] Refactor

---
 mypyc/irbuild/function.py  | 16 ++--------------
 mypyc/irbuild/generator.py | 17 ++++++++++++++---
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index c17233f5cd21..a5af4bcdf8b0 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -76,11 +76,7 @@
     load_env_registers,
     setup_env_class,
 )
-from mypyc.irbuild.generator import (
-    add_methods_to_generator_class,
-    gen_generator_func,
-    gen_generator_func_body,
-)
+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
@@ -266,15 +262,7 @@ def c() -> None:
         )
 
         # Re-enter the FuncItem and visit the body of the function this time.
-        gen_generator_func_body(builder, fn_info, sig)
-
-        # 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()
-
-        add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
+        symtable = gen_generator_func_body(builder, fn_info, sig)
 
         # Evaluate argument defaults in the surrounding scope, since we
         # calculate them *once* when the function definition is evaluated.
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index 16e8f949dd09..f0370ff06b3f 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -12,7 +12,7 @@
 
 from typing import Callable
 
-from mypy.nodes import ARG_OPT, FuncDef, Var
+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
@@ -33,7 +33,7 @@
     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,
@@ -72,7 +72,9 @@ def gen_generator_func(
     return func_ir, func_reg
 
 
-def gen_generator_func_body(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature) -> None:
+def gen_generator_func_body(
+    builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature
+) -> dict[SymbolNode, SymbolTarget]:
     """Generate IR for the body of the generator function.
 
     This will be used by the generated "__next__".
@@ -100,6 +102,15 @@ def gen_generator_func_body(builder: IRBuilder, fn_info: FuncInfo, sig: FuncSign
 
     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()
+
+    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

From 4e60d48c9b0692642c3d28a85cd82439b79859b7 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 30 Apr 2025 16:57:20 +0100
Subject: [PATCH 14/17] Update comment and docstring

---
 mypyc/irbuild/generator.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index f0370ff06b3f..2779d3d92a72 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -75,9 +75,12 @@ def gen_generator_func(
 def gen_generator_func_body(
     builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature
 ) -> dict[SymbolNode, SymbolTarget]:
-    """Generate IR for the body of the generator function.
+    """Generate IR based on the body of a generator function.
 
-    This will be used by the generated "__next__".
+    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)
@@ -102,13 +105,14 @@ def gen_generator_func_body(
 
     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.
+    # 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
 
 

From 3091e55229377ddd54eab1df0d91057eaa1693ba Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Thu, 1 May 2025 15:35:33 +0100
Subject: [PATCH 15/17] Minor return type refactor

---
 mypyc/irbuild/builder.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
index edb873237bf2..7fee21a1c6df 100644
--- a/mypyc/irbuild/builder.py
+++ b/mypyc/irbuild/builder.py
@@ -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)

From 0c284537272bf2e9daf52c42f0cfa903fa8f202a Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Thu, 1 May 2025 15:43:40 +0100
Subject: [PATCH 16/17] Change comment to docstring

---
 mypyc/irbuild/env_class.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py
index 8378b35adaa0..ab786fe71dda 100644
--- a/mypyc/irbuild/env_class.py
+++ b/mypyc/irbuild/env_class.py
@@ -192,12 +192,14 @@ def add_args_to_env(
 
 
 def add_vars_to_env(builder: IRBuilder) -> None:
-    # 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.
+    """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

From 2c495c827cf13631828d32774e7de21ffdedaad4 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Fri, 2 May 2025 11:37:25 +0100
Subject: [PATCH 17/17] Add review

---
 mypyc/irbuild/builder.py   | 27 +++++++++++++++++++++++++++
 mypyc/irbuild/function.py  | 38 ++------------------------------------
 mypyc/irbuild/generator.py | 12 +++++++-----
 3 files changed, 36 insertions(+), 41 deletions(-)

diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
index 7fee21a1c6df..72a5ff4099df 100644
--- a/mypyc/irbuild/builder.py
+++ b/mypyc/irbuild/builder.py
@@ -1497,3 +1497,30 @@ def create_type_params(
         builder.init_type_var(tv, type_param.name, line)
         tvs.append(tv)
     return tvs
+
+
+def calculate_arg_defaults(
+    builder: IRBuilder,
+    fn_info: FuncInfo,
+    func_reg: Value | None,
+    symtable: dict[SymbolNode, SymbolTarget],
+) -> None:
+    """Calculate default argument values and store them.
+
+    They are stored in statics for top level functions and in
+    the function objects for nested functions (while constants are
+    still stored computed on demand).
+    """
+    fitem = fn_info.fitem
+    for arg in fitem.arguments:
+        # Constant values don't get stored but just recomputed
+        if arg.initializer and not is_constant(arg.initializer):
+            value = builder.coerce(
+                builder.accept(arg.initializer), symtable[arg.variable].type, arg.line
+            )
+            if not fn_info.is_nested:
+                name = fitem.fullname + "." + arg.variable.name
+                builder.add(InitStatic(value, name, builder.module_name))
+            else:
+                assert func_reg is not None
+                builder.add(SetAttr(func_reg, arg.variable.name, value, arg.line))
diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py
index a5af4bcdf8b0..cb9a1a3dc4a3 100644
--- a/mypyc/irbuild/function.py
+++ b/mypyc/irbuild/function.py
@@ -25,7 +25,6 @@
     FuncItem,
     LambdaExpr,
     OverloadedFuncDef,
-    SymbolNode,
     TypeInfo,
     Var,
 )
@@ -44,7 +43,6 @@
 from mypyc.ir.ops import (
     BasicBlock,
     GetAttr,
-    InitStatic,
     Integer,
     LoadAddress,
     LoadLiteral,
@@ -62,7 +60,7 @@
     int_rprimitive,
     object_rprimitive,
 )
-from mypyc.irbuild.builder import IRBuilder, SymbolTarget, gen_arg_defaults
+from mypyc.irbuild.builder import IRBuilder, calculate_arg_defaults, gen_arg_defaults
 from mypyc.irbuild.callable_class import (
     add_call_to_callable_class,
     add_get_to_callable_class,
@@ -78,7 +76,6 @@
 )
 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
 from mypyc.primitives.generic_ops import py_setattr_op
 from mypyc.primitives.misc_ops import register_function
@@ -262,11 +259,7 @@ def c() -> None:
         )
 
         # Re-enter the FuncItem and visit the body of the function this time.
-        symtable = gen_generator_func_body(builder, fn_info, sig)
-
-        # 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)
+        gen_generator_func_body(builder, fn_info, sig, func_reg)
     else:
         func_ir, func_reg = gen_func_body(builder, sig, cdef, is_singledispatch)
 
@@ -462,33 +455,6 @@ def handle_non_ext_method(
     builder.add_to_non_ext_dict(non_ext, name, func_reg, fdef.line)
 
 
-def calculate_arg_defaults(
-    builder: IRBuilder,
-    fn_info: FuncInfo,
-    func_reg: Value | None,
-    symtable: dict[SymbolNode, SymbolTarget],
-) -> None:
-    """Calculate default argument values and store them.
-
-    They are stored in statics for top level functions and in
-    the function objects for nested functions (while constants are
-    still stored computed on demand).
-    """
-    fitem = fn_info.fitem
-    for arg in fitem.arguments:
-        # Constant values don't get stored but just recomputed
-        if arg.initializer and not is_constant(arg.initializer):
-            value = builder.coerce(
-                builder.accept(arg.initializer), symtable[arg.variable].type, arg.line
-            )
-            if not fn_info.is_nested:
-                name = fitem.fullname + "." + arg.variable.name
-                builder.add(InitStatic(value, name, builder.module_name))
-            else:
-                assert func_reg is not None
-                builder.add(SetAttr(func_reg, arg.variable.name, value, arg.line))
-
-
 def gen_func_ns(builder: IRBuilder) -> str:
     """Generate a namespace for a nested function using its outer function names."""
     return "_".join(
diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index 2779d3d92a72..74c8d27a6324 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -12,7 +12,7 @@
 
 from typing import Callable
 
-from mypy.nodes import ARG_OPT, FuncDef, SymbolNode, Var
+from mypy.nodes import ARG_OPT, FuncDef, 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
@@ -33,7 +33,7 @@
     Value,
 )
 from mypyc.ir.rtypes import RInstance, int_rprimitive, object_rprimitive
-from mypyc.irbuild.builder import IRBuilder, SymbolTarget, gen_arg_defaults
+from mypyc.irbuild.builder import IRBuilder, calculate_arg_defaults, gen_arg_defaults
 from mypyc.irbuild.context import FuncInfo, GeneratorClass
 from mypyc.irbuild.env_class import (
     add_args_to_env,
@@ -73,8 +73,8 @@ def gen_generator_func(
 
 
 def gen_generator_func_body(
-    builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature
-) -> dict[SymbolNode, SymbolTarget]:
+    builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature, func_reg: Value | None
+) -> None:
     """Generate IR based on the body of a generator function.
 
     Add "__next__", "__iter__" and other generator methods to the generator
@@ -113,7 +113,9 @@ class that implements the function (each function gets a separate class).
 
     add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
 
-    return symtable
+    # 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)
 
 
 def instantiate_generator_class(builder: IRBuilder) -> Value: