Pdl interp

@irdl_op_definition
class GetOperandOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpget_operand-pdl_interpgetoperandop).
    """

    name = "pdl_interp.get_operand"
    index = prop_def(IntegerAttr[I32])
    input_op = operand_def(OperationType)
    value = result_def(ValueType)

    assembly_format = "$index `of` $input_op attr-dict"

    def __init__(self, index: int | IntegerAttr[I32], input_op: SSAValue) -> None:
        if isinstance(index, int):
            index = IntegerAttr(index, i32)
        super().__init__(
            operands=[input_op], properties={"index": index}, result_types=[ValueType()]
        )

@irdl_op_definition
class GetOperandsOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpget_operands-pdl_interpgetoperandsop).
    """

    name = "pdl_interp.get_operands"
    index = opt_prop_def(IntegerAttr[I32])
    input_op = operand_def(OperationType)
    value = result_def(ValueType | RangeType[ValueType])

    # TODO: assembly format doesn't work due to a bug:
    # https://github.com/xdslproject/xdsl/issues/5562
    # assembly_format = "($index^)? `of` $input_op `:` type($value) attr-dict"

    def __init__(
        self,
        index: int | IntegerAttr[I32] | None,
        input_op: SSAValue,
        result_type: ValueType | RangeType[ValueType],
    ) -> None:
        if isinstance(index, int):
            index = IntegerAttr(index, i32)
        super().__init__(
            operands=[input_op],
            properties={"index": index},
            result_types=[result_type],
        )

    @classmethod
    def parse(cls, parser: Parser) -> GetOperandsOp:
        index = parser.parse_optional_integer()
        if index is not None:
            index = IntegerAttr(index, 32)
        parser.parse_characters("of")
        input_op = parser.parse_operand(OperationType())
        parser.parse_punctuation(":")
        result_type = parser.parse_type()
        return GetOperandsOp.build(
            operands=(input_op,),
            properties={"index": index},
            result_types=(result_type,),
        )

    def print(self, printer: Printer):
        if self.index is not None:
            printer.print_string(" ", indent=0)
            self.index.print_without_type(printer)
        printer.print_string(" of ", indent=0)
        printer.print_operand(self.input_op)
        printer.print_string(" : ", indent=0)
        printer.print_attribute(self.value.type)

@irdl_op_definition
class FinalizeOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpfinalize-pdl_interpfinalizeop).
    """

    name = "pdl_interp.finalize"
    traits = traits_def(IsTerminator())

    assembly_format = "attr-dict"

    def __init__(self):
        super().__init__()

@irdl_op_definition
class CheckOperationNameOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpcheck_operation_name-pdl_interpcheckoperationnameop).
    """

    name = "pdl_interp.check_operation_name"
    traits = traits_def(IsTerminator())
    operation_name = prop_def(StringAttr, prop_name="name")
    input_op = operand_def(OperationType)
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = (
        "`of` $input_op `is` $name attr-dict `->` $true_dest `, ` $false_dest"
    )

    def __init__(
        self,
        operation_name: str | StringAttr,
        input_op: SSAValue,
        trueDest: Block,
        falseDest: Block,
    ) -> None:
        if isinstance(operation_name, str):
            operation_name = StringAttr(operation_name)
        super().__init__(
            operands=[input_op],
            properties={"name": operation_name},
            successors=[trueDest, falseDest],
        )

