Stmt

def get_slice(self, name: str) -> slice:
    """Get the slice of the arguments.

    Args:
        name (str): The name of the slice.

    Returns:
        slice: The slice of the arguments.
    """
    index = self.node._name_args_slice[name]
    if isinstance(index, int):
        return slice(index, index + 1)
    return index

def insert(self, idx: int, value: SSAValue) -> None:
    """Insert the argument SSAValue at the specified index.

    Args:
        idx (int): The index to insert the value.
        value (SSAValue): The value to insert.
    """
    args = self.field
    value.add_use(Use(self.node, idx))
    new_args = (*args[:idx], value, *args[idx:])
    self.node._args = new_args
    self.field = new_args

def set_item(self, idx: int, value: SSAValue) -> None:
    """Set the argument SSAVAlue at the specified index.

    Args:
        idx (int): The index of the item to set.
        value (SSAValue): The value to set.
    """
    args = self.field
    args[idx].remove_use(Use(self.node, idx))
    value.add_use(Use(self.node, idx))
    new_args = (*args[:idx], value, *args[idx + 1 :])
    self.node._args = new_args
    self.field = new_args

def check(self) -> None:
    """Check the statement. Raises `Exception` if the statement is not correct.
    This method is called by the `verify` method, which will detect the source
    of the error in the IR. One should always call the `verify` method to verify
    the IR.

    The difference between `check` and `check_type` is that `check` is called
    at any time to check the structure of the IR by `verify`, while `check_type`
    is called after the type inference to check the types of the IR.
    """
    return

def check_type(self) -> None:
    """Check the types of the Block. Raises `Exception` if the types are not correct.
    This method is called by the `verify_type` method, which will detect the source
    of the error in the IR. One should always call the `verify_type` method to verify
    the types of the IR.

    Note:
        This method is generated by the `@statement` decorator. But can be overridden
        if needed.
    """
    raise NotImplementedError(
        "check_type should be implemented in the "
        "statement or generated by the @statement decorator"
    )

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

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

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

def detach(self) -> None:
    """detach the statement from its parent block."""
    if self.parent is None:
        return

    parent: Block = self.parent
    prev_stmt = self.prev_stmt
    next_stmt = self.next_stmt

    if prev_stmt is not None:
        prev_stmt._next_stmt = next_stmt
        self._prev_stmt = None
    else:
        assert (
            parent._first_stmt is self
        ), "Invalid statement, has no prev_stmt but not first_stmt"
        parent._first_stmt = next_stmt

    if next_stmt is not None:
        next_stmt._prev_stmt = prev_stmt
        self._next_stmt = None
    else:
        assert (
            parent._last_stmt is self
        ), "Invalid statement, has no next_stmt but not last_stmt"
        parent._last_stmt = prev_stmt

    self.parent = None
    parent._stmt_len -= 1
    return

def drop_all_references(self) -> None:
    """Remove all the dependency that reference/uses this Statement."""
    self.parent = None
    for idx, arg in enumerate(self._args):
        arg.remove_use(Use(self, idx))
    for region in self._regions:
        region.drop_all_references()

def expect_one_result(self) -> ResultValue:
    """Check if the statement contain only one result, and return it"""
    if len(self._results) != 1:
        raise ValueError(f"expected one result, got {len(self._results)}")
    return self._results[0]

@classmethod
def from_stmt(
    cls,
    other: Statement,
    args: Sequence[SSAValue] | None = None,
    regions: list[Region] | None = None,
    successors: list[Block] | None = None,
    attributes: dict[str, Attribute] | None = None,
) -> Self:
    """Create a similar Statement with new `ResultValue` and without
    attaching to any parent block. This still references to the old successor
    and regions.
    """
    obj = cls.__new__(cls)
    Statement.__init__(
        obj,
        args=args or other._args,
        regions=regions or other._regions,
        successors=successors or other.successors,
        attributes=attributes or other.attributes,
        result_types=[result.type for result in other._results],
        args_slice=other._name_args_slice,
    )
    # inherit the hint:
    for result, other_result in zip(obj._results, other._results):
        result.hints = dict(other_result.hints)

    return obj

def get_attribute(
    self, key: str, default: Attribute | None = None
) -> Attribute | None:
    """Get the attribute or property of the Statement.

    Args:
        key (str): The key of the attribute or property.

    Returns:
        Attribute | None: The attribute or property of the Statement.
    """
    return self.attributes.get(key, default)

@classmethod
def get_present_trait(cls, trait: type[TraitType]) -> TraitType:
    """Just like get_trait, but expects the trait to be there.
    Useful for linter checks, when you know the trait is present."""
    for t in cls.traits:
        if isinstance(t, trait):
            return t
    raise ValueError(f"Trait {trait} not present in statement {cls}")

