from contextlib import contextmanager
from collections import defaultdict, namedtuple
from copy import copy
import warnings

from numba.core import (errors, types, typing, ir, funcdesc, rewrites,
                        typeinfer, config, lowering)

from numba.parfors.parfor import PreParforPass as _parfor_PreParforPass
from numba.parfors.parfor import ParforPass as _parfor_ParforPass
from numba.parfors.parfor import Parfor

from numba.core.compiler_machinery import (FunctionPass, LoweringPass,
                                           AnalysisPass, register_pass)
from numba.core.annotations import type_annotations
from numba.core.ir_utils import (raise_on_unsupported_feature, warn_deprecated,
                                 check_and_legalize_ir, guard,
                                 dead_code_elimination, simplify_CFG,
                                 get_definition,
                                 build_definitions, compute_cfg_from_blocks,
                                 is_operator_or_getitem)
from numba.core import postproc
from llvmlite import binding as llvm


# Outputs of type inference pass
_TypingResults = namedtuple("_TypingResults", [
    "typemap",
    "return_type",
    "calltypes",
    "typing_errors",
])


@contextmanager
def fallback_context(state, msg):
    """
    Wraps code that would signal a fallback to object mode
    """
    try:
        yield
    except Exception as e:
        if not state.status.can_fallback:
            raise
        else:
            # Clear all references attached to the traceback
            e = e.with_traceback(None)
            # this emits a warning containing the error message body in the
            # case of fallback from npm to objmode
            loop_lift = '' if state.flags.enable_looplift else 'OUT'
            msg_rewrite = ("\nCompilation is falling back to object mode "
                           "WITH%s looplifting enabled because %s"
                           % (loop_lift, msg))
            warnings.warn_explicit('%s due to: %s' % (msg_rewrite, e),
                                   errors.NumbaWarning,
                                   state.func_id.filename,
                                   state.func_id.firstlineno)
            raise


def type_inference_stage(typingctx, targetctx, interp, args, return_type,
                         locals={}, raise_errors=True):
    if len(args) != interp.arg_count:
        raise TypeError("Mismatch number of argument types")
    warnings = errors.WarningsFixer(errors.NumbaWarning)
    infer = typeinfer.TypeInferer(typingctx, interp, warnings)
    with typingctx.callstack.register(targetctx.target, infer, interp.func_id,
                                      args):
        # Seed argument types
        for index, (name, ty) in enumerate(zip(interp.arg_names, args)):
            infer.seed_argument(name, index, ty)

        # Seed return type
        if return_type is not None:
            infer.seed_return(return_type)

        # Seed local types
        for k, v in locals.items():
            infer.seed_type(k, v)

        infer.build_constraint()
        # return errors in case of partial typing
        errs = infer.propagate(raise_errors=raise_errors)
        typemap, restype, calltypes = infer.unify(raise_errors=raise_errors)

    # Output all Numba warnings
    warnings.flush()

    return _TypingResults(typemap, restype, calltypes, errs)


class BaseTypeInference(FunctionPass):
    _raise_errors = True

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        """
        Type inference and legalization
        """
        with fallback_context(state, 'Function "%s" failed type inference'
                              % (state.func_id.func_name,)):
            # Type inference
            typemap, return_type, calltypes, errs = type_inference_stage(
                state.typingctx,
                state.targetctx,
                state.func_ir,
                state.args,
                state.return_type,
                state.locals,
                raise_errors=self._raise_errors)
            state.typemap = typemap
            # save errors in case of partial typing
            state.typing_errors = errs
            if self._raise_errors:
                state.return_type = return_type
            state.calltypes = calltypes

        def legalize_return_type(return_type, interp, targetctx):
            """
            Only accept array return type iff it is passed into the function.
            Reject function object return types if in nopython mode.
            """
            if (not targetctx.enable_nrt and
                    isinstance(return_type, types.Array)):
                # Walk IR to discover all arguments and all return statements
                retstmts = []
                caststmts = {}
                argvars = set()
                for bid, blk in interp.blocks.items():
                    for inst in blk.body:
                        if isinstance(inst, ir.Return):
                            retstmts.append(inst.value.name)
                        elif isinstance(inst, ir.Assign):
                            if (isinstance(inst.value, ir.Expr)
                                    and inst.value.op == 'cast'):
                                caststmts[inst.target.name] = inst.value
                            elif isinstance(inst.value, ir.Arg):
                                argvars.add(inst.target.name)

                assert retstmts, "No return statements?"

                for var in retstmts:
                    cast = caststmts.get(var)
                    if cast is None or cast.value.name not in argvars:
                        if self._raise_errors:
                            msg = ("Only accept returning of array passed into "
                                   "the function as argument")
                            raise errors.NumbaTypeError(msg)

            elif (isinstance(return_type, types.Function) or
                    isinstance(return_type, types.Phantom)):
                if self._raise_errors:
                    msg = "Can't return function object ({}) in nopython mode"
                    raise errors.NumbaTypeError(msg.format(return_type))

        with fallback_context(state, 'Function "%s" has invalid return type'
                              % (state.func_id.func_name,)):
            legalize_return_type(state.return_type, state.func_ir,
                                 state.targetctx)
        return True


