dispatcher.py

# -*- coding: utf-8 -*-


import collections
import functools
import sys
import types as pytypes
import uuid
import weakref
from contextlib import ExitStack
from abc import abstractmethod

from numba import _dispatcher
from numba.core import (
    utils, types, errors, typing, serialize, config, compiler, sigutils
)
from numba.core.compiler_lock import global_compiler_lock
from numba.core.typeconv.rules import default_type_manager
from numba.core.typing.templates import fold_arguments
from numba.core.typing.typeof import Purpose, typeof
from numba.core.bytecode import get_code_object
from numba.core.caching import NullCache, FunctionCache
from numba.core import entrypoints
import numba.core.event as ev


class OmittedArg(object):
    """
    A placeholder for omitted arguments with a default value.
    """

    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return "omitted arg(%r)" % (self.value,)

    @property
    def _numba_type_(self):
        return types.Omitted(self.value)


class _FunctionCompiler(object):
    def __init__(self, py_func, targetdescr, targetoptions, locals,
                 pipeline_class):
        self.py_func = py_func
        self.targetdescr = targetdescr
        self.targetoptions = targetoptions
        self.locals = locals
        self.pysig = utils.pysignature(self.py_func)
        self.pipeline_class = pipeline_class
        # Remember key=(args, return_type) combinations that will fail
        # compilation to avoid compilation attempt on them.  The values are
        # the exceptions.
        self._failed_cache = {}

    def fold_argument_types(self, args, kws):
        """
        Given positional and named argument types, fold keyword arguments
        and resolve defaults by inserting types.Omitted() instances.

        A (pysig, argument types) tuple is returned.
        """
        def normal_handler(index, param, value):
            return value

        def default_handler(index, param, default):
            return types.Omitted(default)

        def stararg_handler(index, param, values):
            return types.StarArgTuple(values)
        # For now, we take argument values from the @jit function
        args = fold_arguments(self.pysig, args, kws,
                              normal_handler,
                              default_handler,
                              stararg_handler)
        return self.pysig, args

    def compile(self, args, return_type):
        status, retval = self._compile_cached(args, return_type)
        if status:
            return retval
        else:
            raise retval

    def _compile_cached(self, args, return_type):
        key = tuple(args), return_type
        try:
            return False, self._failed_cache[key]
        except KeyError:
            pass

        try:
            retval = self._compile_core(args, return_type)
        except errors.TypingError as e:
            self._failed_cache[key] = e
            return False, e
        else:
            return True, retval

    def _compile_core(self, args, return_type):
        flags = compiler.Flags()
        self.targetdescr.options.parse_as_flags(flags, self.targetoptions)
        flags = self._customize_flags(flags)

        impl = self._get_implementation(args, {})
        cres = compiler.compile_extra(self.targetdescr.typing_context,
                                      self.targetdescr.target_context,
                                      impl,
                                      args=args, return_type=return_type,
                                      flags=flags, locals=self.locals,
                                      pipeline_class=self.pipeline_class)
        # Check typing error if object mode is used
        if cres.typing_error is not None and not flags.enable_pyobject:
            raise cres.typing_error
        return cres

    def get_globals_for_reduction(self):
        return serialize._get_function_globals_for_reduction(self.py_func)

    def _get_implementation(self, args, kws):
        return self.py_func

    def _customize_flags(self, flags):
        return flags


class _GeneratedFunctionCompiler(_FunctionCompiler):

    def __init__(self, py_func, targetdescr, targetoptions, locals,
                 pipeline_class):
        super(_GeneratedFunctionCompiler, self).__init__(
            py_func, targetdescr, targetoptions, locals, pipeline_class)
        self.impls = set()

    def get_globals_for_reduction(self):
        # This will recursively get the globals used by any nested
        # implementation function.
        return serialize._get_function_globals_for_reduction(self.py_func)

    def _get_implementation(self, args, kws):
        impl = self.py_func(*args, **kws)
        # Check the generating function and implementation signatures are
        # compatible, otherwise compiling would fail later.
        pysig = utils.pysignature(self.py_func)
        implsig = utils.pysignature(impl)
        ok = len(pysig.parameters) == len(implsig.parameters)
        if ok:
            for pyparam, implparam in zip(pysig.parameters.values(),
                                          implsig.parameters.values()):
                # We allow the implementation to omit default values, but
                # if it mentions them, they should have the same value...
                if (pyparam.name != implparam.name or
                    pyparam.kind != implparam.kind or
                    (implparam.default is not implparam.empty and
                     implparam.default != pyparam.default)):
                    ok = False
        if not ok:
            raise TypeError("generated implementation %s should be compatible "
                            "with signature '%s', but has signature '%s'"
                            % (impl, pysig, implsig))
        self.impls.add(impl)
        return impl


_CompileStats = collections.namedtuple(
    '_CompileStats', ('cache_path', 'cache_hits', 'cache_misses'))


class CompilingCounter(object):
    """
    A simple counter that increment in __enter__ and decrement in __exit__.
    """

    def __init__(self):
        self.counter = 0

    def __enter__(self):
        assert self.counter >= 0
        self.counter += 1

    def __exit__(self, *args, **kwargs):
        self.counter -= 1
        assert self.counter >= 0

    def __bool__(self):
        return self.counter > 0

    __nonzero__ = __bool__


