# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

"""
Implement transformation on Numba IR
"""

from collections import namedtuple, defaultdict
import logging
import operator

from numba.cuda.core.analysis import (
    compute_cfg_from_blocks,
    find_top_level_loops,
)
from numba.cuda.core import ir
from numba.cuda.core import errors
from numba.cuda.core import ir_utils
from numba.cuda.core.analysis import compute_use_defs


_logger = logging.getLogger(__name__)


def _extract_loop_lifting_candidates(cfg, blocks):
    """
    Returns a list of loops that are candidate for loop lifting
    """

    # check well-formed-ness of the loop
    def same_exit_point(loop):
        "all exits must point to the same location"
        outedges = set()
        for k in loop.exits:
            succs = set(x for x, _ in cfg.successors(k))
            if not succs:
                # If the exit point has no successor, it contains an return
                # statement, which is not handled by the looplifting code.
                # Thus, this loop is not a candidate.
                _logger.debug("return-statement in loop.")
                return False
            outedges |= succs
        ok = len(outedges) == 1
        _logger.debug("same_exit_point=%s (%s)", ok, outedges)
        return ok

    def one_entry(loop):
        "there is one entry"
        ok = len(loop.entries) == 1
        _logger.debug("one_entry=%s", ok)
        return ok

    def cannot_yield(loop):
        "cannot have yield inside the loop"
        insiders = set(loop.body) | set(loop.entries) | set(loop.exits)
        for blk in map(blocks.__getitem__, insiders):
            for inst in blk.body:
                if isinstance(inst, ir.assign_types):
                    if isinstance(inst.value, ir.yield_types):
                        _logger.debug("has yield")
                        return False
        _logger.debug("no yield")
        return True

    _logger.info("finding looplift candidates")
    # the check for cfg.entry_point in the loop.entries is to prevent a bad
    # rewrite where a prelude for a lifted loop would get written into block -1
    # if a loop entry were in block 0
    candidates = []
    for loop in find_top_level_loops(cfg):
        _logger.debug("top-level loop: %s", loop)
        if (
            same_exit_point(loop)
            and one_entry(loop)
            and cannot_yield(loop)
            and cfg.entry_point() not in loop.entries
        ):
            candidates.append(loop)
            _logger.debug("add candidate: %s", loop)
    return candidates


def find_region_inout_vars(blocks, livemap, callfrom, returnto, body_block_ids):
    """Find input and output variables to a block region."""
    inputs = livemap[callfrom]
    outputs = livemap[returnto]

    # ensure live variables are actually used in the blocks, else remove,
    # saves having to create something valid to run through postproc
    # to achieve similar
    loopblocks = {}
    for k in body_block_ids:
        loopblocks[k] = blocks[k]

    used_vars = set()
    def_vars = set()
    defs = compute_use_defs(loopblocks)
    for vs in defs.usemap.values():
        used_vars |= vs
    for vs in defs.defmap.values():
        def_vars |= vs
    used_or_defined = used_vars | def_vars

    # note: sorted for stable ordering
    inputs = sorted(set(inputs) & used_or_defined)
    outputs = sorted(set(outputs) & used_or_defined & def_vars)
    return inputs, outputs


_loop_lift_info = namedtuple(
    "loop_lift_info", "loop,inputs,outputs,callfrom,returnto"
)


def _loop_lift_get_candidate_infos(cfg, blocks, livemap):
    """
    Returns information on looplifting candidates.
    """
    loops = _extract_loop_lifting_candidates(cfg, blocks)
    loopinfos = []
    for loop in loops:
        [callfrom] = loop.entries  # requirement checked earlier
        an_exit = next(iter(loop.exits))  # anyone of the exit block
        if len(loop.exits) > 1:
            # has multiple exits
            [(returnto, _)] = cfg.successors(
                an_exit
            )  # requirement checked earlier
        else:
            # does not have multiple exits
            returnto = an_exit

        local_block_ids = set(loop.body) | set(loop.entries) | set(loop.exits)
        inputs, outputs = find_region_inout_vars(
            blocks=blocks,
            livemap=livemap,
            callfrom=callfrom,
            returnto=returnto,
            body_block_ids=local_block_ids,
        )

        lli = _loop_lift_info(
            loop=loop,
            inputs=inputs,
            outputs=outputs,
            callfrom=callfrom,
            returnto=returnto,
        )
        loopinfos.append(lli)

    return loopinfos


