[mypyc] Use vectorcall for calls to the type object of a native class

mypyc/codegen/emitclass.py

@@ -17,6 +17,7 @@
     generate_am_send_wrapper,
     generate_bin_op_wrapper,
     generate_bool_wrapper,
+    generate_constructor_wrapper_function,
     generate_contains_wrapper,
     generate_dunder_wrapper,
     generate_get_wrapper,
@@ -252,10 +253,40 @@ def generate_class_reuse(
     context.declarations[name] = HeaderDeclaration(f"CPyThreadLocal {struct_name} *{name};")
 
 
-def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
+def can_use_constructor_for_type_call(cl: ClassIR) -> bool:
+    """Can a call to the type object of cl go straight to the native constructor?
+
+    If so, the type object gets a tp_vectorcall that parses the arguments and calls
+    the native constructor. Otherwise a call goes through tp_new and tp_init.
+    """
+    init = cl.get_method_and_class("__init__")
+    if init is None or "__init__" not in init[1].method_decls:
+        # Without __init__ there are no arguments to parse, and tp_new accepts and
+        # ignores any that are passed. A decorated __init__ isn't called natively.
+        return False
+    if cl.has_method("__new__") or cl.has_non_native_new():
+        # If __new__ returns an instance of a subclass, type.__call__ calls __init__ of
+        # the subclass, but the native constructor doesn't. With a decorated __new__,
+        # the native constructor calls the type object.
+        return False
+    if cl.is_abstract:
+        # The only thing that stops an abstract class from being instantiated is the
+        # check of the self argument in the wrapper of __init__.
+        return False
+    # If __init__ isn't native, the constructor takes *args and **kwargs and calls
+    # tp_init. Building a tuple and a dict for it wouldn't be any faster.
+    return all(arg.kind not in (ARG_STAR, ARG_STAR2) for arg in cl.ctor.sig.args)
+
+
+def generate_class(
+    cl: ClassIR, module: str, emitter: Emitter, source_paths: dict[str, str]
+) -> None:
     """Generate C code for a class.
 
     This is the main entry point to the module.
+
+    Args:
+        source_paths: Map from module names to source file paths
     """
     name = cl.name
     name_prefix = cl.name_prefix(emitter.names)
@@ -415,6 +446,15 @@ def emit_line() -> None:
         generate_new_for_class(cl, new_name, vtable_name, setup_name, init_fn, emitter)
         emit_line()
 
+        if init_fn is not None and can_use_constructor_for_type_call(cl):
+            vectorcall_name = f"{name_prefix}_vectorcall"
+            init_module = init_fn.decl.module_name
+            generate_constructor_wrapper_function(
+                cl, init_fn, vectorcall_name, emitter, source_paths[init_module], init_module
+            )
+            fields["tp_vectorcall"] = vectorcall_name
+            emit_line()
+
         if cl.allow_interpreted_subclasses:
             shadow_vtable_name: str | None = generate_vtables(
                 cl, vtable_setup_name + "_shadow", vtable_name + "_shadow", emitter, shadow=True
@@ -450,7 +490,7 @@ def emit_line() -> None:
             cl.struct_name(emitter.names)
         )
         flags.append("_Py_TPFLAGS_HAVE_VECTORCALL")
-        if not fields.get("tp_vectorcall"):
+        if not fields.get("tp_call"):
             # This is just a placeholder to please CPython. It will be
             # overridden during setup.
             fields["tp_call"] = "PyVectorcall_Call"

mypyc/codegen/emitmodule.py

@@ -755,7 +755,7 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
 
             for cl in module.classes:
                 if cl.is_ext_class:
-                    generate_class(cl, module_name, emitter)
+                    generate_class(cl, module_name, emitter, self.source_paths)
 
             running_flag_classes = {cl.name: cl for cl in module.classes if cl.has_running_flag}
 
@@ -1731,7 +1731,9 @@ def is_fastcall_supported(fn: FuncIR, capi_version: tuple[int, int]) -> bool:
         if fn.name == "__call__":
             # We can use vectorcalls (PEP 590) when supported
             return True
-        # TODO: Support fastcall for __init__ and __new__.
+        # The wrappers of __init__ and __new__ are called from tp_init and tp_new,
+        # which get a tuple and a dict. A call to the type object can skip them, see
+        # can_use_constructor_for_type_call().
         return fn.name != "__init__" and fn.name != "__new__"
     return True
 

mypyc/codegen/emitwrapper.py

@@ -26,7 +26,7 @@
     bitmap_name,
 )
 from mypyc.ir.class_ir import ClassIR