class _DispatcherBase(_dispatcher.Dispatcher):
    """
    Common base class for dispatcher Implementations.
    """

    __numba__ = "py_func"

    def __init__(self, arg_count, py_func, pysig, can_fallback,
                 exact_match_required):
        self._tm = default_type_manager

        # A mapping of signatures to compile results
        self.overloads = collections.OrderedDict()

        self.py_func = py_func
        # other parts of Numba assume the old Python 2 name for code object
        self.func_code = get_code_object(py_func)
        # but newer python uses a different name
        self.__code__ = self.func_code
        # a place to keep an active reference to the types of the active call
        self._types_active_call = []
        # Default argument values match the py_func
        self.__defaults__ = py_func.__defaults__

        argnames = tuple(pysig.parameters)
        default_values = self.py_func.__defaults__ or ()
        defargs = tuple(OmittedArg(val) for val in default_values)
        try:
            lastarg = list(pysig.parameters.values())[-1]
        except IndexError:
            has_stararg = False
        else:
            has_stararg = lastarg.kind == lastarg.VAR_POSITIONAL
        _dispatcher.Dispatcher.__init__(self, self._tm.get_pointer(),
                                        arg_count, self._fold_args,
                                        argnames, defargs,
                                        can_fallback,
                                        has_stararg,
                                        exact_match_required)

        self.doc = py_func.__doc__
        self._compiling_counter = CompilingCounter()
        weakref.finalize(self, self._make_finalizer())

    def _compilation_chain_init_hook(self):
        """
        This will be called ahead of any part of compilation taking place (this
        even includes being ahead of working out the types of the arguments).
        This permits activities such as initialising extension entry points so
        that the compiler knows about additional externally defined types etc
        before it does anything.
        """
        entrypoints.init_all()

    def _reset_overloads(self):
        self._clear()
        self.overloads.clear()

    def _make_finalizer(self):
        """
        Return a finalizer function that will release references to
        related compiled functions.
        """
        overloads = self.overloads
        targetctx = self.targetctx

        # Early-bind utils.shutting_down() into the function's local namespace
        # (see issue #689)
        def finalizer(shutting_down=utils.shutting_down):
            # The finalizer may crash at shutdown, skip it (resources
            # will be cleared by the process exiting, anyway).
            if shutting_down():
                return
            # This function must *not* hold any reference to self:
            # we take care to bind the necessary objects in the closure.
            for cres in overloads.values():
                try:
                    targetctx.remove_user_function(cres.entry_point)
                except KeyError:
                    pass

        return finalizer

    @property
    def signatures(self):
        """
        Returns a list of compiled function signatures.
        """
        return list(self.overloads)

    @property
    def nopython_signatures(self):
        return [cres.signature for cres in self.overloads.values()
                if not cres.objectmode]

    def disable_compile(self, val=True):
        """Disable the compilation of new signatures at call time.
        """
        # If disabling compilation then there must be at least one signature
        assert (not val) or len(self.signatures) > 0
        self._can_compile = not val

    def add_overload(self, cres):
        args = tuple(cres.signature.args)
        sig = [a._code for a in args]
        self._insert(sig, cres.entry_point, cres.objectmode)
        self.overloads[args] = cres

    def fold_argument_types(self, args, kws):
        return self._compiler.fold_argument_types(args, kws)

    def get_call_template(self, args, kws):
        """
        Get a typing.ConcreteTemplate for this dispatcher and the given
        *args* and *kws* types.  This allows to resolve the return type.

        A (template, pysig, args, kws) tuple is returned.
        """
        # XXX how about a dispatcher template class automating the
        # following?

        # Fold keyword arguments and resolve default values
        pysig, args = self._compiler.fold_argument_types(args, kws)
        kws = {}
        # Ensure an overload is available
        if self._can_compile:
            self.compile(tuple(args))

        # Create function type for typing
        func_name = self.py_func.__name__
        name = "CallTemplate({0})".format(func_name)
        # The `key` isn't really used except for diagnosis here,
        # so avoid keeping a reference to `cfunc`.
        call_template = typing.make_concrete_template(
            name, key=func_name, signatures=self.nopython_signatures)
        return call_template, pysig, args, kws

    def get_overload(self, sig):
        """
        Return the compiled function for the given signature.
        """
        args, return_type = sigutils.normalize_signature(sig)
        return self.overloads[tuple(args)].entry_point

    @property
    def is_compiling(self):
        """
        Whether a specialization is currently being compiled.
        """
        return self._compiling_counter

    def _compile_for_args(self, *args, **kws):
        """
        For internal use.  Compile a specialized version of the function
        for the given *args* and *kws*, and return the resulting callable.
        """
        assert not kws
        # call any initialisation required for the compilation chain (e.g.
        # extension point registration).
        self._compilation_chain_init_hook()

        def error_rewrite(e, issue_type):
            """
            Rewrite and raise Exception `e` with help supplied based on the
            specified issue_type.
            """
            if config.SHOW_HELP:
                help_msg = errors.error_extras[issue_type]
                e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
            if config.FULL_TRACEBACKS:
                raise e
            else:
                raise e.with_traceback(None)

        argtypes = []
        for a in args:
            if isinstance(a, OmittedArg):
                argtypes.append(types.Omitted(a.value))
            else:
                argtypes.append(self.typeof_pyval(a))

        return_val = None
        try:
            return_val = self.compile(tuple(argtypes))
        except errors.ForceLiteralArg as e:
            # Received request for compiler re-entry with the list of arguments
            # indicated by e.requested_args.
            # First, check if any of these args are already Literal-ized
            already_lit_pos = [i for i in e.requested_args
                               if isinstance(args[i], types.Literal)]
            if already_lit_pos:
                # Abort compilation if any argument is already a Literal.
                # Letting this continue will cause infinite compilation loop.
                m = ("Repeated literal typing request.\n"
                     "{}.\n"
                     "This is likely caused by an error in typing. "
                     "Please see nested and suppressed exceptions.")
                info = ', '.join('Arg #{} is {}'.format(i, args[i])
                                 for i in sorted(already_lit_pos))
                raise errors.CompilerError(m.format(info))
            # Convert requested arguments into a Literal.
            args = [(types.literal
                     if i in e.requested_args
                     else lambda x: x)(args[i])
                    for i, v in enumerate(args)]
            # Re-enter compilation with the Literal-ized arguments
            return_val = self._compile_for_args(*args)

        except errors.TypingError as e:
            # Intercept typing error that may be due to an argument
            # that failed inferencing as a Numba type
            failed_args = []
            for i, arg in enumerate(args):
                val = arg.value if isinstance(arg, OmittedArg) else arg
                try:
                    tp = typeof(val, Purpose.argument)
                except ValueError as typeof_exc:
                    failed_args.append((i, str(typeof_exc)))
                else:
                    if tp is None:
                        failed_args.append(
                            (i, f"cannot determine Numba type of value {val}"))
            if failed_args:
                # Patch error message to ease debugging
                args_str = "\n".join(
                    f"- argument {i}: {err}" for i, err in failed_args
                )
                msg = (f"{str(e).rstrip()} \n\nThis error may have been caused "
                       f"by the following argument(s):\n{args_str}\n")
                e.patch_message(msg)

            error_rewrite(e, 'typing')
        except errors.UnsupportedError as e:
            # Something unsupported is present in the user code, add help info
            error_rewrite(e, 'unsupported_error')
        except (errors.NotDefinedError, errors.RedefinedError,
                errors.VerificationError) as e:
            # These errors are probably from an issue with either the code
            # supplied being syntactically or otherwise invalid
            error_rewrite(e, 'interpreter')
        except errors.ConstantInferenceError as e:
            # this is from trying to infer something as constant when it isn't
            # or isn't supported as a constant
            error_rewrite(e, 'constant_inference')
        except Exception as e:
            if config.SHOW_HELP:
                if hasattr(e, 'patch_message'):
                    help_msg = errors.error_extras['reportable']
                    e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
            # ignore the FULL_TRACEBACKS config, this needs reporting!
            raise e
        finally:
            self._types_active_call = []
        return return_val

    def inspect_llvm(self, signature=None):
        """Get the LLVM intermediate representation generated by compilation.

        Parameters
        ----------
        signature : tuple of numba types, optional
            Specify a signature for which to obtain the LLVM IR. If None, the
            IR is returned for all available signatures.

        Returns
        -------
        llvm : dict[signature, str] or str
            Either the LLVM IR string for the specified signature, or, if no
            signature was given, a dictionary mapping signatures to LLVM IR
            strings.
        """
        if signature is not None:
            lib = self.overloads[signature].library
            return lib.get_llvm_str()

        return dict((sig, self.inspect_llvm(sig)) for sig in self.signatures)

    def inspect_asm(self, signature=None):
        """Get the generated assembly code.

        Parameters
        ----------
        signature : tuple of numba types, optional
            Specify a signature for which to obtain the assembly code. If
            None, the assembly code is returned for all available signatures.

        Returns
        -------
        asm : dict[signature, str] or str
            Either the assembly code for the specified signature, or, if no
            signature was given, a dictionary mapping signatures to assembly
            code.
        """
        if signature is not None:
            lib = self.overloads[signature].library
            return lib.get_asm_str()

        return dict((sig, self.inspect_asm(sig)) for sig in self.signatures)

    def inspect_types(self, file=None, signature=None,
                      pretty=False, style='default', **kwargs):
        """Print/return Numba intermediate representation (IR)-annotated code.

        Parameters
        ----------
        file : file-like object, optional
            File to which to print. Defaults to sys.stdout if None. Must be
            None if ``pretty=True``.
        signature : tuple of numba types, optional
            Print/return the intermediate representation for only the given
            signature. If None, the IR is printed for all available signatures.
        pretty : bool, optional
            If True, an Annotate object will be returned that can render the
            IR with color highlighting in Jupyter and IPython. ``file`` must
            be None if ``pretty`` is True. Additionally, the ``pygments``
            library must be installed for ``pretty=True``.
        style : str, optional
            Choose a style for rendering. Ignored if ``pretty`` is ``False``.
            This is directly consumed by ``pygments`` formatters. To see a
            list of available styles, import ``pygments`` and run
            ``list(pygments.styles.get_all_styles())``.

        Returns
        -------
        annotated : Annotate object, optional
            Only returned if ``pretty=True``, otherwise this function is only
            used for its printing side effect. If ``pretty=True``, an Annotate
            object is returned that can render itself in Jupyter and IPython.
        """
        overloads = self.overloads
        if signature is not None:
            overloads = {signature: self.overloads[signature]}

        if not pretty:
            if file is None:
                file = sys.stdout

            for ver, res in overloads.items():
                print("%s %s" % (self.py_func.__name__, ver), file=file)
                print('-' * 80, file=file)
                print(res.type_annotation, file=file)
                print('=' * 80, file=file)
        else:
            if file is not None:
                raise ValueError("`file` must be None if `pretty=True`")
            from numba.core.annotations.pretty_annotate import Annotate
            return Annotate(self, signature=signature, style=style)

    def inspect_cfg(self, signature=None, show_wrapper=None, **kwargs):
        """
        For inspecting the CFG of the function.

        By default the CFG of the user function is shown.  The *show_wrapper*
        option can be set to "python" or "cfunc" to show the python wrapper
        function or the *cfunc* wrapper function, respectively.

        Parameters accepted in kwargs
        -----------------------------
        filename : string, optional
            the name of the output file, if given this will write the output to
            filename
        view : bool, optional
            whether to immediately view the optional output file
        highlight : bool, set, dict, optional
            what, if anything, to highlight, options are:
            { incref : bool, # highlight NRT_incref calls
              decref : bool, # highlight NRT_decref calls
              returns : bool, # highlight exits which are normal returns
              raises : bool, # highlight exits which are from raise
              meminfo : bool, # highlight calls to NRT*meminfo
              branches : bool, # highlight true/false branches
             }
            Default is True which sets all of the above to True. Supplying a set
            of strings is also accepted, these are interpreted as key:True with
            respect to the above dictionary. e.g. {'incref', 'decref'} would
            switch on highlighting on increfs and decrefs.
        interleave: bool, set, dict, optional
            what, if anything, to interleave in the LLVM IR, options are:
            { python: bool # interleave python source code with the LLVM IR
              lineinfo: bool # interleave line information markers with the LLVM
                             # IR
            }
            Default is True which sets all of the above to True. Supplying a set
            of strings is also accepted, these are interpreted as key:True with
            respect to the above dictionary. e.g. {'python',} would
            switch on interleaving of python source code in the LLVM IR.
        strip_ir : bool, optional
            Default is False. If set to True all LLVM IR that is superfluous to
            that requested in kwarg `highlight` will be removed.
        show_key : bool, optional
            Default is True. Create a "key" for the highlighting in the rendered
            CFG.
        fontsize : int, optional
            Default is 8. Set the fontsize in the output to this value.
        """
        if signature is not None:
            cres = self.overloads[signature]
            lib = cres.library
            if show_wrapper == 'python':
                fname = cres.fndesc.llvm_cpython_wrapper_name
            elif show_wrapper == 'cfunc':
                fname = cres.fndesc.llvm_cfunc_wrapper_name
            else:
                fname = cres.fndesc.mangled_name
            return lib.get_function_cfg(fname, py_func=self.py_func, **kwargs)

        return dict((sig, self.inspect_cfg(sig, show_wrapper=show_wrapper))
                    for sig in self.signatures)

    def inspect_disasm_cfg(self, signature=None):
        """
        For inspecting the CFG of the disassembly of the function.

        Requires python package: r2pipe
        Requires radare2 binary on $PATH.
        Notebook rendering requires python package: graphviz

        signature : tuple of Numba types, optional
            Print/return the disassembly CFG for only the given signatures.
            If None, the IR is printed for all available signatures.
        """
        if signature is not None:
            cres = self.overloads[signature]
            lib = cres.library
            return lib.get_disasm_cfg(cres.fndesc.mangled_name)

        return dict((sig, self.inspect_disasm_cfg(sig))
                    for sig in self.signatures)

    def get_annotation_info(self, signature=None):
        """
        Gets the annotation information for the function specified by
        signature. If no signature is supplied a dictionary of signature to
        annotation information is returned.
        """
        signatures = self.signatures if signature is None else [signature]
        out = collections.OrderedDict()
        for sig in signatures:
            cres = self.overloads[sig]
            ta = cres.type_annotation
            key = (ta.func_id.filename + ':' + str(ta.func_id.firstlineno + 1),
                   ta.signature)
            out[key] = ta.annotate_raw()[key]
        return out

    def _explain_ambiguous(self, *args, **kws):
        """
        Callback for the C _Dispatcher object.
        """
        assert not kws, "kwargs not handled"
        args = tuple([self.typeof_pyval(a) for a in args])
        # The order here must be deterministic for testing purposes, which
        # is ensured by the OrderedDict.
        sigs = self.nopython_signatures
        # This will raise
        self.typingctx.resolve_overload(self.py_func, sigs, args, kws,
                                        allow_ambiguous=False)

    def _explain_matching_error(self, *args, **kws):
        """
        Callback for the C _Dispatcher object.
        """
        assert not kws, "kwargs not handled"
        args = [self.typeof_pyval(a) for a in args]
        msg = ("No matching definition for argument type(s) %s"
               % ', '.join(map(str, args)))
        raise TypeError(msg)

    def _search_new_conversions(self, *args, **kws):
        """
        Callback for the C _Dispatcher object.
        Search for approximately matching signatures for the given arguments,
        and ensure the corresponding conversions are registered in the C++
        type manager.
        """
        assert not kws, "kwargs not handled"
        args = [self.typeof_pyval(a) for a in args]
        found = False
        for sig in self.nopython_signatures:
            conv = self.typingctx.install_possible_conversions(args, sig.args)
            if conv:
                found = True
        return found

    def __repr__(self):
        return "%s(%s)" % (type(self).__name__, self.py_func)

    def typeof_pyval(self, val):
        """
        Resolve the Numba type of Python value *val*.
        This is called from numba._dispatcher as a fallback if the native code
        cannot decide the type.
        """
        # Not going through the resolve_argument_type() indirection
        # can save a couple µs.
        try:
            tp = typeof(val, Purpose.argument)
        except ValueError:
            tp = types.pyobject
        else:
            if tp is None:
                tp = types.pyobject
        self._types_active_call.append(tp)
        return tp

    def _callback_add_timer(self, duration, cres, lock_name):
        md = cres.metadata
        # md can be None when code is loaded from cache
        if md is not None:
            timers = md.setdefault("timers", {})
            if lock_name not in timers:
                # Only write if the metadata does not exist
                timers[lock_name] = duration
            else:
                msg = f"'{lock_name} metadata is already defined."
                raise AssertionError(msg)

    def _callback_add_compiler_timer(self, duration, cres):
        return self._callback_add_timer(duration, cres,
                                        lock_name="compiler_lock")

    def _callback_add_llvm_timer(self, duration, cres):
        return self._callback_add_timer(duration, cres,
                                        lock_name="llvm_lock")


class _MemoMixin:
    __uuid = None
    # A {uuid -> instance} mapping, for deserialization
    _memo = weakref.WeakValueDictionary()
    # hold refs to last N functions deserialized, retaining them in _memo
    # regardless of whether there is another reference
    _recent = collections.deque(maxlen=config.FUNCTION_CACHE_SIZE)

    @property
    def _uuid(self):
        """
        An instance-specific UUID, to avoid multiple deserializations of
        a given instance.

        Note: this is lazily-generated, for performance reasons.
        """
        u = self.__uuid
        if u is None:
            u = str(uuid.uuid4())
            self._set_uuid(u)
        return u

    def _set_uuid(self, u):
        assert self.__uuid is None
        self.__uuid = u
        self._memo[u] = self
        self._recent.append(self)


class Dispatcher(serialize.ReduceMixin, _MemoMixin, _DispatcherBase):
    """
    Implementation of user-facing dispatcher objects (i.e. created using
    the @jit decorator).
    This is an abstract base class. Subclasses should define the targetdescr
    class attribute.
    """
    _fold_args = True

    __numba__ = 'py_func'

    def __init__(self, py_func, locals={}, targetoptions={},
                 pipeline_class=compiler.Compiler):
        """
        Parameters
        ----------
        py_func: function object to be compiled
        locals: dict, optional
            Mapping of local variable names to Numba types.  Used to override
            the types deduced by the type inference engine.
        targetoptions: dict, optional
            Target-specific config options.
        pipeline_class: type numba.compiler.CompilerBase
            The compiler pipeline type.
        """
        self.typingctx = self.targetdescr.typing_context
        self.targetctx = self.targetdescr.target_context

        pysig = utils.pysignature(py_func)
        arg_count = len(pysig.parameters)
        can_fallback = not targetoptions.get('nopython', False)

        _DispatcherBase.__init__(self, arg_count, py_func, pysig, can_fallback,
                                 exact_match_required=False)

        functools.update_wrapper(self, py_func)

        self.targetoptions = targetoptions
        self.locals = locals
        self._cache = NullCache()
        compiler_class = _FunctionCompiler
        self._compiler = compiler_class(py_func, self.targetdescr,
                                        targetoptions, locals, pipeline_class)
        self._cache_hits = collections.Counter()
        self._cache_misses = collections.Counter()

        self._type = types.Dispatcher(self)
        self.typingctx.insert_global(self, self._type)

    def dump(self, tab=''):
        print(f'{tab}DUMP {type(self).__name__}[{self.py_func.__name__}'
              f', type code={self._type._code}]')
        for cres in self.overloads.values():
            cres.dump(tab=tab + '  ')
        print(f'{tab}END DUMP {type(self).__name__}[{self.py_func.__name__}]')

    @property
    def _numba_type_(self):
        return types.Dispatcher(self)

    def enable_caching(self):
        self._cache = FunctionCache(self.py_func)

    def __get__(self, obj, objtype=None):
        '''Allow a JIT function to be bound as a method to an object'''
        if obj is None:  # Unbound method
            return self
        else:  # Bound method
            return pytypes.MethodType(self, obj)

    def _reduce_states(self):
        """
        Reduce the instance for pickling.  This will serialize
        the original function as well the compilation options and
        compiled signatures, but not the compiled code itself.

        NOTE: part of ReduceMixin protocol
        """
        if self._can_compile:
            sigs = []
        else:
            sigs = [cr.signature for cr in self.overloads.values()]

        return dict(
            uuid=str(self._uuid),
            py_func=self.py_func,
            locals=self.locals,
            targetoptions=self.targetoptions,
            can_compile=self._can_compile,
            sigs=sigs,
        )

    @classmethod
    def _rebuild(cls, uuid, py_func, locals, targetoptions,
                 can_compile, sigs):
        """
        Rebuild an Dispatcher instance after it was __reduce__'d.

        NOTE: part of ReduceMixin protocol
        """
        try:
            return cls._memo[uuid]
        except KeyError:
            pass
        self = cls(py_func, locals, targetoptions)
        # Make sure this deserialization will be merged with subsequent ones
        self._set_uuid(uuid)
        for sig in sigs:
            self.compile(sig)
        self._can_compile = can_compile
        return self

    def compile(self, sig):
        with ExitStack() as scope:
            cres = None

            def cb_compiler(dur):
                if cres is not None:
                    self._callback_add_compiler_timer(dur, cres)

            def cb_llvm(dur):
                if cres is not None:
                    self._callback_add_llvm_timer(dur, cres)

            scope.enter_context(ev.install_timer("numba:compiler_lock",
                                                 cb_compiler))
            scope.enter_context(ev.install_timer("numba:llvm_lock", cb_llvm))
            scope.enter_context(global_compiler_lock)

            if not self._can_compile:
                raise RuntimeError("compilation disabled")
            # Use counter to track recursion compilation depth
            with self._compiling_counter:
                args, return_type = sigutils.normalize_signature(sig)
                # Don't recompile if signature already exists
                existing = self.overloads.get(tuple(args))
                if existing is not None:
                    return existing.entry_point
                # Try to load from disk cache
                cres = self._cache.load_overload(sig, self.targetctx)
                if cres is not None:
                    self._cache_hits[sig] += 1
                    # XXX fold this in add_overload()? (also see compiler.py)
                    if not cres.objectmode:
                        self.targetctx.insert_user_function(cres.entry_point,
                                                            cres.fndesc,
                                                            [cres.library])
                    self.add_overload(cres)
                    return cres.entry_point

                self._cache_misses[sig] += 1
                ev_details = dict(
                    dispatcher=self,
                    args=args,
                    return_type=return_type,
                )
                with ev.trigger_event("numba:compile", data=ev_details):
                    try:
                        cres = self._compiler.compile(args, return_type)
                    except errors.ForceLiteralArg as e:
                        def folded(args, kws):
                            return self._compiler.fold_argument_types(args,
                                                                      kws)[1]
                        raise e.bind_fold_arguments(folded)
                    self.add_overload(cres)
                self._cache.save_overload(sig, cres)
                return cres.entry_point

    def get_compile_result(self, sig):
        """Compile (if needed) and return the compilation result with the
        given signature.

        Returns ``CompileResult``.
        Raises ``NumbaError`` if the signature is incompatible.
        """
        atypes = tuple(sig.args)
        if atypes not in self.overloads:
            if self._can_compile:
                # Compiling may raise any NumbaError
                self.compile(atypes)
            else:
                msg = f"{sig} not available and compilation disabled"
                raise errors.TypingError(msg)
        return self.overloads[atypes]

    def recompile(self):
        """
        Recompile all signatures afresh.
        """
        sigs = list(self.overloads)
        old_can_compile = self._can_compile
        # Ensure the old overloads are disposed of,
        # including compiled functions.
        self._make_finalizer()()
        self._reset_overloads()
        self._cache.flush()
        self._can_compile = True
        try:
            for sig in sigs:
                self.compile(sig)
        finally:
            self._can_compile = old_can_compile

    @property
    def stats(self):
        return _CompileStats(
            cache_path=self._cache.cache_path,
            cache_hits=self._cache_hits,
            cache_misses=self._cache_misses,
        )

    def parallel_diagnostics(self, signature=None, level=1):
        """
        Print parallel diagnostic information for the given signature. If no
        signature is present it is printed for all known signatures. level is
        used to adjust the verbosity, level=1 (default) is minimal verbosity,
        and 2, 3, and 4 provide increasing levels of verbosity.
        """
        def dump(sig):
            ol = self.overloads[sig]
            pfdiag = ol.metadata.get('parfor_diagnostics', None)
            if pfdiag is None:
                msg = "No parfors diagnostic available, is 'parallel=True' set?"
                raise ValueError(msg)
            pfdiag.dump(level)
        if signature is not None:
            dump(signature)
        else:
            [dump(sig) for sig in self.signatures]

    def get_metadata(self, signature=None):
        """
        Obtain the compilation metadata for a given signature.
        """
        if signature is not None:
            return self.overloads[signature].metadata
        else:
            return dict(
                (sig,self.overloads[sig].metadata) for sig in self.signatures
            )

    def get_function_type(self):
        """Return unique function type of dispatcher when possible, otherwise
        return None.

        A Dispatcher instance has unique function type when it
        contains exactly one compilation result and its compilation
        has been disabled (via its disable_compile method).
        """
        if not self._can_compile and len(self.overloads) == 1:
            cres = tuple(self.overloads.values())[0]
            return types.FunctionType(cres.signature)


