[mypyc] Fix C name conflict with conditional nested async funcs (#22330)

Fixes mypyc/mypyc#1182.

Multiple conditional definitions of a nested async function or generator
with the same name generated invalid C (redefinition errors):

```python
async def f(flag: bool) -> str:
    if flag:
        async def inner() -> str:
            return "a"
    else:
        async def inner() -> str:
            return "b"
    return await inner()
```

The environment classes of the nested functions had the same name. Now
we make the names unique, similar to the names of callable classes.

Author: JukkaL

mypyc/irbuild/builder.py

@@ -250,7 +250,9 @@ def __init__(
         self.classes: list[ClassIR] = []
         self.final_names: list[tuple[str, RType]] = []
         self.type_var_names: list[str] = []
-        self.callable_class_names: set[str] = set()
+        # Names of generated classes (callable and environment classes), used to
+        # keep the names unique within the module
+        self.generated_class_names: set[str] = set()
         self.options = options
 
         # Keep track of the number of lambdas instantiated so we avoid name conflicts.
@@ -318,6 +320,27 @@ def __init__(
 
     # High-level control
 
+    def unique_class_name(self, base_name: str) -> str:
+        """Return a unique name for a generated class, based on base_name.
+
+        We allow multiple definitions of a nested function with the same name,
+        since this is valid in if-else blocks. Example:
+
+            if True:
+                def foo():          ---->    foo_obj()
+                    return True
+            else:
+                def foo():          ---->    foo_obj_0()
+                    return False
+        """
+        name = base_name
+        count = 0
+        while name in self.generated_class_names:
+            name = f"{base_name}_{count}"
+            count += 1
+        self.generated_class_names.add(name)
+        return name
+
     def set_module(self, module_name: str, module_path: str) -> None:
         """Set the name and path of the current module.
 

mypyc/irbuild/callable_class.py

@@ -54,22 +54,7 @@ def setup_callable_class(builder: IRBuilder) -> None:
     Return a newly constructed ClassIR representing the callable
     class for the nested function.
     """
-    # Check to see that the name has not already been taken. If so,
-    # rename the class. We allow multiple uses of the same function
-    # name because this is valid in if-else blocks. Example:
-    #
-    #     if True:
-    #         def foo():          ---->    foo_obj()
-    #             return True
-    #     else:
-    #         def foo():          ---->    foo_obj_0()
-    #             return False
-    name = base_name = f"{builder.fn_info.namespaced_name()}_obj"
-    count = 0
-    while name in builder.callable_class_names:
-        name = base_name + "_" + str(count)
-        count += 1
-    builder.callable_class_names.add(name)
+    name = builder.unique_class_name(f"{builder.fn_info.namespaced_name()}_obj")
 
     # Define the actual callable class ClassIR, and set its
     # environment to point at the previously defined environment

mypyc/irbuild/env_class.py

@@ -49,12 +49,8 @@ class is generated, the function environment has not yet been
     Return a ClassIR representing an environment for a function
     containing a nested function.
     """
-    env_class = ClassIR(
-        f"{builder.fn_info.namespaced_name()}_env",
-        builder.module_name,
-        is_generated=True,
-        is_final_class=True,
-    )
+    name = builder.unique_class_name(f"{builder.fn_info.namespaced_name()}_env")
+    env_class = ClassIR(name, builder.module_name, is_generated=True, is_final_class=True)
     env_class.reuse_freed_instance = True
     env_class.attributes[SELF_NAME] = RInstance(env_class)
     if builder.fn_info.is_nested and builder.fn_infos[-2]._env_class is not None:

mypyc/test-data/run-async.test

@@ -943,6 +943,36 @@ async def async_def_contains_decorated_async_def(n: int) -> int:
 async def test_async_def_contains_decorated_async_def() -> None:
     assert await async_def_contains_decorated_async_def(7) == 10
 
+async def conditional_nested_async(flag: int) -> str:
+    if flag < 0:
+        async def inner() -> str:
+            return "first"
+    elif flag > 0:
+        async def inner() -> str:
+            return "second"
+    else:
+        async def inner() -> str:
+            return "third"
+    return await inner()
+
+async def test_conditional_nested_async_defs_with_same_name() -> None:
+    assert await conditional_nested_async(-1) == "first"
+    assert await conditional_nested_async(1) == "second"
+    assert await conditional_nested_async(0) == "third"
+
+def conditional_nested_generator(flag: bool) -> list[int]:
+    if flag:
+        def gen() -> Iterator[int]:
+            yield 1
+    else:
+        def gen() -> Iterator[int]:
+            yield 2
+    return list(gen())
+
+def test_conditional_nested_generators_with_same_name() -> None:
+    assert conditional_nested_generator(True) == [1]
+    assert conditional_nested_generator(False) == [2]
+
 [file asyncio/__init__.pyi]
 def run(x: object) -> object: ...
 

mypyc/test-data/run-functions.test

@@ -1426,6 +1426,24 @@ def test_reassigned_nested_function() -> None:
     assert outer(False) == [-1]
     assert outer(True) == [0, 101]
 
+def conditional_outer(flag: bool) -> int:
+    x = 1
+    if flag:
+        def inner() -> int:
+            def g() -> int:
+                return x + 10
+            return g()
+    else:
+        def inner() -> int:
+            def g() -> int:
+                return x + 20
+            return g()
+    return inner()
+
+def test_conditional_nested_functions_with_nested_functions() -> None:
+    assert conditional_outer(True) == 11
+    assert conditional_outer(False) == 21
+
 [case testStarArgFastPaths]
 from typing import Any, Mapping
 def fn(x: str, y: int) -> str: