Operations

@dataclass(init=False, repr=False)
class IRDLOperation(Operation):
    assembly_format: ClassVar[str | None] = None
    custom_directives: ClassVar[tuple[type[CustomDirective], ...]] = ()

    def __init__(
        self: IRDLOperation,
        *,
        operands: (
            Sequence[SSAValue | Operation | Sequence[SSAValue | Operation] | None]
            | None
        ) = None,
        result_types: Sequence[Attribute | Sequence[Attribute] | None] | None = None,
        properties: Mapping[str, Attribute | None] | None = None,
        attributes: Mapping[str, Attribute | None] | None = None,
        successors: Sequence[Block | Sequence[Block] | None] | None = None,
        regions: (
            Sequence[
                Region
                | Sequence[Operation]
                | Sequence[Block]
                | Sequence[Region | Sequence[Operation] | Sequence[Block]]
                | None
            ]
            | None
        ) = None,
    ):
        if operands is None:
            operands = []
        if result_types is None:
            result_types = []
        if properties is None:
            properties = {}
        if attributes is None:
            attributes = {}
        if successors is None:
            successors = []
        if regions is None:
            regions = []
        irdl_op_init(
            self,
            type(self).get_irdl_definition(),
            operands=operands,
            result_types=result_types,
            properties=properties,
            attributes=attributes,
            successors=successors,
            regions=regions,
        )

    def __post_init__(self):
        op_def = self.get_irdl_definition()
        # Fill in default properties
        for prop_name, prop_def in op_def.properties.items():
            if (
                prop_name not in self.properties
                and not isinstance(prop_def, OptionalDef)
                and prop_def.default_value is not None
            ):
                self.properties[prop_name] = prop_def.default_value

        # Fill in default attributes
        for attr_name, attr_def in op_def.attributes.items():
            if (
                attr_name not in self.attributes
                and not isinstance(attr_def, OptionalDef)
                and attr_def.default_value is not None
            ):
                self.attributes[attr_name] = attr_def.default_value

        return super().__post_init__()

    @classmethod
    def build(
        cls,
        *,
        operands: (
            Sequence[SSAValue | Operation | Sequence[SSAValue | Operation] | None]
            | None
        ) = None,
        result_types: Sequence[Attribute | Sequence[Attribute] | None] | None = None,
        attributes: Mapping[str, Attribute | None] | None = None,
        properties: Mapping[str, Attribute | None] | None = None,
        successors: Sequence[Block | Sequence[Block] | None] | None = None,
        regions: (
            Sequence[
                Region
                | Sequence[Operation]
                | Sequence[Block]
                | Sequence[Region | Sequence[Operation] | Sequence[Block]]
                | None
            ]
            | None
        ) = None,
    ) -> Self:
        """Create a new operation using builders."""
        op = cls.__new__(cls)
        IRDLOperation.__init__(
            op,
            operands=operands,
            result_types=result_types,
            properties=properties,
            attributes=attributes,
            successors=successors,
            regions=regions,
        )
        return op

    @classmethod
    def get_irdl_definition(cls) -> OpDef:
        """Get the IRDL operation definition."""
        ...

    def __eq__(self, other: object) -> bool:
        return self is other

    def __hash__(self) -> int:
        return id(self)

@dataclass
class IRDLOption(ABC):
    """Additional option used in IRDL."""

@dataclass
class AttrSizedSegments(IRDLOption, ABC):
    """
    Expect an attribute on the operation that contains the segment sizes of the
    operand, result, region, or successor lists.
    For instance, the list `[a, b, c, d]` with segment sizes `[1, 3]` will result
    in the `[a], [b, c, d]` lists.
    The attribute must be a dense array of `i32`, its lenght must be equal to the
    number of segments (e.g. the number of operand definitions), and its sum must
    be equal to the number of elements in the list (e.g. the number of operands).
    """

    attribute_name: ClassVar[str]
    as_property: bool = False
    """Name of the attribute containing the segment sizes."""

    def container(self, op: Operation) -> dict[str, Attribute]:
        if self.as_property:
            return op.properties
        else:
            return op.attributes

@dataclass
class AttrSizedOperandSegments(AttrSizedSegments):
    """
    Expect an attribute on the operation that contains the sizes of the operand
    definitions.
    See `AttrSizedSegments` for more information.
    """

    attribute_name = "operandSegmentSizes"
    """Name of the attribute containing the variadic operand sizes."""