class LiftedCode(serialize.ReduceMixin, _MemoMixin, _DispatcherBase):
    """
    Implementation of the hidden dispatcher objects used for lifted code
    (a lifted loop is really compiled as a separate function).
    """
    _fold_args = False
    can_cache = False

    def __init__(self, func_ir, typingctx, targetctx, flags, locals):
        self.func_ir = func_ir
        self.lifted_from = None

        self.typingctx = typingctx
        self.targetctx = targetctx
        self.flags = flags
        self.locals = locals

        _DispatcherBase.__init__(self, self.func_ir.arg_count,
                                 self.func_ir.func_id.func,
                                 self.func_ir.func_id.pysig,
                                 can_fallback=True,
                                 exact_match_required=False)

    def _reduce_states(self):
        """
        Reduce the instance for pickling.  This will serialize
        the original function as well the compilation options and
        compiled signatures, but not the compiled code itself.

        NOTE: part of ReduceMixin protocol
        """
        return dict(
            uuid=self._uuid, func_ir=self.func_ir, flags=self.flags,
            locals=self.locals, extras=self._reduce_extras(),
        )

    def _reduce_extras(self):
        """
        NOTE: sub-class can override to add extra states
        """
        return {}

    @classmethod
    def _rebuild(cls, uuid, func_ir, flags, locals, extras):
        """
        Rebuild an Dispatcher instance after it was __reduce__'d.

        NOTE: part of ReduceMixin protocol
        """
        try:
            return cls._memo[uuid]
        except KeyError:
            pass

        # NOTE: We are assuming that this is must be cpu_target, which is true
        #       for now.
        # TODO: refactor this to not assume on `cpu_target`

        from numba.core import registry
        typingctx = registry.cpu_target.typing_context
        targetctx = registry.cpu_target.target_context

        self = cls(func_ir, typingctx, targetctx, flags, locals, **extras)
        self._set_uuid(uuid)
        return self

    def get_source_location(self):
        """Return the starting line number of the loop.
        """
        return self.func_ir.loc.line

    def _pre_compile(self, args, return_type, flags):
        """Pre-compile actions
        """
        pass

    @abstractmethod
    def compile(self, sig):
        """Lifted code should implement a compilation method that will return
        a CompileResult.entry_point for the given signature."""
        pass

    def _get_dispatcher_for_current_target(self):
        # Lifted code does not honor the target switch currently.
        # No work has been done to check if this can be allowed.
        return self