-from mypyc.ir.func_ir import FUNC_STATICMETHOD, FuncIR, RuntimeArg
+from mypyc.ir.func_ir import FUNC_STATICMETHOD, FuncDecl, FuncIR, RuntimeArg
 from mypyc.ir.rtypes import (
     RInstance,
     RType,
@@ -131,6 +131,23 @@ def make_format_string(func_name: str | None, groups: dict[ArgKind, list[Runtime
     return format
 
 
+def stack_parse_function(
+    real_args: list[RuntimeArg], groups: dict[ArgKind, list[RuntimeArg]]
+) -> str:
+    """Return the CPyArg_ParseStackAndKeywords* variant to use for the given arguments."""
+    # Special case some common signatures
+    if not real_args:
+        # No args
+        return "CPyArg_ParseStackAndKeywordsNoArgs"
+    elif len(real_args) == 1 and len(groups[ARG_POS]) == 1:
+        # Single positional arg
+        return "CPyArg_ParseStackAndKeywordsOneArg"
+    elif len(real_args) == len(groups[ARG_POS]) + len(groups[ARG_OPT]):
+        # No keyword-only args, *args or **kwargs
+        return "CPyArg_ParseStackAndKeywordsSimple"
+    return "CPyArg_ParseStackAndKeywords"
+
+
 def generate_wrapper_function(
     fn: FuncIR, emitter: Emitter, source_path: str, module_name: str
 ) -> None:
@@ -177,20 +194,9 @@ def generate_wrapper_function(
         nargs = "PyVectorcall_NARGS(nargs)"
     else:
         nargs = "nargs"
-    parse_fn = "CPyArg_ParseStackAndKeywords"
-    # Special case some common signatures
-    if not real_args:
-        # No args
-        parse_fn = "CPyArg_ParseStackAndKeywordsNoArgs"
-    elif len(real_args) == 1 and len(groups[ARG_POS]) == 1:
-        # Single positional arg
-        parse_fn = "CPyArg_ParseStackAndKeywordsOneArg"
-    elif len(real_args) == len(groups[ARG_POS]) + len(groups[ARG_OPT]):
-        # No keyword-only args, *args or **kwargs
-        parse_fn = "CPyArg_ParseStackAndKeywordsSimple"
     emitter.emit_lines(
         "if (!{}(args, {}, kwnames, &parser{})) {{".format(
-            parse_fn, nargs, "".join(", " + n for n in arg_ptrs)
+            stack_parse_function(real_args, groups), nargs, "".join(", " + n for n in arg_ptrs)
         ),
         "return NULL;",
         "}",
@@ -210,6 +216,64 @@ def generate_wrapper_function(
     emitter.emit_line("}")
 
 
+def generate_constructor_wrapper_function(
+    cl: ClassIR,
+    init_fn: FuncIR,
+    func_name: str,
+    emitter: Emitter,
+    source_path: str,
+    module_name: str,
+) -> None:
+    """Generate a vectorcall wrapper for the native constructor of a class.
+
+    This is used as tp_vectorcall of the type object, so that calling the type object
+    doesn't go through tp_new and tp_init, which get the arguments as a tuple and a
+    dict and pass them on to the legacy wrapper of __init__.
+
+    The constructor takes the arguments of init_fn, and argument errors are reported
+    the same way as in the wrapper of init_fn. The source path and the module name are
+    those of init_fn, which cl may inherit from a class in another module.
+    """
+    ctor = cl.ctor
+    emitter.emit_line("static PyObject *")
+    emitter.emit_line(
+        f"{func_name}(PyObject *type, PyObject *const *args, size_t nargs, PyObject *kwnames)"
+    )
+    emitter.emit_line("{")
+
+    real_args = list(ctor.sig.args)
+    if ctor.sig.num_bitmap_args:
+        real_args = real_args[: -ctor.sig.num_bitmap_args]
+    groups = make_arg_groups(real_args)
+    assert not groups[ARG_STAR] and not groups[ARG_STAR2], ctor.sig
+    reordered_args = reorder_arg_groups(groups)
+
+    emitter.emit_line(make_static_kwlist(reordered_args))
+    fmt = make_format_string(init_fn.name, groups)
+    emitter.emit_line(f'static CPyArg_Parser parser = {{"{fmt}", kwlist, 0}};')
+    for arg in real_args:
+        emitter.emit_line(
+            "PyObject *obj_{}{};".format(arg.name, " = NULL" if arg.optional else "")
+        )
+    emitter.emit_lines(
+        "if (!{}(args, PyVectorcall_NARGS(nargs), kwnames, &parser{})) {{".format(
+            stack_parse_function(real_args, groups),
+            "".join(f", &obj_{arg.name}" for arg in reordered_args),
+        ),
+        "return NULL;",
+        "}",
+    )
+    for i in range(ctor.sig.num_bitmap_args):
+        name = bitmap_name(i)
+        emitter.emit_line(f"{BITMAP_TYPE} {name} = 0;")
+    traceback_code = generate_traceback_code(init_fn, emitter, source_path, module_name)
+    generate_wrapper_core(
+        ctor, emitter, groups[ARG_OPT] + groups[ARG_NAMED_OPT], traceback_code=traceback_code
+    )
+
+    emitter.emit_line("}")
+
+
 # Legacy generic wrapper functions
 #
 # These take a self object, a Python tuple of positional arguments,
@@ -785,7 +849,7 @@ def generate_contains_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 
 
 def generate_wrapper_core(
-    fn: FuncIR,
+    fn: FuncIR | FuncDecl,
     emitter: Emitter,
     optional_args: list[RuntimeArg] | None = None,
     arg_names: list[str] | None = None,
@@ -892,24 +956,27 @@ def __init__(self, cl: ClassIR | None, emitter: Emitter) -> None:
         self.optional_args: list[RuntimeArg] = []
         self.traceback_code = ""
 
-    def set_target(self, fn: FuncIR) -> None:
+    def set_target(self, fn: FuncIR | FuncDecl) -> None:
         """Set the wrapped function.
 
         It's fine to modify the attributes initialized here later to customize
         the wrapper function.
+
+        The native constructor of a class only has a declaration.
         """
-        self.target_name = fn.name
-        self.target_cname = fn.cname(self.emitter.names)
+        decl = fn.decl if isinstance(fn, FuncIR) else fn
+        self.target_name = decl.name
+        self.target_cname = decl.cname(self.emitter.names)
         # Cached native-call expression so cross-group targets go through the
         # exports table; same as `NATIVE_PREFIX + cname` for in-group calls.
-        self.target_native_call = self.emitter.native_function_call(fn.decl)
-        self.num_bitmap_args = fn.sig.num_bitmap_args
+        self.target_native_call = self.emitter.native_function_call(decl)
+        self.num_bitmap_args = decl.sig.num_bitmap_args
         if self.num_bitmap_args:
-            self.args = fn.args[: -self.num_bitmap_args]
+            self.args = decl.sig.args[: -self.num_bitmap_args]
         else:
-            self.args = fn.args
+            self.args = decl.sig.args
         self.arg_names = [arg.name for arg in self.args]
-        self.ret_type = fn.ret_type
+        self.ret_type = decl.sig.ret_type
 
     def wrapper_name(self) -> str:
         """Return the name of the wrapper function."""