@register_pass(mutates_CFG=True, analysis_only=False)
class NopythonTypeInference(BaseTypeInference):
    _name = "nopython_type_inference"


@register_pass(mutates_CFG=True, analysis_only=False)
class PartialTypeInference(BaseTypeInference):
    _name = "partial_type_inference"
    _raise_errors = False


@register_pass(mutates_CFG=False, analysis_only=False)
class AnnotateTypes(AnalysisPass):
    _name = "annotate_types"

    def __init__(self):
        AnalysisPass.__init__(self)

    def get_analysis_usage(self, AU):
        AU.add_required(IRLegalization)

    def run_pass(self, state):
        """
        Create type annotation after type inference
        """
        func_ir = state.func_ir.copy()
        state.type_annotation = type_annotations.TypeAnnotation(
            func_ir=func_ir,
            typemap=state.typemap,
            calltypes=state.calltypes,
            lifted=state.lifted,
            lifted_from=state.lifted_from,
            args=state.args,
            return_type=state.return_type,
            html_output=config.HTML)

        if config.ANNOTATE:
            print("ANNOTATION".center(80, '-'))
            print(state.type_annotation)
            print('=' * 80)
        if config.HTML:
            with open(config.HTML, 'w') as fout:
                state.type_annotation.html_annotate(fout)

        return False


@register_pass(mutates_CFG=True, analysis_only=False)
class NopythonRewrites(FunctionPass):
    _name = "nopython_rewrites"

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        """
        Perform any intermediate representation rewrites after type
        inference.
        """
        # a bunch of these passes are either making assumptions or rely on some
        # very picky and slightly bizarre state particularly in relation to
        # ir.Del presence. To accommodate, ir.Dels are added ahead of running
        # this pass and stripped at the end.

        # Ensure we have an IR and type information.
        assert state.func_ir
        assert isinstance(getattr(state, 'typemap', None), dict)
        assert isinstance(getattr(state, 'calltypes', None), dict)
        msg = ('Internal error in post-inference rewriting '
               'pass encountered during compilation of '
               'function "%s"' % (state.func_id.func_name,))

        pp = postproc.PostProcessor(state.func_ir)
        pp.run(True)
        with fallback_context(state, msg):
            rewrites.rewrite_registry.apply('after-inference', state)
        pp.remove_dels()
        return True


@register_pass(mutates_CFG=True, analysis_only=False)
class PreParforPass(FunctionPass):

    _name = "pre_parfor_pass"

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        """
        Preprocessing for data-parallel computations.
        """
        # Ensure we have an IR and type information.
        assert state.func_ir
        preparfor_pass = _parfor_PreParforPass(
            state.func_ir,
            state.typemap,
            state.calltypes,
            state.typingctx,
            state.targetctx,
            state.flags.auto_parallel,
            state.parfor_diagnostics.replaced_fns
        )

        preparfor_pass.run()
        return True


# this is here so it pickles and for no other reason
def _reload_parfors():
    """Reloader for cached parfors
    """
    # Re-initialize the parallel backend when load from cache.
    from numba.np.ufunc.parallel import _launch_threads
    _launch_threads()