class LiftedLoop(LiftedCode):
    def _pre_compile(self, args, return_type, flags):
        assert not flags.enable_looplift, "Enable looplift flags is on"

    def compile(self, sig):
        with ExitStack() as scope:
            cres = None

            def cb_compiler(dur):
                if cres is not None:
                    self._callback_add_compiler_timer(dur, cres)

            def cb_llvm(dur):
                if cres is not None:
                    self._callback_add_llvm_timer(dur, cres)

            scope.enter_context(ev.install_timer("numba:compiler_lock",
                                                 cb_compiler))
            scope.enter_context(ev.install_timer("numba:llvm_lock", cb_llvm))
            scope.enter_context(global_compiler_lock)

            # Use counter to track recursion compilation depth
            with self._compiling_counter:
                # XXX this is mostly duplicated from Dispatcher.
                flags = self.flags
                args, return_type = sigutils.normalize_signature(sig)

                # Don't recompile if signature already exists
                # (e.g. if another thread compiled it before we got the lock)
                existing = self.overloads.get(tuple(args))
                if existing is not None:
                    return existing.entry_point

                self._pre_compile(args, return_type, flags)

                # copy the flags, use nopython first
                npm_loop_flags = flags.copy()
                npm_loop_flags.force_pyobject = False

                pyobject_loop_flags = flags.copy()
                pyobject_loop_flags.force_pyobject = True

                # Clone IR to avoid (some of the) mutation in the rewrite pass
                cloned_func_ir = self.func_ir.copy()

                ev_details = dict(
                    dispatcher=self,
                    args=args,
                    return_type=return_type,
                )
                with ev.trigger_event("numba:compile", data=ev_details):
                    # this emulates "object mode fall-back", try nopython, if it
                    # fails, then try again in object mode.
                    try:
                        cres = compiler.compile_ir(typingctx=self.typingctx,
                                                   targetctx=self.targetctx,
                                                   func_ir=cloned_func_ir,
                                                   args=args,
                                                   return_type=return_type,
                                                   flags=npm_loop_flags,
                                                   locals=self.locals,
                                                   lifted=(),
                                                   lifted_from=self.lifted_from,
                                                   is_lifted_loop=True,)
                    except errors.TypingError:
                        cres = compiler.compile_ir(typingctx=self.typingctx,
                                                   targetctx=self.targetctx,
                                                   func_ir=cloned_func_ir,
                                                   args=args,
                                                   return_type=return_type,
                                                   flags=pyobject_loop_flags,
                                                   locals=self.locals,
                                                   lifted=(),
                                                   lifted_from=self.lifted_from,
                                                   is_lifted_loop=True,)
                    # Check typing error if object mode is used
                    if (cres.typing_error is not None):
                        raise cres.typing_error
                    self.add_overload(cres)
                return cres.entry_point


