[mypyc] Call generator helper method directly in await expression (#19376)

Previously calls like `await foo()` were compiled to code that included
code like this (in Python-like pseudocode):
```
a = foo()
...
b = get_coro(a)
...
c = next(b)
```
In the above code, `get_coro(a)` just returns `a` if `foo` is a native
async function, so we now optimize this call away. Also `next(b)` calls
`b.__next__()`, which calls the generated generator helper method
`__mypyc_generator_helper__`. Now we call the helper method directly,
which saves some unnecessary calls.

More importantly, in a follow-up PR I can easily change the way
`__mypyc_generator_helper__` is called, since we now call it directly.
This makes it possible to avoid raising a `StopIteration` exception in
many await expressions. The goal of this PR is to prepare for the latter
optimization. This PR doesn't help performance significantly by itself.

In order to call the helper method directly, I had to generate the
declaration of this method and the generated generator class before the
main irbuild pass, since otherwise a call site could be processed before
we have processed the called generator.

I also improved test coverage of related functionality. We don't have an
IR test for async calls, since the IR is very verbose. I manually
inspected the generated IR to verify that the new code path works both
when calling a top-level function and when calling a method. I'll later
add a mypyc benchmark to ensure that we will notice if the performance
of async calls is degraded.

Author: JukkaL

mypyc/irbuild/function.py

@@ -261,7 +261,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, func_reg)
+        gen_generator_func_body(builder, fn_info, func_reg)
     else:
         func_ir, func_reg = gen_func_body(builder, sig, cdef, is_singledispatch)
 

mypyc/irbuild/generator.py

@@ -13,9 +13,9 @@
 from typing import Callable
 
 from mypy.nodes import ARG_OPT, FuncDef, Var
-from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME, SELF_NAME
+from mypyc.common import ENV_ATTR_NAME, NEXT_LABEL_ATTR_NAME
 from mypyc.ir.class_ir import ClassIR