@dataclass
class AttrSizedResultSegments(AttrSizedSegments):
    """
    Expect an attribute on the operation that contains the sizes of the result
    definitions.
    See `AttrSizedSegments` for more information.
    """

    attribute_name = "resultSegmentSizes"
    """Name of the attribute containing the variadic result sizes."""

@dataclass
class AttrSizedRegionSegments(AttrSizedSegments):
    """
    Expect an attribute on the operation that contains the sizes of the region
    definitions.
    See `AttrSizedSegments` for more information.
    """

    attribute_name = "regionSegmentSizes"
    """Name of the attribute containing the variadic region sizes."""

@dataclass
class AttrSizedSuccessorSegments(AttrSizedSegments):
    """
    Expect an attribute on the operation that contains the sizes of the successor
    definitions.
    See `AttrSizedSegments` for more information.
    """

    attribute_name = "successorSegmentSizes"
    """Name of the attribute containing the variadic successor sizes."""

class SameVariadicSize(IRDLOption):
    """
    All variadic definitions should have the same size.
    """

class SameVariadicResultSize(SameVariadicSize):
    """
    All variadic results should have the same size.
    """

class SameVariadicOperandSize(SameVariadicSize):
    """
    All variadic operands should have the same size.
    """

class SameVariadicRegionSize(SameVariadicSize):
    """
    All variadic regions should have the same size.
    """

class SameVariadicSuccessorSize(SameVariadicSize):
    """
    All variadic successors should have the same size.
    """

@dataclass
class ParsePropInAttrDict(IRDLOption):
    """
    Allows properties to be omitted from the assembly format, causing them
    to be parsed as part of the attribute dictionary.
    This should only be used to ensure MLIR compatibility, it is otherwise
    bad design to use it.
    """

@dataclass
class OperandOrResultDef(ABC):
    """An operand or a result definition. Should not be used directly."""

@dataclass
class VariadicDef(OperandOrResultDef):
    """A variadic operand or result definition. Should not be used directly."""

@dataclass
class OptionalDef(VariadicDef):
    """An optional operand or result definition. Should not be used directly."""

@dataclass(init=False)
class OperandDef(OperandOrResultDef):
    """An IRDL operand definition."""

    constr: RangeConstraint
    """The operand constraint."""

    def __init__(self, attr: Attribute | type[Attribute] | AttrConstraint):
        self.constr = single_range_constr_coercion(attr)

@dataclass(init=False)
class VarOperandDef(OperandDef, VariadicDef):
    """An IRDL variadic operand definition."""

    def __init__(
        self,
        attr: Attribute | type[Attribute] | AttrConstraint | RangeConstraint,
    ):
        self.constr = range_constr_coercion(attr)

@dataclass(init=False)
class OptOperandDef(VarOperandDef, OptionalDef):
    """An IRDL optional operand definition."""

@dataclass(init=False)
class ResultDef(OperandOrResultDef):
    """An IRDL result definition."""

    constr: RangeConstraint
    """The result constraint."""

    def __init__(
        self, attr: Attribute | type[Attribute] | AttrConstraint | RangeConstraint
    ):
        assert not isinstance(attr, RangeConstraint)
        self.constr = single_range_constr_coercion(attr)

@dataclass(init=False)
class VarResultDef(ResultDef, VariadicDef):
    """An IRDL variadic result definition."""

    def __init__(
        self, attr: Attribute | type[Attribute] | AttrConstraint | RangeConstraint
    ):
        self.constr = range_constr_coercion(attr)

@dataclass(init=False)
class OptResultDef(VarResultDef, OptionalDef):
    """An IRDL optional result definition."""

@dataclass(init=True)
class RegionDef:
    """
    An IRDL region definition.
    """

    entry_args: RangeConstraint = field(default_factory=lambda: RangeOf(AnyAttr()))

@dataclass
class VarRegionDef(RegionDef, VariadicDef):
    """An IRDL variadic region definition."""

@dataclass
class OptRegionDef(RegionDef, OptionalDef):
    """An IRDL optional region definition."""

@dataclass
class SingleBlockRegionDef(RegionDef):
    """An IRDL region definition that expects exactly one block."""

class VarSingleBlockRegionDef(RegionDef, VariadicDef):
    """An IRDL variadic region definition that expects exactly one block."""