class LiftedWith(LiftedCode):

    can_cache = True

    def _reduce_extras(self):
        return dict(output_types=self.output_types)

    @property
    def _numba_type_(self):
        return types.Dispatcher(self)

    def get_call_template(self, args, kws):
        """
        Get a typing.ConcreteTemplate for this dispatcher and the given
        *args* and *kws* types.  This enables the resolving of the return type.

        A (template, pysig, args, kws) tuple is returned.
        """
        # Ensure an overload is available
        if self._can_compile:
            self.compile(tuple(args))

        pysig = None
        # Create function type for typing
        func_name = self.py_func.__name__
        name = "CallTemplate({0})".format(func_name)
        # The `key` isn't really used except for diagnosis here,
        # so avoid keeping a reference to `cfunc`.
        call_template = typing.make_concrete_template(
            name, key=func_name, signatures=self.nopython_signatures)
        return call_template, pysig, args, kws

    def compile(self, sig):
        # this is similar to LiftedLoop's compile but does not have the
        # "fallback" to object mode part.
        with ExitStack() as scope:
            cres = None

            def cb_compiler(dur):
                if cres is not None:
                    self._callback_add_compiler_timer(dur, cres)

            def cb_llvm(dur):
                if cres is not None:
                    self._callback_add_llvm_timer(dur, cres)

            scope.enter_context(ev.install_timer("numba:compiler_lock",
                                                 cb_compiler))
            scope.enter_context(ev.install_timer("numba:llvm_lock", cb_llvm))
            scope.enter_context(global_compiler_lock)

            # Use counter to track recursion compilation depth
            with self._compiling_counter:
                # XXX this is mostly duplicated from Dispatcher.
                flags = self.flags
                args, return_type = sigutils.normalize_signature(sig)

                # Don't recompile if signature already exists
                # (e.g. if another thread compiled it before we got the lock)
                existing = self.overloads.get(tuple(args))
                if existing is not None:
                    return existing.entry_point

                self._pre_compile(args, return_type, flags)

                # Clone IR to avoid (some of the) mutation in the rewrite pass
                cloned_func_ir = self.func_ir.copy()

                ev_details = dict(
                    dispatcher=self,
                    args=args,
                    return_type=return_type,
                )
                with ev.trigger_event("numba:compile", data=ev_details):
                    cres = compiler.compile_ir(typingctx=self.typingctx,
                                               targetctx=self.targetctx,
                                               func_ir=cloned_func_ir,
                                               args=args,
                                               return_type=return_type,
                                               flags=flags, locals=self.locals,
                                               lifted=(),
                                               lifted_from=self.lifted_from,
                                               is_lifted_loop=True,)

                    # Check typing error if object mode is used
                    if (cres.typing_error is not None and
                            not flags.enable_pyobject):
                        raise cres.typing_error
                    self.add_overload(cres)
                return cres.entry_point