@irdl_op_definition
class CheckOperandCountOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpcheck_operand_count-pdl_interpcheckoperandcountop).
    """

    name = "pdl_interp.check_operand_count"
    traits = traits_def(IsTerminator())
    input_op = operand_def(OperationType)
    count = prop_def(IntegerAttr[I32])
    compareAtLeast = opt_prop_def(UnitAttr)
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = "`of` $input_op `is` (`at_least` $compareAtLeast^)? $count attr-dict `->` $true_dest `, ` $false_dest"

    def __init__(
        self,
        input_op: SSAValue,
        count: int | IntegerAttr[I32],
        trueDest: Block,
        falseDest: Block,
        compareAtLeast: bool = False,
    ) -> None:
        if isinstance(count, int):
            count = IntegerAttr(count, i32)
        properties = dict[str, Attribute](count=count)
        if compareAtLeast:
            properties["compareAtLeast"] = UnitAttr()
        super().__init__(
            operands=[input_op],
            properties=properties,
            successors=[trueDest, falseDest],
        )

@irdl_op_definition
class CheckResultCountOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpcheck_result_count-pdl_interpcheckresultcountop).
    """

    name = "pdl_interp.check_result_count"
    traits = traits_def(IsTerminator())
    input_op = operand_def(OperationType)
    count = prop_def(IntegerAttr[I32])
    compareAtLeast = opt_prop_def(UnitAttr)
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = "`of` $input_op `is` (`at_least` $compareAtLeast^)? $count attr-dict `->` $true_dest `, ` $false_dest"

    def __init__(
        self,
        input_op: SSAValue,
        count: int | IntegerAttr[I32],
        trueDest: Block,
        falseDest: Block,
        compareAtLeast: bool = False,
    ) -> None:
        if isinstance(count, int):
            count = IntegerAttr(count, i32)
        properties = dict[str, Attribute](count=count)
        if compareAtLeast:
            properties["compareAtLeast"] = UnitAttr()
        super().__init__(
            operands=[input_op],
            properties=properties,
            successors=[trueDest, falseDest],
        )

@irdl_op_definition
class IsNotNullOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpis_not_null-pdl_interpisnotnullop).
    """

    name = "pdl_interp.is_not_null"
    traits = traits_def(IsTerminator())
    value = operand_def(AnyPDLTypeConstr | base(RangeType[AnyPDLType]))
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = (
        "$value `:` type($value) attr-dict `->` $true_dest `, ` $false_dest"
    )

    def __init__(self, value: SSAValue, trueDest: Block, falseDest: Block) -> None:
        super().__init__(operands=[value], successors=[trueDest, falseDest])

@irdl_op_definition
class GetResultOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpget_result-pdl_interpgetresultop).
    """

    name = "pdl_interp.get_result"
    index = prop_def(IntegerAttr[I32])
    input_op = operand_def(OperationType)
    value = result_def(ValueType)

    assembly_format = "$index `of` $input_op attr-dict"

    def __init__(self, index: int | IntegerAttr[I32], input_op: SSAValue) -> None:
        if isinstance(index, int):
            index = IntegerAttr(index, i32)
        super().__init__(
            operands=[input_op], properties={"index": index}, result_types=[ValueType()]
        )

@irdl_op_definition
class GetResultsOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpget_results-pdl_interpgetresultsop).
    """

    name = "pdl_interp.get_results"
    index = opt_prop_def(IntegerAttr[I32])
    input_op = operand_def(OperationType)
    value = result_def(ValueType | RangeType[ValueType])

    # assembly_format = "($index^)? `of` $input_op `:` type($value) attr-dict"
    # TODO: Fix bug preventing this assebmly format from working: https://github.com/xdslproject/xdsl/issues/4136.

    def __init__(
        self,
        index: int | IntegerAttr[I32] | None,
        input_op: SSAValue,
        result_type: ValueType | RangeType[ValueType],
    ) -> None:
        if isinstance(index, int):
            index = IntegerAttr(index, i32)
        super().__init__(
            operands=[input_op],
            properties={"index": index},
            result_types=[result_type],
        )

    @classmethod
    def parse(cls, parser: Parser) -> GetResultsOp:
        index = parser.parse_optional_integer()
        if index is not None:
            index = IntegerAttr(index, 32)
        parser.parse_characters("of")
        input_op = parser.parse_operand(OperationType())
        parser.parse_punctuation(":")
        result_type = parser.parse_type()
        return GetResultsOp.build(
            operands=(input_op,),
            properties={"index": index},
            result_types=(result_type,),
        )

    def print(self, printer: Printer):
        if self.index is not None:
            printer.print_string(" ", indent=0)
            self.index.print_without_type(printer)
        printer.print_string(" of ", indent=0)
        printer.print_operand(self.input_op)
        printer.print_string(" : ", indent=0)
        printer.print_attribute(self.value.type)

@irdl_op_definition
class GetAttributeOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpget_attribute-pdl_interpgetattributeop).
    """

    name = "pdl_interp.get_attribute"
    constraint_name = prop_def(StringAttr, prop_name="name")
    input_op = operand_def(OperationType)
    value = result_def(AttributeType)

    assembly_format = "$name `of` $input_op attr-dict"

    def __init__(self, name: str | StringAttr, input_op: SSAValue) -> None:
        if isinstance(name, str):
            name = StringAttr(name)
        super().__init__(
            operands=[input_op],
            properties={"name": name},
            result_types=[AttributeType()],
        )