def _loop_lift_modify_call_block(liftedloop, block, inputs, outputs, returnto):
    """
    Transform calling block from top-level function to call the lifted loop.
    """
    scope = block.scope
    loc = block.loc
    blk = ir.Block(scope=scope, loc=loc)

    ir_utils.fill_block_with_call(
        newblock=blk,
        callee=liftedloop,
        label_next=returnto,
        inputs=inputs,
        outputs=outputs,
    )
    return blk


def _loop_lift_prepare_loop_func(loopinfo, blocks):
    """
    Inplace transform loop blocks for use as lifted loop.
    """
    entry_block = blocks[loopinfo.callfrom]
    scope = entry_block.scope
    loc = entry_block.loc

    # Lowering assumes the first block to be the one with the smallest offset
    firstblk = min(blocks) - 1
    blocks[firstblk] = ir_utils.fill_callee_prologue(
        block=ir.Block(scope=scope, loc=loc),
        inputs=loopinfo.inputs,
        label_next=loopinfo.callfrom,
    )
    blocks[loopinfo.returnto] = ir_utils.fill_callee_epilogue(
        block=ir.Block(scope=scope, loc=loc),
        outputs=loopinfo.outputs,
    )


def _loop_lift_modify_blocks(
    func_ir, loopinfo, blocks, typingctx, targetctx, flags, locals
):
    """
    Modify the block inplace to call to the lifted-loop.
    Returns a dictionary of blocks of the lifted-loop.
    """
    from numba.cuda.dispatcher import LiftedLoop

    # Copy loop blocks
    loop = loopinfo.loop

    loopblockkeys = set(loop.body) | set(loop.entries)
    if len(loop.exits) > 1:
        # has multiple exits
        loopblockkeys |= loop.exits
    loopblocks = dict((k, blocks[k].copy()) for k in loopblockkeys)
    # Modify the loop blocks
    _loop_lift_prepare_loop_func(loopinfo, loopblocks)
    # Since Python 3.13, [END_FOR, POP_TOP] sequence becomes the start of the
    # block causing the block to have line number of the start of previous loop.
    # Fix this using the loc of the first getiter.
    getiter_exprs = []
    for blk in loopblocks.values():
        getiter_exprs.extend(blk.find_exprs(op="getiter"))
    first_getiter = min(getiter_exprs, key=lambda x: x.loc.line)
    loop_loc = first_getiter.loc
    # Create a new IR for the lifted loop
    lifted_ir = func_ir.derive(
        blocks=loopblocks,
        arg_names=tuple(loopinfo.inputs),
        arg_count=len(loopinfo.inputs),
        force_non_generator=True,
        loc=loop_loc,
    )
    liftedloop = LiftedLoop(lifted_ir, typingctx, targetctx, flags, locals)

    # modify for calling into liftedloop
    callblock = _loop_lift_modify_call_block(
        liftedloop,
        blocks[loopinfo.callfrom],
        loopinfo.inputs,
        loopinfo.outputs,
        loopinfo.returnto,
    )
    # remove blocks
    for k in loopblockkeys:
        del blocks[k]
    # update main interpreter callsite into the liftedloop
    blocks[loopinfo.callfrom] = callblock
    return liftedloop


def _has_multiple_loop_exits(cfg, lpinfo):
    """Returns True if there is more than one exit in the loop.

    NOTE: "common exits" refers to the situation where a loop exit has another
    loop exit as its successor. In that case, we do not need to alter it.
    """
    if len(lpinfo.exits) <= 1:
        return False
    exits = set(lpinfo.exits)
    pdom = cfg.post_dominators()

    # Eliminate blocks that have other blocks as post-dominators.
    processed = set()
    remain = set(exits)  # create a copy to work on
    while remain:
        node = remain.pop()
        processed.add(node)
        exits -= pdom[node] - {node}
        remain = exits - processed

    return len(exits) > 1


def _pre_looplift_transform(func_ir):
    """Canonicalize loops for looplifting."""
    from numba.cuda.core.postproc import PostProcessor

    cfg = compute_cfg_from_blocks(func_ir.blocks)
    # For every loop that has multiple exits, combine the exits into one.
    for loop_info in cfg.loops().values():
        if _has_multiple_loop_exits(cfg, loop_info):
            func_ir, _common_key = _fix_multi_exit_blocks(
                func_ir, loop_info.exits
            )
    # Reset and reprocess the func_ir
    func_ir._reset_analysis_variables()
    PostProcessor(func_ir).run()
    return func_ir


