[mypyc] Add initial support for compiling singledispatch functions

mypyc/irbuild/builder.py

@@ -115,6 +115,7 @@ def __init__(self,
         self.encapsulating_funcs = pbv.encapsulating_funcs
         self.nested_fitems = pbv.nested_funcs.keys()
         self.fdefs_to_decorators = pbv.funcs_to_decorators
+        self.singledispatch_impls = pbv.singledispatch_impls
 
         self.visitor = visitor
 

mypyc/irbuild/context.py

@@ -22,7 +22,8 @@ def __init__(self,
                  is_nested: bool = False,
                  contains_nested: bool = False,
                  is_decorated: bool = False,
-                 in_non_ext: bool = False) -> None:
+                 in_non_ext: bool = False,
+                 is_singledispatch: bool = False) -> None:
         self.fitem = fitem
         self.name = name if not is_decorated else decorator_helper_name(name)
         self.class_name = class_name
@@ -47,6 +48,7 @@ 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

@@ -10,11 +10,11 @@
 instance of the callable class.
 """
 
-from typing import Optional, List, Tuple, Union, Dict
+from typing import NamedTuple, Optional, List, Sequence, Tuple, Union, Dict
 
 from mypy.nodes import (
     ClassDef, FuncDef, OverloadedFuncDef, Decorator, Var, YieldFromExpr, AwaitExpr, YieldExpr,
-    FuncItem, LambdaExpr, SymbolNode, ARG_NAMED, ARG_NAMED_OPT
+    FuncItem, LambdaExpr, SymbolNode, ARG_NAMED, ARG_NAMED_OPT, TypeInfo
 )
 from mypy.types import CallableType, get_proper_type
 
@@ -28,7 +28,9 @@
 )
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.primitives.generic_ops import py_setattr_op, next_raw_op, iter_op
-from mypyc.primitives.misc_ops import check_stop_op, yield_from_except_op, coro_op, send_op
+from mypyc.primitives.misc_ops import (
+    check_stop_op, yield_from_except_op, coro_op, send_op, slow_isinstance_op
+)
 from mypyc.primitives.dict_ops import dict_set_item_op
 from mypyc.common import SELF_NAME, LAMBDA_NAME, decorator_helper_name
 from mypyc.sametype import is_same_method_signature
@@ -84,7 +86,10 @@ def transform_decorator(builder: IRBuilder, dec: Decorator) -> None:
         decorated_func = load_decorated_func(builder, dec.func, func_reg)
         builder.assign(get_func_target(builder, dec.func), decorated_func, dec.func.line)
         func_reg = decorated_func
-    else:
+    # If the prebuild pass didn't put this function in the function to decorators map (for example
+    # if this is a registered singledispatch implementation with no other decorators), we should
+    # treat this function as a regular function, not a decorated function
+    elif dec.func in builder.fdefs_to_decorators:
         # Obtain the the function name in order to construct the name of the helper function.
         name = dec.func.fullname.split('.')[-1]
         helper_name = decorator_helper_name(name)
@@ -206,6 +211,7 @@ def c() -> None:
     is_nested = fitem in builder.nested_fitems or isinstance(fitem, LambdaExpr)
     contains_nested = fitem in builder.encapsulating_funcs.keys()
     is_decorated = fitem in builder.fdefs_to_decorators
+    is_singledispatch = fitem in builder.singledispatch_impls
     in_non_ext = False
     class_name = None
     if cdef:
@@ -214,7 +220,8 @@ def c() -> None:
         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_nested, contains_nested, is_decorated, in_non_ext,
+                           is_singledispatch))
 
     # 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
@@ -247,6 +254,9 @@ 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
@@ -628,6 +638,23 @@ def gen_glue(builder: IRBuilder, sig: FuncSignature, target: FuncIR,
         return gen_glue_method(builder, sig, target, cls, base, fdef.line, do_py_ops)
 
 
+class ArgInfo(NamedTuple):
+    args: List[Value]
+    arg_names: List[Optional[str]]
+    arg_kinds: List[int]
+
+
+def get_args(builder: IRBuilder, rt_args: Sequence[RuntimeArg], line: int) -> ArgInfo:
+    # The environment operates on Vars, so we make some up
+    fake_vars = [(Var(arg.name), arg.type) for arg in rt_args]
+    args = [builder.read(builder.add_local_reg(var, type, is_arg=True), line)
+            for var, type in fake_vars]
+    arg_names = [arg.name if arg.kind in (ARG_NAMED, ARG_NAMED_OPT) else None
+                 for arg in rt_args]
+    arg_kinds = [concrete_arg_kind(arg.kind) for arg in rt_args]
+    return ArgInfo(args, arg_names, arg_kinds)
+
+
 def gen_glue_method(builder: IRBuilder, sig: FuncSignature, target: FuncIR,
                     cls: ClassIR, base: ClassIR, line: int,
                     do_pycall: bool,
@@ -664,13 +691,8 @@ def f(builder: IRBuilder, x: object) -> int: ...
     if target.decl.kind == FUNC_NORMAL:
         rt_args[0] = RuntimeArg(sig.args[0].name, RInstance(cls))
 
-    # The environment operates on Vars, so we make some up
-    fake_vars = [(Var(arg.name), arg.type) for arg in rt_args]
-    args = [builder.read(builder.add_local_reg(var, type, is_arg=True), line)
-            for var, type in fake_vars]
-    arg_names = [arg.name if arg.kind in (ARG_NAMED, ARG_NAMED_OPT) else None
-                 for arg in rt_args]
-    arg_kinds = [concrete_arg_kind(arg.kind) for arg in rt_args]
+    arg_info = get_args(builder, rt_args, line)
+    args, arg_kinds, arg_names = arg_info.args, arg_info.arg_kinds, arg_info.arg_names
 
     if do_pycall:
         retval = builder.builder.py_method_call(
@@ -739,3 +761,35 @@ def get_func_target(builder: IRBuilder, fdef: FuncDef) -> AssignmentTarget:
         return builder.lookup(fdef)
 
     return builder.add_local_reg(fdef, object_rprimitive)
+
+
+def check_if_isinstance(builder: IRBuilder, obj: Value, typ: TypeInfo, line: int) -> Value:
+    if typ in builder.mapper.type_to_ir:
+        class_ir = builder.mapper.type_to_ir[typ]
+        return builder.builder.isinstance_native(obj, class_ir, line)
+    else:
+        class_obj = builder.load_module_attr_by_fullname(typ.fullname, line)
+        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)
+    impls = builder.singledispatch_impls[fitem]
+    line = fitem.line
+    current_func_decl = builder.mapper.func_to_decl[fitem]
+    arg_info = get_args(builder, current_func_decl.sig.args, line)
+    for dispatch_type, impl in impls:
+        func_decl = builder.mapper.func_to_decl[impl]
+        call_impl, next_impl = BasicBlock(), BasicBlock()
+        should_call_impl = check_if_isinstance(builder, arg_info.args[0], dispatch_type, line)
+        builder.add_bool_branch(should_call_impl, call_impl, next_impl)
+
+        # Call the registered implementation
+        builder.activate_block(call_impl)
+
+        ret_val = builder.builder.call(
+            func_decl, arg_info.args, arg_info.arg_kinds, arg_info.arg_names, line
+        )
+        builder.nonlocal_control[-1].gen_return(builder, ret_val, line)
+        builder.activate_block(next_impl)

mypyc/irbuild/prebuildvisitor.py

@@ -1,7 +1,10 @@
-from typing import Dict, List, Set
+from mypy.types import Instance, get_proper_type
+from typing import DefaultDict, Dict, List, NamedTuple, Set, Optional, Tuple
+from collections import defaultdict
 
 from mypy.nodes import (
-    Decorator, Expression, FuncDef, FuncItem, LambdaExpr, NameExpr, SymbolNode, Var, MemberExpr
+    Decorator, Expression, FuncDef, FuncItem, LambdaExpr, NameExpr, SymbolNode, Var, MemberExpr,
+    CallExpr, RefExpr, TypeInfo
 )
 from mypy.traverser import TraverserVisitor
 
@@ -50,6 +53,10 @@ def __init__(self) -> None:
         # Map function to its non-special decorators.
         self.funcs_to_decorators: Dict[FuncDef, List[Expression]] = {}
 
+        # Map of main singledispatch function to list of registered implementations
+        self.singledispatch_impls: DefaultDict[
+            FuncDef, List[Tuple[TypeInfo, FuncDef]]] = defaultdict(list)
+
     def visit_decorator(self, dec: Decorator) -> None:
         if dec.decorators:
             # Only add the function being decorated if there exist
@@ -63,6 +70,20 @@ def visit_decorator(self, dec: Decorator) -> None:
                 # Property setters are not treated as decorated methods.
                 self.prop_setters.add(dec.func)
             else:
+                removed: List[int] = []
+                for i, d in enumerate(dec.decorators):
+                    impl = get_singledispatch_register_call_info(d, dec.func)
+                    if impl is not None:
+                        self.singledispatch_impls[impl.singledispatch_func].append(
+                            (impl.dispatch_type, dec.func))
+                        removed.append(i)
+                for i in reversed(removed):
+                    del dec.decorators[i]
+                # if the only decorators are register calls, we shouldn't treat this
+                # as a decorated function because there aren't any decorators to apply
+                if not dec.decorators:
+                    return
+
                 self.funcs_to_decorators[dec.func] = dec.decorators
         super().visit_decorator(dec)
 
@@ -141,3 +162,45 @@ def add_free_variable(self, symbol: SymbolNode) -> None:
         # and mark is as a non-local symbol within that function.
         func = self.symbols_to_funcs[symbol]
         self.free_variables.setdefault(func, set()).add(symbol)
+
+
+class RegisteredImpl(NamedTuple):
+    singledispatch_func: FuncDef
+    dispatch_type: TypeInfo
+
+
+def get_singledispatch_register_call_info(decorator: Expression, func: FuncDef
+                                          ) -> Optional[RegisteredImpl]:
+    # @fun.register(complex)
+    # def g(arg): ...
+    if (isinstance(decorator, CallExpr) and len(decorator.args) == 1
+            and isinstance(decorator.args[0], RefExpr)):
+        callee = decorator.callee
+        dispatch_type = decorator.args[0].node
+        if not isinstance(dispatch_type, TypeInfo):
+            return None
+
+        if isinstance(callee, MemberExpr):
+            return registered_impl_from_possible_register_call(callee, dispatch_type)
+    # @fun.register
+    # def g(arg: int): ...
+    elif isinstance(decorator, MemberExpr):
+        # we don't know if this is a register call yet, so we can't be sure that the function
+        # actually has arguments
+        if not func.arguments:
+            return None
+        arg_type = get_proper_type(func.arguments[0].variable.type)
+        if not isinstance(arg_type, Instance):
+            return None
+        info = arg_type.type
+        return registered_impl_from_possible_register_call(decorator, info)
+    return None
+
+
+def registered_impl_from_possible_register_call(expr: MemberExpr, dispatch_type: TypeInfo
+                                                ) -> Optional[RegisteredImpl]:
+    if expr.name == 'register' and isinstance(expr.expr, NameExpr):
+        node = expr.expr.node
+        if isinstance(node, Decorator):
+            return RegisteredImpl(node.func, dispatch_type)
+    return None

mypyc/primitives/misc_ops.py

@@ -141,7 +141,7 @@
     is_borrowed=True)
 
 # isinstance(obj, cls)
-function_op(
+slow_isinstance_op = function_op(
     name='builtins.isinstance',
     arg_types=[object_rprimitive, object_rprimitive],
     return_type=c_int_rprimitive,

mypyc/test-data/run-singledispatch.test

@@ -2,7 +2,7 @@
 # Most of these tests are marked as xfails because mypyc doesn't support singledispatch yet
 # (These tests will be re-enabled when mypyc supports singledispatch)
 
-[case testSpecializedImplementationUsed-xfail]
+[case testSpecializedImplementationUsed]
 from functools import singledispatch
 
 @singledispatch
@@ -17,7 +17,7 @@ def test_specialize() -> None:
     assert fun('a')
     assert not fun(3)
 
-[case testSubclassesOfExpectedTypeUseSpecialized-xfail]
+[case testSubclassesOfExpectedTypeUseSpecialized]
 from functools import singledispatch
 class A: pass
 class B(A): pass
@@ -76,7 +76,7 @@ def test_singledispatch() -> None:
     assert fun('a') == 'str'
     assert fun({'a': 'b'}) == 'default'
 
-[case testCanRegisterCompiledClasses-xfail]
+[case testCanRegisterCompiledClasses]
 from functools import singledispatch
 class A: pass
 
@@ -136,21 +136,6 @@ def fun_specialized(arg: int) -> bool:
 def test_singledispatch() -> None:
     assert fun_specialized('a')
 
-[case testTypeAnnotationsDisagreeWithRegisterArgument-xfail]
-from functools import singledispatch
-
-@singledispatch
-def fun(arg) -> bool:
-    return False
-
-@fun.register(int)
-def fun_specialized(arg: str) -> bool:
-    return True
-
-def test_singledispatch() -> None:
-    assert fun(3) # type: ignore
-    assert not fun('a')
-
 [case testNoneIsntATypeWhenUsedAsArgumentToRegister-xfail]
 from functools import singledispatch
 
@@ -385,3 +370,55 @@ def test_verify() -> None:
     assert verify_list(MypyFile(), 5, ['a', 'b']) == ['in TypeInfo', 'hello']
     assert verify_list(TypeInfo(), str, ['a', 'b']) == ['in TypeInfo', 'hello']
     assert verify_list(TypeVarExpr(), 'a', ['x', 'y']) == ['x', 'y']
+
+[case testArgsInRegisteredImplNamedDifferentlyFromMainFunction]
+from functools import singledispatch
+
+@singledispatch
+def f(a) -> bool:
+    return False
+
+@f.register
+def g(b: int) -> bool:
+    return True
+
+def test_singledispatch():
+    assert f(5)
+    assert not f('a')
+
+[case testKeywordArguments-xfail]
+from functools import singledispatch
+
+@singledispatch
+def f(arg, *, kwarg: bool = False) -> bool:
+    return not kwarg
+
+@f.register
+def g(arg: int, *, kwarg: bool = True) -> bool:
+    return kwarg
+
+def test_keywords():
+    assert f('a')
+    assert f('a', kwarg=False)
+    assert not f('a', kwarg=True)
+
+    assert f(1)
+    assert f(1, kwarg=True)
+    assert not f(1, kwarg=False)
+
+[case testGeneratorAndMultipleTypesOfIterable-xfail]
+from functools import singledispatch
+from typing import *
+
+@singledispatch
+def f(arg: Any) -> Iterable[int]:
+    yield 1
+
+@f.register
+def g(arg: str) -> Iterable[int]:
+    return [0]
+
+def test_iterables():
+    assert f(1) != [1]
+    assert list(f(1)) == [1]
+    assert f('a') == [0]