@register_pass(mutates_CFG=True, analysis_only=False)
class ParforPass(FunctionPass):

    _name = "parfor_pass"

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        """
        Convert data-parallel computations into Parfor nodes
        """
        # Ensure we have an IR and type information.
        assert state.func_ir
        parfor_pass = _parfor_ParforPass(state.func_ir,
                                         state.typemap,
                                         state.calltypes,
                                         state.return_type,
                                         state.typingctx,
                                         state.targetctx,
                                         state.flags.auto_parallel,
                                         state.flags,
                                         state.metadata,
                                         state.parfor_diagnostics)
        parfor_pass.run()

        # check the parfor pass worked and warn if it didn't
        has_parfor = False
        for blk in state.func_ir.blocks.values():
            for stmnt in blk.body:
                if isinstance(stmnt, Parfor):
                    has_parfor = True
                    break
            else:
                continue
            break

        if not has_parfor:
            # parfor calls the compiler chain again with a string
            if not (config.DISABLE_PERFORMANCE_WARNINGS or
                    state.func_ir.loc.filename == '<string>'):
                url = ("https://numba.readthedocs.io/en/stable/user/"
                       "parallel.html#diagnostics")
                msg = ("\nThe keyword argument 'parallel=True' was specified "
                       "but no transformation for parallel execution was "
                       "possible.\n\nTo find out why, try turning on parallel "
                       "diagnostics, see %s for help." % url)
                warnings.warn(errors.NumbaPerformanceWarning(msg,
                                                             state.func_ir.loc))

        # Add reload function to initialize the parallel backend.
        state.reload_init.append(_reload_parfors)
        return True


@register_pass(mutates_CFG=False, analysis_only=True)
class DumpParforDiagnostics(AnalysisPass):

    _name = "dump_parfor_diagnostics"

    def __init__(self):
        AnalysisPass.__init__(self)

    def run_pass(self, state):
        if state.flags.auto_parallel.enabled:
            if config.PARALLEL_DIAGNOSTICS:
                if state.parfor_diagnostics is not None:
                    state.parfor_diagnostics.dump(config.PARALLEL_DIAGNOSTICS)
                else:
                    raise RuntimeError("Diagnostics failed.")
        return True


@register_pass(mutates_CFG=True, analysis_only=False)
class NativeLowering(LoweringPass):

    _name = "native_lowering"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        if state.library is None:
            codegen = state.targetctx.codegen()
            state.library = codegen.create_library(state.func_id.func_qualname)
            # Enable object caching upfront, so that the library can
            # be later serialized.
            state.library.enable_object_caching()

        library = state.library
        targetctx = state.targetctx
        interp = state.func_ir  # why is it called this?!
        typemap = state.typemap
        restype = state.return_type
        calltypes = state.calltypes
        flags = state.flags
        metadata = state.metadata
        pre_stats = llvm.passmanagers.dump_refprune_stats()

        msg = ("Function %s failed at nopython "
               "mode lowering" % (state.func_id.func_name,))
        with fallback_context(state, msg):
            # Lowering
            fndesc = \
                funcdesc.PythonFunctionDescriptor.from_specialized_function(
                    interp, typemap, restype, calltypes,
                    mangler=targetctx.mangler, inline=flags.forceinline,
                    noalias=flags.noalias, abi_tags=[flags.get_mangle_string()])

            with targetctx.push_code_library(library):
                lower = lowering.Lower(targetctx, library, fndesc, interp,
                                       metadata=metadata)
                lower.lower()
                if not flags.no_cpython_wrapper:
                    lower.create_cpython_wrapper(flags.release_gil)

                if not flags.no_cfunc_wrapper:
                    # skip cfunc wrapper generation if unsupported
                    # argument or return types are used
                    for t in state.args:
                        if isinstance(t, (types.Omitted, types.Generator)):
                            break
                    else:
                        if isinstance(restype,
                                      (types.Optional, types.Generator)):
                            pass
                        else:
                            lower.create_cfunc_wrapper()

                env = lower.env
                call_helper = lower.call_helper
                del lower

            from numba.core.compiler import _LowerResult  # TODO: move this
            if flags.no_compile:
                state['cr'] = _LowerResult(fndesc, call_helper,
                                           cfunc=None, env=env)
            else:
                # Prepare for execution
                # Insert native function for use by other jitted-functions.
                # We also register its library to allow for inlining.
                cfunc = targetctx.get_executable(library, fndesc, env)
                targetctx.insert_user_function(cfunc, fndesc, [library])
                state['cr'] = _LowerResult(fndesc, call_helper,
                                           cfunc=cfunc, env=env)

            # capture pruning stats
            post_stats = llvm.passmanagers.dump_refprune_stats()
            metadata['prune_stats'] = post_stats - pre_stats

            # Save the LLVM pass timings
            metadata['llvm_pass_timings'] = library.recorded_timings
        return True