def loop_lifting(func_ir, typingctx, targetctx, flags, locals):
    """
    Loop lifting transformation.

    Given a interpreter `func_ir` returns a 2 tuple of
    `(toplevel_interp, [loop0_interp, loop1_interp, ....])`
    """
    func_ir = _pre_looplift_transform(func_ir)
    blocks = func_ir.blocks.copy()
    cfg = compute_cfg_from_blocks(blocks)
    loopinfos = _loop_lift_get_candidate_infos(
        cfg, blocks, func_ir.variable_lifetime.livemap
    )
    loops = []
    if loopinfos:
        _logger.debug(
            "loop lifting this IR with %d candidates:\n%s",
            len(loopinfos),
            func_ir.dump_to_string(),
        )
    for loopinfo in loopinfos:
        lifted = _loop_lift_modify_blocks(
            func_ir, loopinfo, blocks, typingctx, targetctx, flags, locals
        )
        loops.append(lifted)

    # Make main IR
    main = func_ir.derive(blocks=blocks)

    return main, loops


def canonicalize_cfg_single_backedge(blocks):
    """
    Rewrite loops that have multiple backedges.
    """
    cfg = compute_cfg_from_blocks(blocks)
    newblocks = blocks.copy()

    def new_block_id():
        return max(newblocks.keys()) + 1

    def has_multiple_backedges(loop):
        count = 0
        for k in loop.body:
            blk = blocks[k]
            edges = blk.terminator.get_targets()
            # is a backedge?
            if loop.header in edges:
                count += 1
                if count > 1:
                    # early exit
                    return True
        return False

    def yield_loops_with_multiple_backedges():
        for lp in cfg.loops().values():
            if has_multiple_backedges(lp):
                yield lp

    def replace_target(term, src, dst):
        def replace(target):
            return dst if target == src else target

        if isinstance(term, ir.branch_types):
            return ir.Branch(
                cond=term.cond,
                truebr=replace(term.truebr),
                falsebr=replace(term.falsebr),
                loc=term.loc,
            )
        elif isinstance(term, ir.jump_types):
            return ir.Jump(target=replace(term.target), loc=term.loc)
        else:
            assert not term.get_targets()
            return term

    def rewrite_single_backedge(loop):
        """
        Add new tail block that gathers all the backedges
        """
        header = loop.header
        tailkey = new_block_id()
        for blkkey in loop.body:
            blk = newblocks[blkkey]
            if header in blk.terminator.get_targets():
                newblk = blk.copy()
                # rewrite backedge into jumps to new tail block
                newblk.body[-1] = replace_target(
                    blk.terminator, header, tailkey
                )
                newblocks[blkkey] = newblk
        # create new tail block
        entryblk = newblocks[header]
        tailblk = ir.Block(scope=entryblk.scope, loc=entryblk.loc)
        # add backedge
        tailblk.append(ir.Jump(target=header, loc=tailblk.loc))
        newblocks[tailkey] = tailblk

    for loop in yield_loops_with_multiple_backedges():
        rewrite_single_backedge(loop)

    return newblocks


def canonicalize_cfg(blocks):
    """
    Rewrite the given blocks to canonicalize the CFG.
    Returns a new dictionary of blocks.
    """
    return canonicalize_cfg_single_backedge(blocks)


