mypyc/irbuild/mapper.py

"""Maintain a mapping from mypy concepts to IR/compiled concepts."""

from typing import Dict, Optional, Union
from mypy.ordered_dict import OrderedDict

from mypy.nodes import FuncDef, TypeInfo, SymbolNode, ARG_STAR, ARG_STAR2
from mypy.types import (
    Instance, Type, CallableType, LiteralType, TypedDictType, UnboundType, PartialType,
    UninhabitedType, Overloaded, UnionType, TypeType, AnyType, NoneTyp, TupleType, TypeVarType,
    get_proper_type
)

from mypyc.ir.ops import LiteralsMap
from mypyc.ir.rtypes import (
    RType, RUnion, RTuple, RInstance, object_rprimitive, dict_rprimitive, tuple_rprimitive,
    none_rprimitive, int_rprimitive, float_rprimitive, str_rprimitive, bool_rprimitive,
    list_rprimitive, set_rprimitive
)
from mypyc.ir.func_ir import FuncSignature, FuncDecl, RuntimeArg
from mypyc.ir.class_ir import ClassIR


class Mapper:
    """Keep track of mappings from mypy concepts to IR concepts.

    For example, we keep track of how the mypy TypeInfos of compiled
    classes map to class IR objects.

    This state is shared across all modules being compiled in all
    compilation groups.
    """

    def __init__(self, group_map: Dict[str, Optional[str]]) -> None:
        self.group_map = group_map
        self.type_to_ir = {}  # type: Dict[TypeInfo, ClassIR]
        self.func_to_decl = {}  # type: Dict[SymbolNode, FuncDecl]
        # LiteralsMap maps literal values to a static name. Each
        # compilation group has its own LiteralsMap. (Since they can't
        # share literals.)
        self.literals = {
            v: OrderedDict() for v in group_map.values()
        }  # type: Dict[Optional[str], LiteralsMap]

    def type_to_rtype(self, typ: Optional[Type]) -> RType:
        if typ is None:
            return object_rprimitive

        typ = get_proper_type(typ)
        if isinstance(typ, Instance):
            if typ.type.fullname == 'builtins.int':
                return int_rprimitive
            elif typ.type.fullname == 'builtins.float':
                return float_rprimitive
            elif typ.type.fullname == 'builtins.str':
                return str_rprimitive
            elif typ.type.fullname == 'builtins.bool':
                return bool_rprimitive
            elif typ.type.fullname == 'builtins.list':
                return list_rprimitive
            # Dict subclasses are at least somewhat common and we
            # specifically support them, so make sure that dict operations
            # get optimized on them.
            elif any(cls.fullname == 'builtins.dict' for cls in typ.type.mro):
                return dict_rprimitive
            elif typ.type.fullname == 'builtins.set':
                return set_rprimitive
            elif typ.type.fullname == 'builtins.tuple':
                return tuple_rprimitive  # Varying-length tuple
            elif typ.type in self.type_to_ir:
                inst = RInstance(self.type_to_ir[typ.type])
                # Treat protocols as Union[protocol, object], so that we can do fast
                # method calls in the cases where the protocol is explicitly inherited from
                # and fall back to generic operations when it isn't.
                if typ.type.is_protocol:
                    return RUnion([inst, object_rprimitive])
                else:
                    return inst
            else:
                return object_rprimitive
        elif isinstance(typ, TupleType):
            # Use our unboxed tuples for raw tuples but fall back to
            # being boxed for NamedTuple.
            if typ.partial_fallback.type.fullname == 'builtins.tuple':
                return RTuple([self.type_to_rtype(t) for t in typ.items])
            else:
                return tuple_rprimitive
        elif isinstance(typ, CallableType):
            return object_rprimitive
        elif isinstance(typ, NoneTyp):
            return none_rprimitive
        elif isinstance(typ, UnionType):
            return RUnion([self.type_to_rtype(item)
                           for item in typ.items])
        elif isinstance(typ, AnyType):
            return object_rprimitive
        elif isinstance(typ, TypeType):
            return object_rprimitive
        elif isinstance(typ, TypeVarType):
            # Erase type variable to upper bound.
            # TODO: Erase to union if object has value restriction?
            return self.type_to_rtype(typ.upper_bound)
        elif isinstance(typ, PartialType):
            assert typ.var.type is not None
            return self.type_to_rtype(typ.var.type)
        elif isinstance(typ, Overloaded):
            return object_rprimitive
        elif isinstance(typ, TypedDictType):
            return dict_rprimitive
        elif isinstance(typ, LiteralType):
            return self.type_to_rtype(typ.fallback)
        elif isinstance(typ, (UninhabitedType, UnboundType)):
            # Sure, whatever!
            return object_rprimitive

        # I think we've covered everything that is supposed to
        # actually show up, so anything else is a bug somewhere.
        assert False, 'unexpected type %s' % type(typ)

    def get_arg_rtype(self, typ: Type, kind: int) -> RType:
        if kind == ARG_STAR:
            return tuple_rprimitive
        elif kind == ARG_STAR2:
            return dict_rprimitive
        else:
            return self.type_to_rtype(typ)

    def fdef_to_sig(self, fdef: FuncDef) -> FuncSignature:
        if isinstance(fdef.type, CallableType):
            arg_types = [self.get_arg_rtype(typ, kind)
                         for typ, kind in zip(fdef.type.arg_types, fdef.type.arg_kinds)]
            ret = self.type_to_rtype(fdef.type.ret_type)
        else:
            # Handle unannotated functions
            arg_types = [object_rprimitive for arg in fdef.arguments]
            ret = object_rprimitive

        args = [RuntimeArg(arg_name, arg_type, arg_kind)
                for arg_name, arg_kind, arg_type in zip(fdef.arg_names, fdef.arg_kinds, arg_types)]

        # We force certain dunder methods to return objects to support letting them
        # return NotImplemented. It also avoids some pointless boxing and unboxing,
        # since tp_richcompare needs an object anyways.
        if fdef.name in ('__eq__', '__ne__', '__lt__', '__gt__', '__le__', '__ge__'):
            ret = object_rprimitive
        return FuncSignature(args, ret)

    def literal_static_name(self, module: str,
                            value: Union[int, float, complex, str, bytes]) -> str:
        # Literals are shared between modules in a compilation group
        # but not outside the group.
        literals = self.literals[self.group_map.get(module)]

        # Include type to distinguish between 1 and 1.0, and so on.
        key = (type(value), value)
        if key not in literals:
            if isinstance(value, str):
                prefix = 'unicode_'
            else:
                prefix = type(value).__name__ + '_'
            literals[key] = prefix + str(len(literals))
        return literals[key]