[mypyc] Use native calls to singledispatch functions (#10981)

mypyc/irbuild/builder.py

@@ -857,6 +857,16 @@ def call_refexpr_with_args(
         callee_node = callee.node
         if isinstance(callee_node, OverloadedFuncDef):
             callee_node = callee_node.impl
+        # TODO: use native calls for any decorated functions which have all their decorators
+        # removed, not just singledispatch functions (which we don't do now just in case those
+        # decorated functions are callable classes or cannot be called without the python API for
+        # some other reason)
+        if (
+            isinstance(callee_node, Decorator)
+            and callee_node.func not in self.fdefs_to_decorators
+            and callee_node.func in self.singledispatch_impls
+        ):
+            callee_node = callee_node.func
         if (callee_node is not None
                 and callee.fullname is not None
                 and callee_node in self.mapper.func_to_decl

mypyc/irbuild/function.py

@@ -947,14 +947,40 @@ def gen_dispatch_func_ir(
 
     generate_singledispatch_dispatch_function(builder, main_func_name, fitem)
     args, _, blocks, _, fn_info = builder.leave()
-    dispatch_func_ir = add_call_to_callable_class(builder, args, blocks, sig, fn_info)
+    dispatch_callable_class = add_call_to_callable_class(builder, args, blocks, sig, fn_info)
+    builder.functions.append(dispatch_callable_class)
     add_get_to_callable_class(builder, fn_info)
     add_register_method_to_callable_class(builder, fn_info)
     func_reg = instantiate_callable_class(builder, fn_info)
+    dispatch_func_ir = generate_dispatch_glue_native_function(
+        builder, fitem, dispatch_callable_class.decl, dispatch_name
+    )
 
     return dispatch_func_ir, func_reg
 
 
+def generate_dispatch_glue_native_function(
+    builder: IRBuilder,
+    fitem: FuncDef,
+    callable_class_decl: FuncDecl,
+    dispatch_name: str,
+) -> FuncIR:
+    line = fitem.line
+    builder.enter()
+    # We store the callable class in the globals dict for this function
+    callable_class = builder.load_global_str(dispatch_name, line)
+    decl = builder.mapper.func_to_decl[fitem]
+    arg_info = get_args(builder, decl.sig.args, line)
+    args = [callable_class] + arg_info.args
+    arg_kinds = [ArgKind.ARG_POS] + arg_info.arg_kinds
+    arg_names = arg_info.arg_names
+    arg_names.insert(0, 'self')
+    ret_val = builder.builder.call(callable_class_decl, args, arg_kinds, arg_names, line)
+    builder.add(Return(ret_val))
+    arg_regs, _, blocks, _, fn_info = builder.leave()
+    return FuncIR(decl, arg_regs, blocks)
+
+
 def generate_singledispatch_callable_class_ctor(builder: IRBuilder) -> None:
     """Create an __init__ that sets registry and dispatch_cache to empty dicts"""
     line = -1

mypyc/test-data/irbuild-singledispatch.test

@@ -28,25 +28,6 @@ L0:
     r4 = PyObject_SetAttr(__mypyc_self__, r3, r2)
     r5 = r4 >= 0 :: signed
     return 1
-def f_obj.__get__(__mypyc_self__, instance, owner):
-    __mypyc_self__, instance, owner, r0 :: object
-    r1 :: bit
-    r2 :: object
-L0:
-    r0 = load_address _Py_NoneStruct
-    r1 = instance == r0
-    if r1 goto L1 else goto L2 :: bool
-L1:
-    return __mypyc_self__
-L2:
-    r2 = PyMethod_New(__mypyc_self__, instance)
-    return r2
-def f_obj.register(__mypyc_self__, cls, func):
-    __mypyc_self__ :: __main__.f_obj
-    cls, func, r0 :: object
-L0:
-    r0 = CPySingledispatch_RegisterFunction(__mypyc_self__, cls, func)
-    return r0
 def f_obj.__call__(__mypyc_self__, arg):
     __mypyc_self__ :: __main__.f_obj
     arg :: object
@@ -114,8 +95,168 @@ L7:
     r22 = PyObject_CallFunctionObjArgs(r6, arg, 0)
     r23 = unbox(bool, r22)
     return r23
+def f_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def f_obj.register(__mypyc_self__, cls, func):
+    __mypyc_self__ :: __main__.f_obj
+    cls, func, r0 :: object
+L0:
+    r0 = CPySingledispatch_RegisterFunction(__mypyc_self__, cls, func)
+    return r0
+def f(arg):
+    arg :: object
+    r0 :: dict
+    r1 :: str
+    r2 :: object
+    r3 :: bool
+L0:
+    r0 = __main__.globals :: static
+    r1 = 'f'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = f_obj.__call__(r2, arg)
+    return r3
 def g(arg):
     arg :: int
 L0:
     return 1
 
+
+[case testCallsToSingledispatchFunctionsAreNative]
+from functools import singledispatch
+
+@singledispatch
+def f(x: object) -> None:
+    pass
+
+def test():
+    f('a')
+[out]
+def __mypyc_singledispatch_main_function_f__(x):
+    x :: object
+L0:
+    return 1
+def f_obj.__init__(__mypyc_self__):
+    __mypyc_self__ :: __main__.f_obj
+    r0 :: dict
+    r1 :: bool
+    r2 :: dict
+    r3 :: str
+    r4 :: int32
+    r5 :: bit
+L0:
+    r0 = PyDict_New()
+    __mypyc_self__.registry = r0; r1 = is_error
+    r2 = PyDict_New()
+    r3 = 'dispatch_cache'
+    r4 = PyObject_SetAttr(__mypyc_self__, r3, r2)
+    r5 = r4 >= 0 :: signed
+    return 1
+def f_obj.__call__(__mypyc_self__, x):
+    __mypyc_self__ :: __main__.f_obj
+    x :: object
+    r0 :: ptr
+    r1 :: object
+    r2 :: dict
+    r3, r4 :: object
+    r5 :: bit
+    r6, r7 :: object
+    r8 :: str
+    r9 :: object
+    r10 :: dict
+    r11 :: object
+    r12 :: int32
+    r13 :: bit
+    r14 :: object
+    r15 :: ptr
+    r16 :: object
+    r17 :: bit
+    r18 :: int
+    r19 :: object
+    r20 :: None
+L0:
+    r0 = get_element_ptr x ob_type :: PyObject
+    r1 = load_mem r0 :: builtins.object*
+    keep_alive x
+    r2 = __mypyc_self__.dispatch_cache
+    r3 = CPyDict_GetWithNone(r2, r1)
+    r4 = load_address _Py_NoneStruct
+    r5 = r3 != r4
+    if r5 goto L1 else goto L2 :: bool
+L1:
+    r6 = r3
+    goto L3
+L2:
+    r7 = functools :: module
+    r8 = '_find_impl'
+    r9 = CPyObject_GetAttr(r7, r8)
+    r10 = __mypyc_self__.registry
+    r11 = PyObject_CallFunctionObjArgs(r9, r1, r10, 0)
+    r12 = CPyDict_SetItem(r2, r1, r11)
+    r13 = r12 >= 0 :: signed
+    r6 = r11
+L3:
+    r14 = load_address PyLong_Type
+    r15 = get_element_ptr r6 ob_type :: PyObject
+    r16 = load_mem r15 :: builtins.object*
+    keep_alive r6
+    r17 = r16 == r14
+    if r17 goto L4 else goto L5 :: bool
+L4:
+    r18 = unbox(int, r6)
+    unreachable
+L5:
+    r19 = PyObject_CallFunctionObjArgs(r6, x, 0)
+    r20 = unbox(None, r19)
+    return r20
+def f_obj.__get__(__mypyc_self__, instance, owner):
+    __mypyc_self__, instance, owner, r0 :: object
+    r1 :: bit
+    r2 :: object
+L0:
+    r0 = load_address _Py_NoneStruct
+    r1 = instance == r0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    return __mypyc_self__
+L2:
+    r2 = PyMethod_New(__mypyc_self__, instance)
+    return r2
+def f_obj.register(__mypyc_self__, cls, func):
+    __mypyc_self__ :: __main__.f_obj
+    cls, func, r0 :: object
+L0:
+    r0 = CPySingledispatch_RegisterFunction(__mypyc_self__, cls, func)
+    return r0
+def f(x):
+    x :: object
+    r0 :: dict
+    r1 :: str
+    r2 :: object
+    r3 :: None
+L0:
+    r0 = __main__.globals :: static
+    r1 = 'f'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = f_obj.__call__(r2, x)
+    return r3
+def test():
+    r0 :: str
+    r1 :: None
+    r2 :: object
+L0:
+    r0 = 'a'
+    r1 = f(r0)
+    r2 = box(None, 1)
+    return r2
+