def with_lifting(func_ir, typingctx, targetctx, flags, locals):
    """With-lifting transformation

    Rewrite the IR to extract all withs.
    Only the top-level withs are extracted.
    Returns the (the_new_ir, the_lifted_with_ir)
    """
    from numba.cuda.core import postproc

    def dispatcher_factory(func_ir, objectmode=False, **kwargs):
        from numba.cuda.dispatcher import LiftedWith, ObjModeLiftedWith

        myflags = flags.copy()
        if objectmode:
            # Lifted with-block cannot looplift
            myflags.enable_looplift = False
            # Lifted with-block uses object mode
            myflags.enable_pyobject = True
            myflags.force_pyobject = True
            myflags.no_cpython_wrapper = False
            cls = ObjModeLiftedWith
        else:
            cls = LiftedWith
        return cls(func_ir, typingctx, targetctx, myflags, locals, **kwargs)

    # find where with-contexts regions are
    withs, func_ir = find_setupwiths(func_ir)

    if not withs:
        return func_ir, []

    postproc.PostProcessor(func_ir).run()  # ensure we have variable lifetime
    assert func_ir.variable_lifetime
    vlt = func_ir.variable_lifetime
    blocks = func_ir.blocks.copy()
    cfg = vlt.cfg
    # For each with-regions, mutate them according to
    # the kind of contextmanager
    sub_irs = []
    for blk_start, blk_end in withs:
        body_blocks = _cfg_nodes_in_region(cfg, blk_start, blk_end).copy()
        _legalize_with_head(blocks[blk_start])
        # Find the contextmanager
        cmkind, extra = _get_with_contextmanager(func_ir, blocks, blk_start)
        # Mutate the body and get new IR
        sub = cmkind.mutate_with_body(
            func_ir,
            blocks,
            blk_start,
            blk_end,
            body_blocks,
            dispatcher_factory,
            extra,
        )
        sub_irs.append(sub)
    if not sub_irs:
        # Unchanged
        new_ir = func_ir
    else:
        new_ir = func_ir.derive(blocks)
    return new_ir, sub_irs


def _get_with_contextmanager(func_ir, blocks, blk_start):
    """Get the global object used for the context manager"""
    _illegal_cm_msg = "Illegal use of context-manager."

    def get_var_dfn(var):
        """Get the definition given a variable"""
        return func_ir.get_definition(var)

    def get_ctxmgr_obj(var_ref):
        """Return the context-manager object and extra info.

        The extra contains the arguments if the context-manager is used
        as a call.
        """
        # If the contextmanager used as a Call
        dfn = func_ir.get_definition(var_ref)
        if isinstance(dfn, ir.expr_types) and dfn.op == "call":
            args = [get_var_dfn(x) for x in dfn.args]
            kws = {k: get_var_dfn(v) for k, v in dfn.kws}
            extra = {"args": args, "kwargs": kws}
            var_ref = dfn.func
        else:
            extra = None

        ctxobj = ir_utils.guard(ir_utils.find_outer_value, func_ir, var_ref)

        # check the contextmanager object
        if ctxobj is ir.UNDEFINED:
            raise errors.CompilerError(
                "Undefined variable used as context manager",
                loc=blocks[blk_start].loc,
            )

        if ctxobj is None:
            raise errors.CompilerError(_illegal_cm_msg, loc=dfn.loc)

        return ctxobj, extra

    # Scan the start of the with-region for the contextmanager
    for stmt in blocks[blk_start].body:
        if isinstance(stmt, ir.enterwith_types):
            var_ref = stmt.contextmanager
            ctxobj, extra = get_ctxmgr_obj(var_ref)
            if not hasattr(ctxobj, "mutate_with_body"):
                raise errors.CompilerError(
                    "Unsupported context manager in use",
                    loc=blocks[blk_start].loc,
                )
            return ctxobj, extra
    # No contextmanager found?
    raise errors.CompilerError(
        "malformed with-context usage",
        loc=blocks[blk_start].loc,
    )


def _legalize_with_head(blk):
    """Given *blk*, the head block of the with-context, check that it doesn't
    do anything else.
    """
    counters = defaultdict(int)
    for stmt in blk.body:
        # The counters dict is keyed on the IR node type. As the rest of the
        # function pops out specific node types, we normalize these node types
        # to be the Numba-CUDA versions of them so that we don't have to have
        # more complicated logic looking for both Numba and Numba-CUDA IR nodes
        # of the same kind.
        if isinstance(stmt, ir.enterwith_types):
            stmt_type = ir.EnterWith
        elif isinstance(stmt, ir.jump_types):
            stmt_type = ir.Jump
        elif isinstance(stmt, ir.del_types):
            stmt_type = ir.Del
        else:
            stmt_type = type(stmt)

        counters[stmt_type] += 1

    if counters.pop(ir.EnterWith) != 1:
        raise errors.CompilerError(
            "with's head-block must have exactly 1 ENTER_WITH",
            loc=blk.loc,
        )
    if counters.pop(ir.Jump, 0) != 1:
        raise errors.CompilerError(
            "with's head-block must have exactly 1 JUMP",
            loc=blk.loc,
        )
    # Can have any number of del
    counters.pop(ir.Del, None)
    # There MUST NOT be any other statements
    if counters:
        raise errors.CompilerError(
            "illegal statements in with's head-block",
            loc=blk.loc,
        )


