Generate dispatch function separately from original function

Instead of adding the logic to dispatch to the correct function to the
main singledispatch function itself, generate 2 separate functions: one
that compiles the logic of the undecorated function normally, and a
separate dispatch function that checks the first argument's type and
dispatches either to one of the registered implementations or the
fallback.

That makes it easier to generate code when the main singledispatch
function is a generator or some other kind of special function, as the
code generation can generate the main singledispatch function the same
way it would otherwise.

Because we're no longer changing how that main singledispatch function
is generated based on whether it's a singledispatch function, we don't
need to store whether the current function is a singledispatch function
in FuncInfo.

Author: pranavrajpal

mypyc/irbuild/context.py

@@ -22,8 +22,7 @@ def __init__(self,
                  is_nested: bool = False,
                  contains_nested: bool = False,
                  is_decorated: bool = False,
-                 in_non_ext: bool = False,
-                 is_singledispatch: bool = False) -> None:
+                 in_non_ext: bool = False) -> None:
         self.fitem = fitem
         self.name = name if not is_decorated else decorator_helper_name(name)
         self.class_name = class_name
@@ -48,7 +47,6 @@ def __init__(self,
         self.contains_nested = contains_nested
         self.is_decorated = is_decorated
         self.in_non_ext = in_non_ext
-        self.is_singledispatch = is_singledispatch
 
         # TODO: add field for ret_type: RType = none_rprimitive
 

mypyc/irbuild/function.py

@@ -219,9 +219,12 @@ def c() -> None:
         in_non_ext = not ir.is_ext_class
         class_name = cdef.name
 
-    builder.enter(FuncInfo(fitem, name, class_name, gen_func_ns(builder),
-                           is_nested, contains_nested, is_decorated, in_non_ext,
-                           is_singledispatch))
+    if is_singledispatch:
+        func_name = '__mypyc_singledispatch_main_function_{}__'.format(name)
+    else:
+        func_name = name
+    builder.enter(FuncInfo(fitem, func_name, class_name, gen_func_ns(builder),
+                           is_nested, contains_nested, is_decorated, in_non_ext))
 
     # Functions that contain nested functions need an environment class to store variables that
     # are free in their nested functions. Generator functions need an environment class to
@@ -254,9 +257,6 @@ def c() -> None:
     if builder.fn_info.contains_nested and not builder.fn_info.is_generator:
         finalize_env_class(builder)
 
-    if builder.fn_info.is_singledispatch:
-        add_singledispatch_registered_impls(builder)
-
     builder.ret_types[-1] = sig.ret_type
 
     # Add all variables and functions that are declared/defined within this
@@ -313,6 +313,20 @@ def c() -> None:
     # calculate them *once* when the function definition is evaluated.
     calculate_arg_defaults(builder, fn_info, func_reg, symtable)
 
+    if is_singledispatch:
+        # add the generated main singledispatch function
+        builder.functions.append(func_ir)
+        # create a dispatch function (a function that checks the first argument type and dispatches
+        # to the correct implementation)
+        builder.enter()
+        assert isinstance(fitem, FuncDef)
+        generate_singledispatch_dispatch_function(builder, fn_info.name, fitem)
+        args, _, blocks, _, fn_info = builder.leave()
+        dispatch_name = decorator_helper_name(name) if is_decorated else name
+        func_decl = FuncDecl(dispatch_name, None, builder.module_name, sig)
+        dispatch_func_ir = FuncIR(func_decl, args, blocks)
+        return dispatch_func_ir, None
+
     return (func_ir, func_reg)
 
 
@@ -772,9 +786,11 @@ def check_if_isinstance(builder: IRBuilder, obj: Value, typ: TypeInfo, line: int
         return builder.call_c(slow_isinstance_op, [obj, class_obj], line)
 
 
-def add_singledispatch_registered_impls(builder: IRBuilder) -> None:
-    fitem = builder.fn_info.fitem
-    assert isinstance(fitem, FuncDef)
+def generate_singledispatch_dispatch_function(
+    builder: IRBuilder,
+    main_singledispatch_function_name: str,
+    fitem: FuncDef,
+) -> None:
     impls = builder.singledispatch_impls[fitem]
     line = fitem.line
     current_func_decl = builder.mapper.func_to_decl[fitem]
@@ -801,3 +817,5 @@ def gen_func_call_and_return(func_name: str) -> None:
 
         gen_func_call_and_return(impl.name)
         builder.activate_block(next_impl)
+
+    gen_func_call_and_return(main_singledispatch_function_name)