@classmethod
def get_trait(cls, trait: type[TraitType]) -> TraitType | None:
    """Get the trait of the Statement."""
    for t in cls.traits:
        if isinstance(t, trait):
            return t
    return None

@classmethod
def has_trait(cls, trait_type: type[Trait["Statement"]]) -> bool:
    """Check if the Statement has a specific trait.

    Args:
        trait_type (type[StmtTrait]): The type of trait to check for.

    Returns:
        bool: True if the class has the specified trait, False otherwise.
    """
    for trait in cls.traits:
        if isinstance(trait, trait_type):
            return True
    return False

def insert_after(self, stmt: Statement) -> None:
    """Insert the current Statement after the input Statement.

    Args:
        stmt (Statement): Input Statement.

    Example:
        The following example demonstrates how to insert a Statement after another Statement.
        After `insert_after` is called, `stmt1` will be inserted after `stmt2`, which appears in IR in the order (stmt2 -> stmt1)
        ```python
        stmt1 = Statement()
        stmt2 = Statement()
        stmt1.insert_after(stmt2)
        ```
    """
    if self._next_stmt is not None and self._prev_stmt is not None:
        raise ValueError(
            f"Cannot insert before a statement that is already in a block: {self.name}"
        )

    if stmt._next_stmt is not None:
        stmt._next_stmt._prev_stmt = self

    self._prev_stmt = stmt
    self._next_stmt = stmt._next_stmt

    self.parent = stmt.parent
    stmt._next_stmt = self

    if self.parent:
        self.parent._stmt_len += 1

        if self._next_stmt is None:
            self.parent._last_stmt = self

def insert_before(self, stmt: Statement) -> None:
    """Insert the current Statement before the input Statement.

    Args:
        stmt (Statement): Input Statement.

    Example:
        The following example demonstrates how to insert a Statement before another Statement.
        After `insert_before` is called, `stmt1` will be inserted before `stmt2`, which appears in IR in the order (stmt1 -> stmt2)
        ```python
        stmt1 = Statement()
        stmt2 = Statement()
        stmt1.insert_before(stmt2)
        ```
    """
    if self._next_stmt is not None and self._prev_stmt is not None:
        raise ValueError(
            f"Cannot insert before a statement that is already in a block: {self.name}"
        )

    if stmt._prev_stmt is not None:
        stmt._prev_stmt._next_stmt = self

    self._next_stmt = stmt
    self._prev_stmt = stmt._prev_stmt

    self.parent = stmt.parent
    stmt._prev_stmt = self

    if self.parent:
        self.parent._stmt_len += 1

        if self._prev_stmt is None:
            self.parent._first_stmt = self

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

    Args:
        other (Self): The other Statelemt 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 IRNode is structurally equal to the other.
    """
    if context is None:
        context = {}

    if self.name != other.name:
        return False

    if (
        len(self.args) != len(other.args)
        or len(self.regions) != len(other.regions)
        or len(self.successors) != len(other.successors)
        or self.attributes != other.attributes
    ):
        return False

    if (
        self.parent is not None
        and other.parent is not None
        and context.get(self.parent) != other.parent
    ):
        return False

    if not all(
        context.get(arg, arg) == other_arg
        for arg, other_arg in zip(self.args, other.args)
    ):
        return False

    if not all(
        context.get(successor, successor) == other_successor
        for successor, other_successor in zip(self.successors, other.successors)
    ):
        return False

    if not all(
        region.is_structurally_equal(other_region, context)
        for region, other_region in zip(self.regions, other.regions)
    ):
        return False

    for result, other_result in zip(self._results, other._results):
        context[result] = other_result

    return True

def replace_by(self, stmt: Statement) -> None:
    """Replace the current Statement by the input Statement.

    Args:
        stmt (Statement): Input Statement.
    """
    stmt.insert_before(self)
    for result, old_result in zip(stmt._results, self._results):
        old_result.replace_by(result)
        if old_result.name:
            result.name = old_result.name
    self.delete()

def verify(self) -> None:
    try:
        self.check()
    except ValidationError as e:
        raise e
    except Exception as e:
        raise ValidationError(self, *e.args) from e

def verify_type(self) -> None:
    """Verify the type of the statement.

    Note:
        This API should be called after all the types are figured out (by typeinfer)
    """
    try:
        self.check_type()
    except TypeCheckError as e:
        raise e
    except Exception as e:
        raise TypeCheckError(self, *e.args) from e

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

    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.
        include_self (bool, optional): If the walk should include the Statement itself. Defaults to True.

    Yields:
        Iterator[Statement]: An iterator that yield Statements of Blocks in the Region, in the specified order.
    """
    if include_self and not region_first:
        yield self

    for region in reversed(self.regions) if reverse else self.regions:
        yield from region.walk(reverse=reverse, region_first=region_first)

    if include_self and region_first:
        yield self