@irdl_op_definition
class GetAttributeTypeOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpget_attribute_type-pdl_interpgetattributetypeop).
    """

    name = "pdl_interp.get_attribute_type"
    value = operand_def(AttributeType)
    result = result_def(TypeType)

    assembly_format = "`of` $value attr-dict"

    def __init__(self, value: SSAValue) -> None:
        super().__init__(
            operands=[value],
            result_types=[TypeType()],
        )

@irdl_op_definition
class CheckAttributeOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpcheck_attribute-pdl_interpcheckattributeop).
    """

    name = "pdl_interp.check_attribute"
    traits = traits_def(IsTerminator())
    constantValue = prop_def()
    attribute = operand_def(AttributeType)
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = (
        "$attribute `is` $constantValue attr-dict `->` $true_dest `, ` $false_dest"
    )

    def __init__(
        self,
        constantValue: Attribute,
        attribute: SSAValue,
        trueDest: Block,
        falseDest: Block,
    ) -> None:
        super().__init__(
            operands=[attribute],
            properties={"constantValue": constantValue},
            successors=[trueDest, falseDest],
        )

@irdl_op_definition
class CheckTypeOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpcheck_type-pdl_interpchecktypeop).
    """

    name = "pdl_interp.check_type"
    traits = traits_def(IsTerminator())
    type = prop_def(TypeAttribute)
    value = operand_def(TypeType)
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = "$value `is` $type attr-dict `->` $true_dest `, ` $false_dest"

    def __init__(
        self, type: TypeAttribute, value: SSAValue, trueDest: Block, falseDest: Block
    ) -> None:
        super().__init__(
            operands=[value],
            properties={"type": type},
            successors=[trueDest, falseDest],
        )

@irdl_op_definition
class CheckTypesOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpcheck_types-pdl_interpchecktypesop).
    """

    name = "pdl_interp.check_types"
    traits = traits_def(IsTerminator())
    types = prop_def(ArrayAttr[TypeAttribute])
    value = operand_def(RangeType[TypeType])
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = "$value `are` $types attr-dict `->` $true_dest `, ` $false_dest"

    def __init__(
        self,
        types: ArrayAttr[TypeAttribute] | Sequence[TypeAttribute],
        value: SSAValue,
        trueDest: Block,
        falseDest: Block,
    ) -> None:
        if not isinstance(types, ArrayAttr):
            types = ArrayAttr(types)
        super().__init__(
            operands=[value],
            properties={"types": types},
            successors=[trueDest, falseDest],
        )

@irdl_op_definition
class AreEqualOp(IRDLOperation):
    """
    See external [documentation](https://mlir.llvm.org/docs/Dialects/PDLInterpOps/#pdl_interpare_equal-pdl_interpareequalop).
    """

    name = "pdl_interp.are_equal"
    traits = traits_def(IsTerminator())
    T: ClassVar = VarConstraint("T", AnyPDLTypeConstr | base(RangeType[AnyPDLType]))
    lhs = operand_def(T)
    rhs = operand_def(T)
    true_dest = successor_def()
    false_dest = successor_def()

    assembly_format = (
        "operands `:` type($lhs) attr-dict `->` $true_dest `, ` $false_dest"
    )

    def __init__(
        self, lhs: SSAValue, rhs: SSAValue, trueDest: Block, falseDest: Block
    ) -> None:
        super().__init__(operands=[lhs, rhs], successors=[trueDest, falseDest])