[mypyc] Use native calls in singledispatch dispatch functions

mypyc/irbuild/function.py

@@ -833,21 +833,31 @@ def generate_singledispatch_dispatch_function(
     current_func_decl = builder.mapper.func_to_decl[fitem]
     arg_info = get_args(builder, current_func_decl.sig.args, line)
 
-    def gen_func_call_and_return(func_name: str, fullname: Optional[str] = None) -> None:
-        func = load_func(builder, func_name, fullname, line)
-        ret_val = builder.builder.py_call(
-            func, arg_info.args, line, arg_info.arg_kinds, arg_info.arg_names
-        )
+    def gen_func_call_and_return(
+        func_name: str,
+        fdef: FuncDef,
+        fullname: Optional[str] = None
+    ) -> None:
+        if is_decorated(builder, fdef):
+            func = load_func(builder, func_name, fullname, line)
+            ret_val = builder.builder.py_call(
+                func, arg_info.args, line, arg_info.arg_kinds, arg_info.arg_names
+            )
+        else:
+            func_decl = builder.mapper.func_to_decl[fdef]
+            ret_val = builder.builder.call(
+                func_decl, arg_info.args, arg_info.arg_kinds, arg_info.arg_names, line
+            )
         coerced = builder.coerce(ret_val, current_func_decl.sig.ret_type, line)
         builder.nonlocal_control[-1].gen_return(builder, coerced, line)
 
     # Add all necessary imports of other modules that have registered functions in other modules
     # We're doing this in a separate pass over the implementations because that avoids the
     # complexity and code size implications of generating this import before every call to a
     # registered implementation that might need this imported
-    # TODO: avoid adding imports if we use native calls for all of the registered implementations
-    # in a module (once we add support for using native calls for registered implementations)
     for _, impl in impls:
+        if not is_decorated(builder, impl):
+            continue
         module_name = impl.fullname.rsplit('.')[0]
         if module_name not in builder.imports:
             # We need to generate an import here because the module needs to be imported before we
@@ -868,15 +878,15 @@ def gen_func_call_and_return(func_name: str, fullname: Optional[str] = None) ->
         # The shortname of a function is just '{class}.{func_name}', and we don't support
         # singledispatchmethod yet, so that is always the same as the function name
         name = short_id_from_name(impl.name, impl.name, impl.line)
-        gen_func_call_and_return(name, fullname=impl.fullname)
+        gen_func_call_and_return(name, impl, fullname=impl.fullname)
         builder.activate_block(next_impl)
 
     # We don't pass fullname here because getting the fullname of the main generated singledispatch
     # function isn't easy, and we don't need it because the fullname is only needed for making sure
     # we load the function from another module instead of the globals dict if it's defined in
     # another module, which will never be true for the main singledispatch function (it's always
     # generated in the same module as the dispatch function)
-    gen_func_call_and_return(main_singledispatch_function_name)
+    gen_func_call_and_return(main_singledispatch_function_name, fitem)
 
 
 def gen_dispatch_func_ir(

mypyc/test-data/irbuild-singledispatch.test

@@ -0,0 +1,47 @@
+[case testNativeCallsUsedInDispatchFunction]
+from functools import singledispatch
+@singledispatch
+def f(arg) -> bool:
+    return False
+
+@f.register
+def g(arg: int) -> bool:
+    return True
+[out]
+def __mypyc___mypyc_singledispatch_main_function_f___decorator_helper__(arg):
+    arg :: object
+L0:
+    return 0
+def __mypyc_f_decorator_helper__(arg):
+    arg, r0 :: object
+    r1 :: int32
+    r2 :: bit
+    r3 :: bool
+    r4 :: int
+    r5 :: bool
+    r6 :: dict
+    r7 :: str
+    r8, r9 :: object
+    r10 :: bool
+L0:
+    r0 = load_address PyLong_Type
+    r1 = PyObject_IsInstance(arg, r0)
+    r2 = r1 >= 0 :: signed
+    r3 = truncate r1: int32 to builtins.bool
+    if r3 goto L1 else goto L2 :: bool
+L1:
+    r4 = unbox(int, arg)
+    r5 = g(r4)
+    return r5
+L2:
+    r6 = __main__.globals :: static
+    r7 = '__mypyc___mypyc_singledispatch_main_function_f___decorator_helper__'
+    r8 = CPyDict_GetItem(r6, r7)
+    r9 = PyObject_CallFunctionObjArgs(r8, arg, 0)
+    r10 = unbox(bool, r9)
+    return r10
+def g(arg):
+    arg :: int
+L0:
+    return 1
+

mypyc/test/test_irbuild.py

@@ -34,6 +34,7 @@
     'irbuild-unreachable.test',
     'irbuild-isinstance.test',
     'irbuild-dunders.test',
+    'irbuild-singledispatch.test',
 ]