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

* Wrap singledispatch callable class in glue function

mypyc doesn't have support for calling callable classes without going
through the python API, so instead of having to special case the
singledispatch callable class when we want to generate a native call to
it, we create a glue function that loads the callable class from the
globals dict and then makes a native call to that callable class's
__call__ method.

* Use native calls for singledispatch functions

Use a native call to the glue function that we generate for
singledispatch functions instead of calling those singledispatch
functions through the python API.

* Add irbuild test for native calls to singledispatch functions

Add an irbuild test to make sure we use native calls to the glue
functions for singledispatch functions, and that the glue functions make
native calls to the __call__ methods on the callable classes.

* Fix irbuild test

Author: pranavrajpal

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
+