def _cfg_nodes_in_region(cfg, region_begin, region_end):
    """Find the set of CFG nodes that are in the given region"""
    region_nodes = set()
    stack = [region_begin]
    while stack:
        tos = stack.pop()
        succlist = list(cfg.successors(tos))
        # a single block function will have a empty successor list
        if succlist:
            succs, _ = zip(*succlist)
            nodes = set(
                [
                    node
                    for node in succs
                    if node not in region_nodes and node != region_end
                ]
            )
            stack.extend(nodes)
            region_nodes |= nodes

    return region_nodes


def find_setupwiths(func_ir):
    """Find all top-level with.

    Returns a list of ranges for the with-regions.
    """

    def find_ranges(blocks):
        cfg = compute_cfg_from_blocks(blocks)
        sus_setups, sus_pops = set(), set()
        # traverse the cfg and collect all suspected SETUP_WITH and POP_BLOCK
        # statements so that we can iterate over them
        for label, block in blocks.items():
            for stmt in block.body:
                if ir_utils.is_setup_with(stmt):
                    sus_setups.add(label)
                if ir_utils.is_pop_block(stmt):
                    sus_pops.add(label)

        # now that we do have the statements, iterate through them in reverse
        # topo order and from each start looking for pop_blocks
        setup_with_to_pop_blocks_map = defaultdict(set)
        for setup_block in cfg.topo_sort(sus_setups, reverse=True):
            # begin pop_block, search
            to_visit, seen = [], []
            to_visit.append(setup_block)
            while to_visit:
                # get whatever is next and record that we have seen it
                block = to_visit.pop()
                seen.append(block)
                # go through the body of the block, looking for statements
                for stmt in blocks[block].body:
                    # raise detected before pop_block
                    if ir_utils.is_raise(stmt):
                        raise errors.CompilerError(
                            "unsupported control flow due to raise "
                            "statements inside with block"
                        )
                    # if a pop_block, process it
                    if ir_utils.is_pop_block(stmt) and block in sus_pops:
                        # record the jump target of this block belonging to this setup
                        setup_with_to_pop_blocks_map[setup_block].add(block)
                        # remove the block from blocks to be matched
                        sus_pops.remove(block)
                        # stop looking, we have reached the frontier
                        break
                    # if we are still here, by the block terminator,
                    # add all its targets to the to_visit stack, unless we
                    # have seen them already
                    if ir_utils.is_terminator(stmt):
                        to_visit.extend(
                            t for t in stmt.get_targets() if t not in seen
                        )

        return setup_with_to_pop_blocks_map

    blocks = func_ir.blocks
    # initial find, will return a dictionary, mapping indices of blocks
    # containing SETUP_WITH statements to a set of indices of blocks containing
    # POP_BLOCK statements
    with_ranges_dict = find_ranges(blocks)
    # rewrite the CFG in case there are multiple POP_BLOCK statements for one
    # with
    func_ir = consolidate_multi_exit_withs(with_ranges_dict, blocks, func_ir)
    # here we need to turn the withs back into a list of tuples so that the
    # rest of the code can cope
    with_ranges_tuple = [(s, list(p)[0]) for (s, p) in with_ranges_dict.items()]

    # check for POP_BLOCKS with multiple outgoing edges and reject
    for _, p in with_ranges_tuple:
        targets = blocks[p].terminator.get_targets()
        if len(targets) != 1:
            raise errors.CompilerError(
                "unsupported control flow: with-context contains branches "
                "(i.e. break/return/raise) that can leave the block "
            )
    # now we check for returns inside with and reject them
    for _, p in with_ranges_tuple:
        target_block = blocks[p]
        if ir_utils.is_return(
            func_ir.blocks[target_block.terminator.get_targets()[0]].terminator
        ):
            _rewrite_return(func_ir, p)

    # now we need to rewrite the tuple such that we have SETUP_WITH matching the
    # successor of the block that contains the POP_BLOCK.
    with_ranges_tuple = [
        (s, func_ir.blocks[p].terminator.get_targets()[0])
        for (s, p) in with_ranges_tuple
    ]

    # finally we check for nested with statements and reject them
    with_ranges_tuple = _eliminate_nested_withs(with_ranges_tuple)

    return with_ranges_tuple, func_ir