@register_pass(mutates_CFG=False, analysis_only=True)
class NoPythonSupportedFeatureValidation(AnalysisPass):
    """NoPython Mode check: Validates the IR to ensure that features in use are
    in a form that is supported"""

    _name = "nopython_supported_feature_validation"

    def __init__(self):
        AnalysisPass.__init__(self)

    def run_pass(self, state):
        raise_on_unsupported_feature(state.func_ir, state.typemap)
        warn_deprecated(state.func_ir, state.typemap)
        return False


@register_pass(mutates_CFG=False, analysis_only=True)
class IRLegalization(AnalysisPass):

    _name = "ir_legalization"

    def __init__(self):
        AnalysisPass.__init__(self)

    def run_pass(self, state):
        # NOTE: this function call must go last, it checks and fixes invalid IR!
        check_and_legalize_ir(state.func_ir, flags=state.flags)
        return True


@register_pass(mutates_CFG=True, analysis_only=False)
class NoPythonBackend(LoweringPass):

    _name = "nopython_backend"

    def __init__(self):
        LoweringPass.__init__(self)

    def run_pass(self, state):
        """
        Back-end: Generate LLVM IR from Numba IR, compile to machine code
        """
        lowered = state['cr']
        signature = typing.signature(state.return_type, *state.args)

        from numba.core.compiler import compile_result
        state.cr = compile_result(
            typing_context=state.typingctx,
            target_context=state.targetctx,
            entry_point=lowered.cfunc,
            typing_error=state.status.fail_reason,
            type_annotation=state.type_annotation,
            library=state.library,
            call_helper=lowered.call_helper,
            signature=signature,
            objectmode=False,
            lifted=state.lifted,
            fndesc=lowered.fndesc,
            environment=lowered.env,
            metadata=state.metadata,
            reload_init=state.reload_init,
        )
        return True