class ObjModeLiftedWith(LiftedWith):
    def __init__(self, *args, **kwargs):
        self.output_types = kwargs.pop('output_types', None)
        super(LiftedWith, self).__init__(*args, **kwargs)
        if not self.flags.force_pyobject:
            raise ValueError("expecting `flags.force_pyobject`")
        if self.output_types is None:
            raise TypeError('`output_types` must be provided')
        # switch off rewrites, they have no effect
        self.flags.no_rewrites = True

    @property
    def _numba_type_(self):
        return types.ObjModeDispatcher(self)

    def get_call_template(self, args, kws):
        """
        Get a typing.ConcreteTemplate for this dispatcher and the given
        *args* and *kws* types.  This enables the resolving of the return type.

        A (template, pysig, args, kws) tuple is returned.
        """
        assert not kws
        self._legalize_arg_types(args)
        # Coerce to object mode
        args = [types.ffi_forced_object] * len(args)

        if self._can_compile:
            self.compile(tuple(args))

        signatures = [typing.signature(self.output_types, *args)]
        pysig = None
        func_name = self.py_func.__name__
        name = "CallTemplate({0})".format(func_name)
        call_template = typing.make_concrete_template(
            name, key=func_name, signatures=signatures)

        return call_template, pysig, args, kws

    def _legalize_arg_types(self, args):
        for i, a in enumerate(args, start=1):
            if isinstance(a, types.List):
                msg = (
                    'Does not support list type inputs into '
                    'with-context for arg {}'
                )
                raise errors.TypingError(msg.format(i))
            elif isinstance(a, types.Dispatcher):
                msg = (
                    'Does not support function type inputs into '
                    'with-context for arg {}'
                )
                raise errors.TypingError(msg.format(i))

    @global_compiler_lock
    def compile(self, sig):
        args, _ = sigutils.normalize_signature(sig)
        sig = (types.ffi_forced_object,) * len(args)
        return super().compile(sig)


# Initialize typeof machinery
_dispatcher.typeof_init(
    OmittedArg,
    dict((str(t), t._code) for t in types.number_domain))