def _rewrite_return(func_ir, target_block_label):
    """Rewrite a return block inside a with statement.

    Arguments
    ---------

    func_ir: Function IR
      the CFG to transform
    target_block_label: int
      the block index/label of the block containing the POP_BLOCK statement


    This implements a CFG transformation to insert a block between two other
    blocks.

    The input situation is:

    ┌───────────────┐
    │   top         │
    │   POP_BLOCK   │
    │   bottom      │
    └───────┬───────┘
            │
    ┌───────▼───────┐
    │               │
    │    RETURN     │
    │               │
    └───────────────┘

    If such a pattern is detected in IR, it means there is a `return` statement
    within a `with` context. The basic idea is to rewrite the CFG as follows:

    ┌───────────────┐
    │   top         │
    │   POP_BLOCK   │
    │               │
    └───────┬───────┘
            │
    ┌───────▼───────┐
    │               │
    │     bottom    │
    │               │
    └───────┬───────┘
            │
    ┌───────▼───────┐
    │               │
    │    RETURN     │
    │               │
    └───────────────┘

    We split the block that contains the `POP_BLOCK` statement into two blocks.
    Everything from the beginning of the block up to and including the
    `POP_BLOCK` statement is considered the 'top' and everything below is
    considered 'bottom'. Finally the jump statements are re-wired to make sure
    the CFG remains valid.

    """
    # the block itself from the index
    target_block = func_ir.blocks[target_block_label]
    # get the index of the block containing the return
    target_block_successor_label = target_block.terminator.get_targets()[0]
    # the return block
    target_block_successor = func_ir.blocks[target_block_successor_label]

    # create the new return block with an appropriate label
    max_label = ir_utils.find_max_label(func_ir.blocks)
    new_label = max_label + 1
    # create the new return block
    new_block_loc = target_block_successor.loc
    new_block_scope = ir.Scope(None, loc=new_block_loc)
    new_block = ir.Block(new_block_scope, loc=new_block_loc)

    # Split the block containing the POP_BLOCK into top and bottom
    # Block must be of the form:
    # -----------------
    # <some stmts>
    # POP_BLOCK
    # <some more stmts>
    # JUMP
    # -----------------
    top_body, bottom_body = [], []
    pop_blocks = [*target_block.find_insts(ir.popblock_types)]
    assert len(pop_blocks) == 1
    assert len([*target_block.find_insts(ir.jump_types)]) == 1
    assert isinstance(target_block.body[-1], ir.jump_types)
    pb_marker = pop_blocks[0]
    pb_is = target_block.body.index(pb_marker)
    top_body.extend(target_block.body[:pb_is])
    top_body.append(ir.Jump(target_block_successor_label, target_block.loc))
    bottom_body.extend(target_block.body[pb_is:-1])
    bottom_body.append(ir.Jump(new_label, target_block.loc))

    # get the contents of the return block
    return_body = func_ir.blocks[target_block_successor_label].body
    # finally, re-assign all blocks
    new_block.body.extend(return_body)
    target_block_successor.body.clear()
    target_block_successor.body.extend(bottom_body)
    target_block.body.clear()
    target_block.body.extend(top_body)

    # finally, append the new return block and rebuild the IR properties
    func_ir.blocks[new_label] = new_block
    func_ir._definitions = ir_utils.build_definitions(func_ir.blocks)
    return func_ir


def _eliminate_nested_withs(with_ranges):
    known_ranges = []

    def within_known_range(start, end, known_ranges):
        for a, b in known_ranges:
            # FIXME: this should be a comparison in topological order, right
            # now we are comparing the integers of the blocks, stuff probably
            # works by accident.
            if start > a and end < b:
                return True
        return False

    for s, e in sorted(with_ranges):
        if not within_known_range(s, e, known_ranges):
            known_ranges.append((s, e))

    return known_ranges


def consolidate_multi_exit_withs(withs: dict, blocks, func_ir):
    """Modify the FunctionIR to merge the exit blocks of with constructs."""
    for k in withs:
        vs: set = withs[k]
        if len(vs) > 1:
            func_ir, common = _fix_multi_exit_blocks(
                func_ir,
                vs,
                split_condition=ir_utils.is_pop_block,
            )
            withs[k] = {common}
    return func_ir