@register_pass(mutates_CFG=True, analysis_only=False)
class InlineOverloads(FunctionPass):
    """
    This pass will inline a function wrapped by the numba.extending.overload
    decorator directly into the site of its call depending on the value set in
    the 'inline' kwarg to the decorator.

    This is a typed pass. CFG simplification and DCE are performed on
    completion.
    """

    _name = "inline_overloads"

    def __init__(self):
        FunctionPass.__init__(self)

    _DEBUG = False

    def run_pass(self, state):
        """Run inlining of overloads
        """
        if self._DEBUG:
            print('before overload inline'.center(80, '-'))
            print(state.func_id.unique_name)
            print(state.func_ir.dump())
            print(''.center(80, '-'))
        from numba.core.inline_closurecall import (InlineWorker,
                                                   callee_ir_validator)
        inline_worker = InlineWorker(state.typingctx,
                                     state.targetctx,
                                     state.locals,
                                     state.pipeline,
                                     state.flags,
                                     callee_ir_validator,
                                     state.typemap,
                                     state.calltypes,
                                     )
        modified = False
        work_list = list(state.func_ir.blocks.items())
        # use a work list, look for call sites via `ir.Expr.op == call` and
        # then pass these to `self._do_work` to make decisions about inlining.
        while work_list:
            label, block = work_list.pop()
            for i, instr in enumerate(block.body):
                # TO-DO: other statements (setitem)
                if isinstance(instr, ir.Assign):
                    expr = instr.value
                    if isinstance(expr, ir.Expr):
                        workfn = self._do_work_expr

                        if guard(workfn, state, work_list, block, i, expr,
                                 inline_worker):
                            modified = True
                            break  # because block structure changed

        if self._DEBUG:
            print('after overload inline'.center(80, '-'))
            print(state.func_id.unique_name)
            print(state.func_ir.dump())
            print(''.center(80, '-'))

        if modified:
            # Remove dead blocks, this is safe as it relies on the CFG only.
            cfg = compute_cfg_from_blocks(state.func_ir.blocks)
            for dead in cfg.dead_nodes():
                del state.func_ir.blocks[dead]
            # clean up blocks
            dead_code_elimination(state.func_ir,
                                  typemap=state.typemap)
            # clean up unconditional branches that appear due to inlined
            # functions introducing blocks
            state.func_ir.blocks = simplify_CFG(state.func_ir.blocks)

        if self._DEBUG:
            print('after overload inline DCE'.center(80, '-'))
            print(state.func_id.unique_name)
            print(state.func_ir.dump())
            print(''.center(80, '-'))
        return True

    def _get_attr_info(self, state, expr):
        recv_type = state.typemap[expr.value.name]
        recv_type = types.unliteral(recv_type)
        matched = state.typingctx.find_matching_getattr_template(
            recv_type, expr.attr,
        )
        if not matched:
            return None

        template = matched['template']
        if getattr(template, 'is_method', False):
            # The attribute template is representing a method.
            # Don't inline the getattr.
            return None

        templates = [template]
        sig = typing.signature(matched['return_type'], recv_type)
        arg_typs = sig.args
        is_method = False

        return templates, sig, arg_typs, is_method

    def _get_callable_info(self, state, expr):

        def get_func_type(state, expr):
            func_ty = None
            if expr.op == 'call':
                # check this is a known and typed function
                try:
                    func_ty = state.typemap[expr.func.name]
                except KeyError:
                    # e.g. Calls to CUDA Intrinsic have no mapped type
                    # so KeyError
                    return None
                if not hasattr(func_ty, 'get_call_type'):
                    return None

            elif is_operator_or_getitem(expr):
                func_ty = state.typingctx.resolve_value_type(expr.fn)
            else:
                return None

            return func_ty

        if expr.op == 'call':
            # try and get a definition for the call, this isn't always
            # possible as it might be a eval(str)/part generated
            # awaiting update etc. (parfors)
            to_inline = None
            try:
                to_inline = state.func_ir.get_definition(expr.func)
            except Exception:
                return None

            # do not handle closure inlining here, another pass deals with that
            if getattr(to_inline, 'op', False) == 'make_function':
                return None

        func_ty = get_func_type(state, expr)
        if func_ty is None:
            return None

        sig = state.calltypes[expr]
        if not sig:
            return None

        templates, arg_typs, is_method = None, None, False
        if getattr(func_ty, 'template', None) is not None:
            # @overload_method
            is_method = True
            templates = [func_ty.template]
            arg_typs = (func_ty.template.this,) + sig.args
        else:
            # @overload case
            templates = getattr(func_ty, 'templates', None)
            arg_typs = sig.args

        return templates, sig, arg_typs, is_method

    def _do_work_expr(self, state, work_list, block, i, expr, inline_worker):

        def select_template(templates, args):
            if templates is None:
                return None

            impl = None
            for template in templates:
                inline_type = getattr(template, '_inline', None)
                if inline_type is None:
                    # inline not defined
                    continue
                if args not in template._inline_overloads:
                    # skip overloads not matching signature
                    continue
                if not inline_type.is_never_inline:
                    try:
                        impl = template._overload_func(*args)
                        if impl is None:
                            raise Exception  # abort for this template
                        break
                    except Exception:
                        continue
            else:
                return None

            return template, inline_type, impl

        inlinee_info = None
        if expr.op == 'getattr':
            inlinee_info = self._get_attr_info(state, expr)
        else:
            inlinee_info = self._get_callable_info(state, expr)

        if not inlinee_info:
            return False

        templates, sig, arg_typs, is_method = inlinee_info
        inlinee = select_template(templates, arg_typs)
        if inlinee is None:
            return False
        template, inlinee_type, impl = inlinee

        return self._run_inliner(
            state, inlinee_type, sig, template, arg_typs, expr, i, impl, block,
            work_list, is_method, inline_worker,
        )

    def _run_inliner(
        self, state, inline_type, sig, template, arg_typs, expr, i, impl, block,
        work_list, is_method, inline_worker,
    ):

        do_inline = True
        if not inline_type.is_always_inline:
            from numba.core.typing.templates import _inline_info
            caller_inline_info = _inline_info(state.func_ir,
                                              state.typemap,
                                              state.calltypes,
                                              sig)

            # must be a cost-model function, run the function
            iinfo = template._inline_overloads[arg_typs]['iinfo']
            if inline_type.has_cost_model:
                do_inline = inline_type.value(expr, caller_inline_info, iinfo)
            else:
                assert 'unreachable'

        if do_inline:
            if is_method:
                if not self._add_method_self_arg(state, expr):
                    return False
            arg_typs = template._inline_overloads[arg_typs]['folded_args']
            iinfo = template._inline_overloads[arg_typs]['iinfo']
            freevars = iinfo.func_ir.func_id.func.__code__.co_freevars
            _, _, _, new_blocks = inline_worker.inline_ir(state.func_ir,
                                                          block,
                                                          i,
                                                          iinfo.func_ir,
                                                          freevars,
                                                          arg_typs=arg_typs)
            if work_list is not None:
                for blk in new_blocks:
                    work_list.append(blk)
            return True
        else:
            return False

    def _add_method_self_arg(self, state, expr):
        func_def = guard(get_definition, state.func_ir, expr.func)
        if func_def is None:
            return False
        expr.args.insert(0, func_def.value)
        return True