class OptSingleBlockRegionDef(RegionDef, OptionalDef):
    """An IRDL optional region definition that expects exactly one block."""

@dataclass
class AttrOrPropDef:
    """An IRDL attribute or property definition."""

    constr: AttrConstraint
    default_value: Attribute | None = None

@dataclass
class AttributeDef(AttrOrPropDef):
    """An IRDL attribute definition."""

@dataclass
class OptAttributeDef(AttributeDef, OptionalDef):
    """An IRDL attribute definition for an optional attribute."""

@dataclass
class PropertyDef(AttrOrPropDef):
    """An IRDL property definition."""

@dataclass
class OptPropertyDef(PropertyDef, OptionalDef):
    """An IRDL property definition for an optional property."""

class SuccessorDef:
    """An IRDL successor definition."""

class VarSuccessorDef(SuccessorDef, VariadicDef):
    """An IRDL variadic successor definition."""

class OptSuccessorDef(SuccessorDef, OptionalDef):
    """An IRDL optional successor definition."""

@dataclass(kw_only=True)
class OpDef:
    """The internal IRDL definition of an operation."""

    name: str = field(kw_only=False)
    operands: list[tuple[str, OperandDef]] = field(
        default_factory=list[tuple[str, OperandDef]]
    )
    results: list[tuple[str, ResultDef]] = field(
        default_factory=list[tuple[str, ResultDef]]
    )
    properties: dict[str, PropertyDef] = field(default_factory=dict[str, PropertyDef])
    attributes: dict[str, AttributeDef] = field(default_factory=dict[str, AttributeDef])
    regions: list[tuple[str, RegionDef]] = field(
        default_factory=list[tuple[str, RegionDef]]
    )
    successors: list[tuple[str, SuccessorDef]] = field(
        default_factory=list[tuple[str, SuccessorDef]]
    )
    options: list[IRDLOption] = field(default_factory=list[IRDLOption])
    traits: OpTraits = field(default_factory=lambda: traits_def())

    accessor_names: dict[str, tuple[str, Literal["attribute", "property"]]] = field(
        default_factory=dict[str, tuple[str, Literal["attribute", "property"]]]
    )
    """
    Mapping from the accessor name to the attribute or property name.
    In some cases, the attribute name is not a valid Python identifier,
    or is already used by the operation, so we need to use a different name.
    """
    assembly_format: str | None = field(default=None)
    custom_directives: dict[str, type[CustomDirective]] = field(
        default_factory=lambda: {}
    )

    @staticmethod
    def from_pyrdl(pyrdl_def: type[IRDLOperationInvT]) -> OpDef:
        """Decorator used on classes to define a new operation definition."""

        type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint] | None = None

        # If the operation inherit from `Generic`, this means that it specializes a
        # generic operation. Retrieve the mapping from `TypeVar` to pyrdl constraints.
        if issubclass(pyrdl_def, Generic):
            type_var_mapping = {
                k: irdl_to_attr_constraint(v)
                for k, v in get_type_var_mapping_with_defaults(pyrdl_def).items()
            }

        def wrong_field_exception(field_name: str) -> PyRDLOpDefinitionError:
            raise PyRDLOpDefinitionError(
                f"{pyrdl_def.__name__}.{field_name} is neither a function,"
                "operand, result, region, or attribute definition. "
                "Operands should be defined with type hints of "
                "operand_def(<Constraint>), results with "
                "result_def(<Constraint>), regions with "
                "region_def(), attributes with "
                "attr_def(<Constraint>), properties with prop_def(<Constraint>), "
                "and constants (indicated by uppercase field names) as ClassVar."
            )

        op_def = OpDef(pyrdl_def.name)

        # If an operation subclass overrides a superclass field, only keep the
        # definition of the subclass, with the exception of `traits`, which are
        # processed for all superclasses.
        field_names = set[str]()

        traits_defs: list[OpTraits] = []

        # Get all fields of the class, including the parent classes
        for parent_cls in pyrdl_def.mro():
            # Do not collect fields from Generic, as Generic will not contain
            # IRDL definitions, and contains ClassVar fields that are not
            # allowed in IRDL definitions.
            if parent_cls == Generic:
                continue
            if parent_cls in Operation.mro():
                continue

            clsdict = parent_cls.__dict__

            annotations = parent_cls.__annotations__

            for field_name in annotations:
                if field_name not in clsdict:
                    if is_const_classvar(
                        field_name, annotations[field_name], PyRDLOpDefinitionError
                    ):
                        continue
                    raise wrong_field_exception(field_name)

            for field_name in clsdict:
                if field_name in ("name", "assembly_format", "custom_directives"):
                    continue
                if field_name in _OPERATION_DICT_KEYS:
                    # Fields that are already in Operation (i.e. operands, results, ...)
                    continue
                if field_name in field_names:
                    # already registered value for field name
                    continue
                if field_name in annotations and is_const_classvar(
                    field_name, annotations[field_name], PyRDLOpDefinitionError
                ):
                    continue

                value = clsdict[field_name]

                # Check that all fields of the operation definition are either already
                # in Operation, or are class functions or methods.

                if field_name == "irdl_options":
                    if not isa(value, tuple[IRDLOption, ...]):
                        if not isinstance(value, tuple):
                            raise PyRDLOpDefinitionError(
                                f"All values `irdl_options` must be a `tuple`, got `{type(value).__name__}`."
                            )
                        raise PyRDLOpDefinitionError(
                            "All values in irdl_options should inherit IRDLOption"
                        )
                    op_def.options.extend(value)
                    for option in value:
                        if isinstance(option, AttrSizedSegments):
                            defs = (
                                op_def.properties
                                if option.as_property
                                else op_def.attributes
                            )
                            def_name = "property" if option.as_property else "attribute"
                            if option.attribute_name in defs:
                                raise PyRDLOpDefinitionError(
                                    f"pyrdl operation definition '{pyrdl_def.__name__}' "
                                    f"has a '{option.attribute_name}' {def_name}, which "
                                    "is incompatible with the "
                                    f"{option} option."
                                )
                            from xdsl.dialects.builtin import DenseArrayBase

                            if option.as_property:
                                prop_def = PropertyDef(
                                    irdl_to_attr_constraint(DenseArrayBase)
                                )
                                op_def.properties[option.attribute_name] = prop_def
                            else:
                                attr_def = AttributeDef(
                                    irdl_to_attr_constraint(DenseArrayBase)
                                )
                                op_def.attributes[option.attribute_name] = attr_def
                    continue

                if field_name == "traits":
                    traits = value
                    if not isinstance(traits, OpTraits):
                        raise PyRDLOpDefinitionError(
                            f"pyrdl operation definition '{pyrdl_def.__name__}' "
                            "traits field should be an instance of"
                            f"'{OpTraits.__name__}'."
                        )
                    traits_defs.append(traits)
                    continue

                # Dunder fields are allowed (i.e. __orig_bases__, __annotations__, ...)
                # They are used by Python to store information about the class, so they
                # should not be considered as part of the operation definition.
                # Also, they can provide a possiblea escape hatch.
                if field_name[:2] == "__" and field_name[-2:] == "__":
                    continue

                # Methods, properties, and functions are allowed
                if isinstance(
                    value, FunctionType | PropertyType | classmethod | staticmethod
                ):
                    continue

                # Get attribute constraints from a list of pyrdl constraints
                def get_constraint(
                    pyrdl_constr: IRDLAttrConstraint,
                ) -> AttrConstraint:
                    constraint = irdl_list_to_attr_constraint(
                        (pyrdl_constr,),
                        allow_type_var=True,
                    )
                    if type_var_mapping is not None:
                        constraint = constraint.mapping_type_vars(type_var_mapping)
                    return constraint

                # Get attribute constraints from a list of pyrdl constraints
                def get_range_constraint(
                    pyrdl_constr: RangeConstraint | IRDLAttrConstraint,
                ) -> RangeConstraint:
                    if isinstance(pyrdl_constr, RangeConstraint):
                        # Pyright does not know the type of the generic range constraint
                        return cast(RangeConstraint, pyrdl_constr)
                    return RangeOf(get_constraint(pyrdl_constr))

                field_names.add(field_name)

                match value:
                    case _ResultFieldDef():
                        if not issubclass(value.cls, VariadicDef):
                            if isinstance(value.param, RangeConstraint):
                                raise TypeError(
                                    "Cannot use a RangeConstraint in result_def, use an "
                                    "AttrConstraint or var_result_def or "
                                    "opt_result_def instead."
                                )
                            constraint = get_constraint(value.param)
                            result_def = value.cls(constraint)
                        else:
                            constraint = get_range_constraint(value.param)
                            result_def = value.cls(constraint)
                        op_def.results.append((field_name, result_def))
                        continue
                    case _OperandFieldDef():
                        if not issubclass(value.cls, VariadicDef):
                            if isinstance(value.param, RangeConstraint):
                                raise TypeError(
                                    "Cannot use a RangeConstraint in operand_def, use an "
                                    "AttrConstraint or var_operand_def or "
                                    "opt_operand_def instead."
                                )
                            constraint = get_constraint(value.param)
                            operand_def = value.cls(constraint)
                        else:
                            constraint = get_range_constraint(value.param)
                            operand_def = cast(type[VarOperandDef], value.cls)(
                                constraint
                            )
                        op_def.operands.append((field_name, operand_def))
                        continue
                    case _AttributeFieldDef():
                        constraint = get_constraint(value.param)
                        attribute_def = value.cls(constraint, value.default_value)
                        ir_name = field_name if value.ir_name is None else value.ir_name
                        op_def.attributes[ir_name] = attribute_def
                        op_def.accessor_names[field_name] = (ir_name, "attribute")
                        continue
                    case _PropertyFieldDef():
                        constraint = get_constraint(value.param)
                        property_def = value.cls(constraint, value.default_value)
                        ir_name = field_name if value.ir_name is None else value.ir_name
                        op_def.properties[ir_name] = property_def
                        op_def.accessor_names[field_name] = (ir_name, "property")
                        continue
                    case _RegionFieldDef():
                        constraint = get_range_constraint(value.entry_args)
                        region_def = value.cls(constraint)
                        op_def.regions.append((field_name, region_def))
                        continue
                    case _SuccessorFieldDef():
                        successor_def = value.cls()
                        op_def.successors.append((field_name, successor_def))
                        continue
                    case _:
                        pass

                raise wrong_field_exception(field_name)

        if traits_defs:
            if len(traits_defs) == 1:
                op_def.traits = traits_defs[0]
            else:
                op_def.traits = OpTraits(
                    lambda: tuple(
                        trait for traits in traits_defs for trait in traits.gen_traits()
                    )
                )
        op_def.assembly_format = pyrdl_def.assembly_format
        op_def.custom_directives = {
            directive.__name__: directive for directive in pyrdl_def.custom_directives
        }
        assert inspect.ismethod(Operation.parse)
        if op_def.assembly_format is not None and (
            pyrdl_def.print != Operation.print
            or not inspect.ismethod(pyrdl_def.parse)
            or pyrdl_def.parse.__func__ != Operation.parse.__func__
        ):
            raise PyRDLOpDefinitionError(
                "Cannot define both an assembly format (with the assembly_format "
                "variable) and the print and parse methods."
            )

        return op_def

    def verify(self, op: Operation):
        """Given an IRDL definition, verify that an operation satisfies its invariants."""

        # Mapping from type variables to their concrete types.
        constraint_context = VerificationContext()

        # Verify operands.
        irdl_op_verify_arg_list(op, self, VarIRConstruct.OPERAND, constraint_context)

        # Verify results.
        irdl_op_verify_arg_list(op, self, VarIRConstruct.RESULT, constraint_context)

        # Verify regions.
        irdl_op_verify_regions(op, self, constraint_context)

        # Verify successors.
        verify_variadic_size(op, self, VarIRConstruct.SUCCESSOR)

        # Verify properties.
        for prop_name, attr_def in self.properties.items():
            if prop_name not in op.properties:
                if isinstance(attr_def, OptPropertyDef):
                    continue
                raise VerifyException(
                    f"property '{prop_name}' expected in operation '{op.name}'"
                )
            attr_def.constr.verify(op.properties[prop_name], constraint_context)

        for prop_name in op.properties.keys():
            if prop_name not in self.properties:
                raise VerifyException(
                    f"property '{prop_name}' is not defined by the operation '{op.name}'. "
                    "Use the dictionary attribute to add arbitrary information "
                    "to the operation."
                )

        # Verify attributes.
        for attr_name, attr_def in self.attributes.items():
            if attr_name not in op.attributes:
                if isinstance(attr_def, OptAttributeDef):
                    continue
                raise VerifyException(
                    f"attribute '{attr_name}' expected in operation '{op.name}'"
                )
            attr_def.constr.verify(op.attributes[attr_name], constraint_context)

        # Verify traits.
        for trait in self.traits:
            trait.verify(op)