-from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
+from mypyc.ir.func_ir import FuncDecl, FuncIR
 from mypyc.ir.ops import (
     NO_TRACEBACK_LINE_NO,
     BasicBlock,
@@ -78,17 +78,15 @@ def gen_generator_func(
     return func_ir, func_reg
 
 
-def gen_generator_func_body(
-    builder: IRBuilder, fn_info: FuncInfo, sig: FuncSignature, func_reg: Value | None
-) -> None:
+def gen_generator_func_body(builder: IRBuilder, fn_info: FuncInfo, 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
     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)
+    builder.enter(fn_info, ret_type=object_rprimitive)
     setup_env_for_generator_class(builder)
 
     load_outer_envs(builder, builder.fn_info.generator_class)
@@ -117,7 +115,7 @@ class that implements the function (each function gets a separate class).
 
     args, _, blocks, ret_type, fn_info = builder.leave()
 
-    add_methods_to_generator_class(builder, fn_info, sig, args, blocks, fitem.is_coroutine)
+    add_methods_to_generator_class(builder, fn_info, args, blocks, fitem.is_coroutine)
 
     # Evaluate argument defaults in the surrounding scope, since we
     # calculate them *once* when the function definition is evaluated.
@@ -153,10 +151,9 @@ def instantiate_generator_class(builder: IRBuilder) -> Value:
 
 
 def setup_generator_class(builder: IRBuilder) -> ClassIR:
-    name = f"{builder.fn_info.namespaced_name()}_gen"
-
-    generator_class_ir = ClassIR(name, builder.module_name, is_generated=True, is_final_class=True)
-    generator_class_ir.reuse_freed_instance = True
+    mapper = builder.mapper
+    assert isinstance(builder.fn_info.fitem, FuncDef)
+    generator_class_ir = mapper.fdef_to_generator[builder.fn_info.fitem]
     if builder.fn_info.can_merge_generator_and_env_classes():
         builder.fn_info.env_class = generator_class_ir
     else:
@@ -216,46 +213,25 @@ def add_raise_exception_blocks_to_generator_class(builder: IRBuilder, line: int)
 def add_methods_to_generator_class(
     builder: IRBuilder,
     fn_info: FuncInfo,
-    sig: FuncSignature,
     arg_regs: list[Register],
     blocks: list[BasicBlock],
     is_coroutine: bool,
 ) -> None:
-    helper_fn_decl = add_helper_to_generator_class(builder, arg_regs, blocks, sig, fn_info)
-    add_next_to_generator_class(builder, fn_info, helper_fn_decl, sig)
-    add_send_to_generator_class(builder, fn_info, helper_fn_decl, sig)
+    helper_fn_decl = add_helper_to_generator_class(builder, arg_regs, blocks, fn_info)
+    add_next_to_generator_class(builder, fn_info, helper_fn_decl)
+    add_send_to_generator_class(builder, fn_info, helper_fn_decl)
     add_iter_to_generator_class(builder, fn_info)
-    add_throw_to_generator_class(builder, fn_info, helper_fn_decl, sig)
+    add_throw_to_generator_class(builder, fn_info, helper_fn_decl)
     add_close_to_generator_class(builder, fn_info)
     if is_coroutine:
         add_await_to_generator_class(builder, fn_info)
 
 
 def add_helper_to_generator_class(
-    builder: IRBuilder,
-    arg_regs: list[Register],
-    blocks: list[BasicBlock],
-    sig: FuncSignature,
-    fn_info: FuncInfo,
+    builder: IRBuilder, arg_regs: list[Register], blocks: list[BasicBlock], fn_info: FuncInfo
 ) -> FuncDecl:
     """Generates a helper method for a generator class, called by '__next__' and 'throw'."""
-    sig = FuncSignature(
-        (
-            RuntimeArg(SELF_NAME, object_rprimitive),
-            RuntimeArg("type", object_rprimitive),
-            RuntimeArg("value", object_rprimitive),
-            RuntimeArg("traceback", object_rprimitive),
-            RuntimeArg("arg", object_rprimitive),
-        ),
-        sig.ret_type,
-    )
-    helper_fn_decl = FuncDecl(
-        "__mypyc_generator_helper__",
-        fn_info.generator_class.ir.name,
-        builder.module_name,
-        sig,
-        internal=True,
-    )
+    helper_fn_decl = fn_info.generator_class.ir.method_decls["__mypyc_generator_helper__"]
     helper_fn_ir = FuncIR(
         helper_fn_decl, arg_regs, blocks, fn_info.fitem.line, traceback_name=fn_info.fitem.name
     )
@@ -272,9 +248,7 @@ def add_iter_to_generator_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
         builder.add(Return(builder.self()))
 
 
-def add_next_to_generator_class(
-    builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature
-) -> None:
+def add_next_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl) -> None:
     """Generates the '__next__' method for a generator class."""
     with builder.enter_method(fn_info.generator_class.ir, "__next__", object_rprimitive, fn_info):
         none_reg = builder.none_object()
@@ -289,9 +263,7 @@ def add_next_to_generator_class(
         builder.add(Return(result))
 
 
-def add_send_to_generator_class(
-    builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature
-) -> None:
+def add_send_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl) -> None:
     """Generates the 'send' method for a generator class."""
     with builder.enter_method(fn_info.generator_class.ir, "send", object_rprimitive, fn_info):
         arg = builder.add_argument("arg", object_rprimitive)
@@ -307,9 +279,7 @@ def add_send_to_generator_class(
         builder.add(Return(result))
 
 
-def add_throw_to_generator_class(
-    builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl, sig: FuncSignature
-) -> None:
+def add_throw_to_generator_class(builder: IRBuilder, fn_info: FuncInfo, fn_decl: FuncDecl) -> None:
     """Generates the 'throw' method for a generator class."""
     with builder.enter_method(fn_info.generator_class.ir, "throw", object_rprimitive, fn_info):
         typ = builder.add_argument("type", object_rprimitive)

mypyc/irbuild/main.py

@@ -25,7 +25,7 @@ def f(x: int) -> int:
 from typing import Any, Callable, TypeVar, cast
 
 from mypy.build import Graph
-from mypy.nodes import ClassDef, Expression, MypyFile
+from mypy.nodes import ClassDef, Expression, FuncDef, MypyFile
 from mypy.state import state
 from mypy.types import Type
 from mypyc.analysis.attrdefined import analyze_always_defined_attrs
@@ -37,7 +37,11 @@ def f(x: int) -> int:
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.mapper import Mapper
 from mypyc.irbuild.prebuildvisitor import PreBuildVisitor
-from mypyc.irbuild.prepare import build_type_map, find_singledispatch_register_impls
+from mypyc.irbuild.prepare import (
+    build_type_map,
+    create_generator_class_if_needed,
+    find_singledispatch_register_impls,
+)
 from mypyc.irbuild.visitor import IRBuilderVisitor
 from mypyc.irbuild.vtable import compute_vtable
 from mypyc.options import CompilerOptions
@@ -76,6 +80,15 @@ def build_ir(
         pbv = PreBuildVisitor(errors, module, singledispatch_info.decorators_to_remove, types)
         module.accept(pbv)
 
+        # Declare generator classes for nested async functions and generators.
+        for fdef in pbv.nested_funcs:
+            if isinstance(fdef, FuncDef):
+                # Make generator class name sufficiently unique.
+                suffix = f"___{fdef.line}"
+                create_generator_class_if_needed(
+                    module.fullname, None, fdef, mapper, name_suffix=suffix
+                )
+
         # Construct and configure builder objects (cyclic runtime dependency).
         visitor = IRBuilderVisitor()
         builder = IRBuilder(

mypyc/irbuild/mapper.py

@@ -64,6 +64,8 @@ def __init__(self, group_map: dict[str, str | None]) -> None:
         self.type_to_ir: dict[TypeInfo, ClassIR] = {}
         self.func_to_decl: dict[SymbolNode, FuncDecl] = {}
         self.symbol_fullnames: set[str] = set()
+        # The corresponding generator class that implements a generator/async function
+        self.fdef_to_generator: dict[FuncDef, ClassIR] = {}
 
     def type_to_rtype(self, typ: Type | None) -> RType:
         if typ is None:
@@ -171,7 +173,14 @@ def fdef_to_sig(self, fdef: FuncDef, strict_dunders_typing: bool) -> FuncSignatu
                 for typ, kind in zip(fdef.type.arg_types, fdef.type.arg_kinds)
             ]
             arg_pos_onlys = [name is None for name in fdef.type.arg_names]
-            ret = self.type_to_rtype(fdef.type.ret_type)
+            # TODO: We could probably support decorators sometimes (static and class method?)
+            if (fdef.is_coroutine or fdef.is_generator) and not fdef.is_decorated:
+                # Give a more precise type for generators, so that we can optimize
+                # code that uses them. They return a generator object, which has a
+                # specific class. Without this, the type would have to be 'object'.
+                ret: RType = RInstance(self.fdef_to_generator[fdef])
+            else:
+                ret = self.type_to_rtype(fdef.type.ret_type)
         else:
             # Handle unannotated functions
             arg_types = [object_rprimitive for _ in fdef.arguments]

mypyc/irbuild/prepare.py

@@ -38,7 +38,7 @@
 from mypy.semanal import refers_to_fullname
 from mypy.traverser import TraverserVisitor
 from mypy.types import Instance, Type, get_proper_type
-from mypyc.common import PROPSET_PREFIX, get_id_from_name
+from mypyc.common import PROPSET_PREFIX, SELF_NAME, get_id_from_name
 from mypyc.crash import catch_errors
 from mypyc.errors import Errors
 from mypyc.ir.class_ir import ClassIR
@@ -51,7 +51,14 @@
     RuntimeArg,
 )
 from mypyc.ir.ops import DeserMaps
-from mypyc.ir.rtypes import RInstance, RType, dict_rprimitive, none_rprimitive, tuple_rprimitive
+from mypyc.ir.rtypes import (
+    RInstance,
+    RType,
+    dict_rprimitive,
+    none_rprimitive,
+    object_rprimitive,
+    tuple_rprimitive,
+)
 from mypyc.irbuild.mapper import Mapper
 from mypyc.irbuild.util import (
     get_func_def,
@@ -115,7 +122,7 @@ def build_type_map(
 
     # Collect all the functions also. We collect from the symbol table
     # so that we can easily pick out the right copy of a function that
-    # is conditionally defined.
+    # is conditionally defined. This doesn't include nested functions!
     for module in modules:
         for func in get_module_func_defs(module):
             prepare_func_def(module.fullname, None, func, mapper, options)
@@ -179,6 +186,8 @@ def prepare_func_def(
     mapper: Mapper,
     options: CompilerOptions,
 ) -> FuncDecl:
+    create_generator_class_if_needed(module_name, class_name, fdef, mapper)
+
     kind = (
         FUNC_STATICMETHOD
         if fdef.is_static
@@ -190,6 +199,38 @@ def prepare_func_def(
     return decl
 
 
+def create_generator_class_if_needed(
+    module_name: str, class_name: str | None, fdef: FuncDef, mapper: Mapper, name_suffix: str = ""
+) -> None:
+    """If function is a generator/async function, declare a generator class.
+
+    Each generator and async function gets a dedicated class that implements the
+    generator protocol with generated methods.
+    """
+    if fdef.is_coroutine or fdef.is_generator:
+        name = "_".join(x for x in [fdef.name, class_name] if x) + "_gen" + name_suffix
+        cir = ClassIR(name, module_name, is_generated=True, is_final_class=True)
+        cir.reuse_freed_instance = True
+        mapper.fdef_to_generator[fdef] = cir
+
+        helper_sig = FuncSignature(
+            (
+                RuntimeArg(SELF_NAME, object_rprimitive),
+                RuntimeArg("type", object_rprimitive),
+                RuntimeArg("value", object_rprimitive),
+                RuntimeArg("traceback", object_rprimitive),
+                RuntimeArg("arg", object_rprimitive),
+            ),
+            object_rprimitive,
+        )
+
+        # The implementation of most generator functionality is behind this magic method.
+        helper_fn_decl = FuncDecl(
+            "__mypyc_generator_helper__", name, module_name, helper_sig, internal=True
+        )
+        cir.method_decls[helper_fn_decl.name] = helper_fn_decl
+
+
 def prepare_method_def(
     ir: ClassIR,
     module_name: str,

mypyc/irbuild/statement.py

@@ -48,11 +48,13 @@
 )
 from mypyc.common import TEMP_ATTR_NAME
 from mypyc.ir.ops import (
+    ERR_NEVER,
     NAMESPACE_MODULE,
     NO_TRACEBACK_LINE_NO,
     Assign,
     BasicBlock,
     Branch,
+    Call,
     InitStatic,
     Integer,
     LoadAddress,
@@ -930,16 +932,41 @@ def emit_yield_from_or_await(
     to_yield_reg = Register(object_rprimitive)
     received_reg = Register(object_rprimitive)
 
-    get_op = coro_op if is_await else iter_op
-    if isinstance(get_op, PrimitiveDescription):
-        iter_val = builder.primitive_op(get_op, [val], line)
+    helper_method = "__mypyc_generator_helper__"
+    if (
+        isinstance(val, (Call, MethodCall))
+        and isinstance(val.type, RInstance)
+        and val.type.class_ir.has_method(helper_method)
+    ):
+        # This is a generated native generator class, and we can use a fast path.
+        # This allows two optimizations:
+        # 1) No need to call CPy_GetCoro() or iter() since for native generators
+        #    it just returns the generator object (implemented here).
+        # 2) Instead of calling next(), call generator helper method directly,
+        #    since next() just calls __next__ which calls the helper method.
+        iter_val: Value = val
     else:
-        iter_val = builder.call_c(get_op, [val], line)
+        get_op = coro_op if is_await else iter_op
+        if isinstance(get_op, PrimitiveDescription):
+            iter_val = builder.primitive_op(get_op, [val], line)
+        else:
+            iter_val = builder.call_c(get_op, [val], line)
 
     iter_reg = builder.maybe_spill_assignable(iter_val)
 
     stop_block, main_block, done_block = BasicBlock(), BasicBlock(), BasicBlock()
-    _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], line)
+
+    if isinstance(iter_reg.type, RInstance) and iter_reg.type.class_ir.has_method(helper_method):
+        # Second fast path optimization: call helper directly (see also comment above).
+        obj = builder.read(iter_reg)
+        nn = builder.none_object()
+        m = MethodCall(obj, helper_method, [nn, nn, nn, nn], line)
+        # Generators have custom error handling, so disable normal error handling.
+        m.error_kind = ERR_NEVER
+        _y_init = builder.add(m)
+    else:
+        _y_init = builder.call_c(next_raw_op, [builder.read(iter_reg)], line)
+
     builder.add(Branch(_y_init, stop_block, main_block, Branch.IS_ERROR))
 
     # Try extracting a return value from a StopIteration and return it.
@@ -948,7 +975,7 @@ def emit_yield_from_or_await(
     builder.assign(result, builder.call_c(check_stop_op, [], line), line)
     # Clear the spilled iterator/coroutine so that it will be freed.
     # Otherwise, the freeing of the spilled register would likely be delayed.
-    err = builder.add(LoadErrorValue(object_rprimitive))
+    err = builder.add(LoadErrorValue(iter_reg.type))
     builder.assign(iter_reg, err, line)
     builder.goto(done_block)