@register_pass(mutates_CFG=False, analysis_only=False)
class DeadCodeElimination(FunctionPass):
    """
    Does dead code elimination
    """

    _name = "dead_code_elimination"

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        dead_code_elimination(state.func_ir, state.typemap)
        return True


@register_pass(mutates_CFG=False, analysis_only=False)
class PreLowerStripPhis(FunctionPass):
    """Remove phi nodes (ir.Expr.phi) introduced by SSA.

    This is needed before Lowering because the phi nodes in Numba IR do not
    match the semantics of phi nodes in LLVM IR. In Numba IR, phi nodes may
    expand into multiple LLVM instructions.
    """

    _name = "strip_phis"

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        state.func_ir = self._strip_phi_nodes(state.func_ir)
        state.func_ir._definitions = build_definitions(state.func_ir.blocks)
        # Rerun postprocessor to update metadata
        post_proc = postproc.PostProcessor(state.func_ir)
        post_proc.run(emit_dels=False)

        # Ensure we are not in objectmode generator
        if (state.func_ir.generator_info is not None
                and state.typemap is not None):
            # Rebuild generator type
            # TODO: move this into PostProcessor
            gentype = state.return_type
            state_vars = state.func_ir.generator_info.state_vars
            state_types = [state.typemap[k] for k in state_vars]
            state.return_type = types.Generator(
                gen_func=gentype.gen_func,
                yield_type=gentype.yield_type,
                arg_types=gentype.arg_types,
                state_types=state_types,
                has_finalizer=gentype.has_finalizer,
            )
        return True

    def _strip_phi_nodes(self, func_ir):
        """Strip Phi nodes from ``func_ir``

        For each phi node, put incoming value to their respective incoming
        basic-block at possibly the latest position (i.e. after the latest
        assignment to the corresponding variable).
        """
        exporters = defaultdict(list)
        phis = set()
        # Find all variables that needs to be exported
        for label, block in func_ir.blocks.items():
            for assign in block.find_insts(ir.Assign):
                if isinstance(assign.value, ir.Expr):
                    if assign.value.op == 'phi':
                        phis.add(assign)
                        phi = assign.value
                        for ib, iv in zip(phi.incoming_blocks,
                                          phi.incoming_values):
                            exporters[ib].append((assign.target, iv))

        # Rewrite the blocks with the new exporting assignments
        newblocks = {}
        for label, block in func_ir.blocks.items():
            newblk = copy(block)
            newblocks[label] = newblk

            # strip phis
            newblk.body = [stmt for stmt in block.body if stmt not in phis]

            # insert exporters
            for target, rhs in exporters[label]:
                # If RHS is undefined
                if rhs is ir.UNDEFINED:
                    # Put in a NULL initializer, set the location to be in what
                    # will eventually materialize as the prologue.
                    rhs = ir.Expr.null(loc=func_ir.loc)

                assign = ir.Assign(
                    target=target,
                    value=rhs,
                    loc=rhs.loc
                )
                # Insert at the earliest possible location; i.e. after the
                # last assignment to rhs
                assignments = [stmt for stmt in newblk.find_insts(ir.Assign)
                               if stmt.target == rhs]
                if assignments:
                    last_assignment = assignments[-1]
                    newblk.insert_after(assign, last_assignment)
                else:
                    newblk.prepend(assign)

        func_ir.blocks = newblocks
        return func_ir