def _fix_multi_exit_blocks(func_ir, exit_nodes, *, split_condition=None):
    """Modify the FunctionIR to create a single common exit node given the
    original exit nodes.

    Parameters
    ----------
    func_ir :
        The FunctionIR. Mutated inplace.
    exit_nodes :
        The original exit nodes. A sequence of block keys.
    split_condition : callable or None
        If not None, it is a callable with the signature
        `split_condition(statement)` that determines if the `statement` is the
        splitting point (e.g. `POP_BLOCK`) in an exit node.
        If it's None, the exit node is not split.
    """

    # Convert the following:
    #
    #     |           |
    # +-------+   +-------+
    # | exit0 |   | exit1 |
    # +-------+   +-------+
    #     |           |
    # +-------+   +-------+
    # | after0|   | after1|
    # +-------+   +-------+
    #     |           |
    #
    # To roughly:
    #
    #     |           |
    # +-------+   +-------+
    # | exit0 |   | exit1 |
    # +-------+   +-------+
    #     |           |
    #     +-----+-----+
    #           |
    #      +---------+
    #      | common  |
    #      +---------+
    #           |
    #       +-------+
    #       | post  |
    #       +-------+
    #           |
    #     +-----+-----+
    #     |           |
    # +-------+   +-------+
    # | after0|   | after1|
    # +-------+   +-------+

    blocks = func_ir.blocks
    # Getting the scope
    any_blk = min(func_ir.blocks.values())
    scope = any_blk.scope
    # Getting the maximum block label
    max_label = max(func_ir.blocks) + 1
    # Define the new common block for the new exit.
    common_block = ir.Block(any_blk.scope, loc=ir.unknown_loc)
    common_label = max_label
    max_label += 1
    blocks[common_label] = common_block
    # Define the new block after the exit.
    post_block = ir.Block(any_blk.scope, loc=ir.unknown_loc)
    post_label = max_label
    max_label += 1
    blocks[post_label] = post_block

    # Adjust each exit node
    remainings = []
    for i, k in enumerate(exit_nodes):
        blk = blocks[k]

        # split the block if needed
        if split_condition is not None:
            for stmt in blk.body:
                if split_condition(stmt):
                    break
        else:
            # no splitting
            pt = -1

        before = blk.body[:pt]
        after = blk.body[pt:]
        remainings.append(after)

        # Add control-point variable to mark which exit block this is.
        blk.body = before
        loc = blk.loc
        blk.body.append(
            ir.Assign(
                value=ir.Const(i, loc=loc),
                target=scope.get_or_define("$cp", loc=loc),
                loc=loc,
            )
        )
        # Replace terminator with a jump to the common block
        assert not blk.is_terminated
        blk.body.append(ir.Jump(common_label, loc=ir.unknown_loc))

    if split_condition is not None:
        # Move the splitting statement to the common block
        common_block.body.append(remainings[0][0])
    assert not common_block.is_terminated
    # Append jump from common block to post block
    common_block.body.append(ir.Jump(post_label, loc=loc))

    # Make if-else tree to jump to target
    remain_blocks = range(max_label, max_label + len(remainings))

    switch_block = post_block
    loc = ir.unknown_loc
    for i, remain in enumerate(remainings):
        match_expr = scope.redefine("$cp_check", loc=loc)
        match_rhs = scope.redefine("$cp_rhs", loc=loc)

        # Do comparison to match control-point variable to the exit block
        switch_block.body.append(
            ir.Assign(value=ir.Const(i, loc=loc), target=match_rhs, loc=loc),
        )

        # Add assignment for the comparison
        switch_block.body.append(
            ir.Assign(
                value=ir.Expr.binop(
                    fn=operator.eq,
                    lhs=scope.get("$cp"),
                    rhs=match_rhs,
                    loc=loc,
                ),
                target=match_expr,
                loc=loc,
            ),
        )

        # Insert jump to the next case
        [jump_target] = remain[-1].get_targets()
        switch_block.body.append(
            ir.Branch(match_expr, jump_target, remain_blocks[i], loc=loc),
        )
        switch_block = ir.Block(scope=scope, loc=loc)
        blocks[remain_blocks[i]] = switch_block

    # Add the final jump
    switch_block.body.append(ir.Jump(jump_target, loc=loc))

    return func_ir, common_label
