Region

def __getitem__(self, block: Block) -> int:
    """Get the index of a block within the region.

    Args:
        block (Block): The block to get the index of.

    Raises:
        ValueError: If the block does not belong to the region.

    Returns:
        int: The index of the block within the region.
    """
    if block.parent is not self:
        raise ValueError("Block does not belong to the region")
    return self._block_idx[block]

def clone(self, ssamap: dict[SSAValue, SSAValue] | None = None) -> Region:
    """Clone a region. This will clone all blocks and statements in the region.
    `SSAValue` defined outside the region will not be cloned unless provided in `ssamap`.

    Note:
        Blocks must be in definition order (e.g. reverse post-order of the CFG)
        so that every statement result is cloned before it is referenced.
        Use :class:`kirin.rewrite.SortBlocks` to ensure this after passes
        that may reorder blocks.
    """
    ret = Region()
    successor_map: dict[Block, Block] = {}
    _ssamap = ssamap or {}

    # Collect all SSA values defined inside this region (statement results).
    # Block args are handled separately below.
    in_region_defs: set[SSAValue] = set()
    for block in self.blocks:
        for stmt in block.stmts:
            in_region_defs.update(stmt.results)

    # First pass: create cloned blocks and block args (order doesn't matter).
    for block in self.blocks:
        new_block = Block()
        ret.blocks.append(new_block)
        successor_map[block] = new_block
        for arg in block.args:
            new_arg = new_block.args.append_from(arg.type, arg.name)
            _ssamap[arg] = new_arg

    def _map_arg(arg: SSAValue) -> SSAValue:
        if arg in _ssamap:
            return _ssamap[arg]
        if arg in in_region_defs:
            raise ValueError(
                f"Region.clone: in-region SSA value {arg!r} used before "
                f"its defining statement was cloned. This indicates a "
                f"block ordering issue — run SortBlocks before cloning."
            )
        # Value defined outside the region — keep as-is.
        return arg

    # Second pass: clone statements in block-list order.
    for block in self.blocks:
        for stmt in block.stmts:
            new_stmt = stmt.from_stmt(
                stmt,
                args=[_map_arg(arg) for arg in stmt.args],
                regions=[region.clone(_ssamap) for region in stmt.regions],
                successors=[
                    successor_map[successor] for successor in stmt.successors
                ],
            )
            successor_map[block].stmts.append(new_stmt)
            for result, new_result in zip(stmt.results, new_stmt.results):
                _ssamap[result] = new_result
                new_result.name = result.name

    return ret

def delete(self, safe: bool = True) -> None:
    """Delete the Region completely from the IR graph.

    Note:
        This method will detach + remove references of the Region.

    Args:
        safe (bool, optional): If True, raise error if there is anything that still reference components in the Region. Defaults to True.
    """
    self.detach()
    self.drop_all_references()

def detach(self, index: int | None = None) -> None:
    """Detach this Region from the IR tree graph.

    Note:
        Detach only detach the Region from the IR graph. It does not remove uses that reference the Region.
    """
    # already detached
    if self.parent_node is None:
        return

    if index is not None:
        region_idx = index
    else:
        region_idx = self.region_index

    del self.parent_node._regions[region_idx]
    self.parent_node = None

def drop_all_references(self) -> None:
    """Remove all the dependency that reference/uses this Region."""
    self.parent_node = None
    for block in self._blocks:
        block.drop_all_references()

def is_structurally_equal(
    self,
    other: IRNode,
    context: dict[IRNode | SSAValue, IRNode | SSAValue] | None = None,
) -> bool:
    """Check if the Region is structurally equal to another Region.

    Args:
        other (IRNode): The other node to compare with.
        context (dict[IRNode  |  SSAValue, IRNode  |  SSAValue] | None, optional): A map of IRNode/SSAValue to hint that they are equivalent so the check will treat them as equivalent. Defaults to None.

    Returns:
        bool: True if the Region is structurally equal to the other Region.
    """
    if not isinstance(other, Region):
        return False

    if context is None:
        context = {}
    context[self] = other
    if len(self.blocks) != len(other.blocks):
        return False

    for block, other_block in zip(self.blocks, other.blocks):
        context[block] = other_block

    if not all(
        block.is_structurally_equal(other_block, context)
        for block, other_block in zip(self.blocks, other.blocks)
    ):
        return False

    return True

def stmts(self) -> Iterator[Statement]:
    """Iterate over all the Statements in the Region. This does not walk into nested Regions.

    Yields:
        Iterator[Statement]: An iterator that yield Statements of Blocks in the Region.
    """
    for block in self.blocks:
        yield from block.stmts

def verify(self) -> None:
    """Verify the correctness of the Region.

    Raises:
        ValidationError: If the Region is not correct.
    """
    from kirin.ir.nodes.stmt import Statement

    if not isinstance(self.parent_node, Statement):
        raise ValidationError(
            self, "expect Region to have a parent of type Statement"
        )

    for block in self.blocks:
        block.verify()

def verify_type(self) -> None:
    for block in self.blocks:
        block.verify_type()

def walk(
    self, *, reverse: bool = False, region_first: bool = False
) -> Iterator[Statement]:
    """Traversal the Statements of Blocks in the Region.

    Args:
        reverse (bool, optional): If walk in the reversed manner. Defaults to False.
        region_first (bool, optional): If the walk should go through the Statement first or the Region of a Statement first. Defaults to False.

    Yields:
        Iterator[Statement]: An iterator that yield Statements of Blocks in the Region, in the specified order.
    """
    for block in reversed(self.blocks) if reverse else self.blocks:
        yield from block.walk(reverse=reverse, region_first=region_first)

def __setitem__(  # ty: ignore[invalid-method-override]  # generic TypeVar resolution
    self, idx: int | slice, block_or_blocks: Block | Sequence[Block]
) -> None:
    """Replace/Set the Blocks of the Region.

    Args:
        idx (int | slice): The index or slice to replace the [`Blocks`][kirin.ir.Block].
        block_or_blocks (Block | Sequence[Block]): The Block or Blocks to replace the Blocks.

    """
    if isinstance(idx, int) and isinstance(block_or_blocks, Block):
        self.field[idx].detach()
        block_or_blocks.attach(self.node)
        self.field[idx] = block_or_blocks
        self.node._block_idx[block_or_blocks] = idx
    elif isinstance(idx, slice) and not isinstance(block_or_blocks, Block):
        blocks = list(block_or_blocks)
        for block in blocks:
            block.attach(self.node)
        self.field[idx] = blocks
        self.node._block_idx = {
            block: i for i, block in enumerate(self.field)
        }  # reindex
    else:
        raise ValueError("Invalid assignment")

def append(self, value: Block) -> None:
    """Append a Block to the Region.

    Args:
        value (Block): The block to be appended.
    """
    value.attach(self.node)
    self.node._block_idx[value] = len(self.field)
    self.field.append(value)

def insert(self, idx: int, value: Block) -> None:
    """Inserts a Block at the specified index.

    Args:
        idx (int): The index at which to insert the block.
        value (Block): The block to be inserted.
    """
    value.attach(self.node)
    self.field.insert(idx, value)
    for i, value in enumerate(self.field[idx:], idx):
        self.node._block_idx[value] = i