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.
sthagen · Jul 9, 2021 · 59555e4 · 59555e4
1 parent 099b047
commit 59555e4
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 12 deletions.
diff --git a/mypyc/irbuild/context.py b/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
 

diff --git a/mypyc/irbuild/function.py b/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)