Index

class DialectInterface(ABC):
    """
    A base class for dialects' interfaces.
    They usually define functionality which is dialect specific to some transformation.

    For example DialectInlinerInterface defines which dialect operations can be inlined and how.
    Dialects will implement this interface and the inlining transformation will query them through the base interface.

    The design logic tries to follow MLIR's dialect interfaces closely
    https://mlir.llvm.org/docs/Interfaces/#dialect-interfaces
    """

class IsTerminator(OpTrait):
    """
    This trait provides verification and functionality for operations that are
    known to be terminators.

    See external [documentation](https://mlir.llvm.org/docs/Traits/#terminator).
    """

    def verify(self, op: Operation) -> None:
        """Check that the operation satisfies the IsTerminator trait requirements."""
        if op.parent is not None and op.parent.last_op != op:
            raise VerifyException(
                f"'{op.name}' must be the last operation in its parent block"
            )

class NoTerminator(OpTrait):
    """
    Allow an operation to have single block regions with no terminator.

    See external [documentation](https://mlir.llvm.org/docs/Traits/#terminator).
    """

    def verify(self, op: Operation) -> None:
        for region in op.regions:
            if len(region.blocks) > 1:
                raise VerifyException(
                    f"'{op.name}' does not contain single-block regions"
                )

@dataclass(frozen=True)
class OpTrait:
    """
    A trait attached to an operation definition.
    Traits can be used to define operation invariants, additional semantic information,
    or to group operations that have similar properties.
    Note that traits are the merge of traits and interfaces in MLIR.
    """

    def verify(self, op: Operation) -> None:
        """Check that the operation satisfies the trait requirements."""
        pass

class VerifyException(DiagnosticException):
    pass

@dataclass
class Dialect:
    """Contains the operations and attributes of a specific dialect"""

    _name: str

    _operations: list[type[Operation]] = field(
        default_factory=list[type["Operation"]], init=True, repr=True
    )
    _attributes: list[type[Attribute]] = field(
        default_factory=list[type["Attribute"]], init=True, repr=True
    )
    _interfaces: list[DialectInterface] = field(
        default_factory=list[DialectInterface], init=True, repr=True
    )

    @property
    def operations(self) -> Iterator[type[Operation]]:
        return iter(self._operations)

    @property
    def attributes(self) -> Iterator[type[Attribute]]:
        return iter(self._attributes)

    @property
    def name(self) -> str:
        return self._name

    @staticmethod
    def split_name(name: str) -> tuple[str, str]:
        try:
            names = name.split(".", 1)
            first, second = names
            return (first, second)
        except ValueError as e:
            raise ValueError(f"Invalid operation or attribute name {name}.") from e

    def get_interface(
        self, interface: type[DialectInterfaceT]
    ) -> DialectInterfaceT | None:
        """
        Return a class that implements the 'interface' if it exists.
        """
        for i in self._interfaces:
            if isinstance(i, interface):
                return i
        return None

    def has_interface(self, interface: type[DialectInterfaceT]) -> bool:
        return self.get_interface(interface) is not None

@dataclass(frozen=True)
class Attribute(ABC):
    """
    A compile-time value.
    Attributes are used to represent SSA variable types, and can be attached
    on operations to give extra information.
    """

    name: ClassVar[str] = field(init=False, repr=False)
    """The attribute name should be a static field in the attribute classes."""

    def __post_init__(self):
        self._verify()
        if not isinstance(self, Data | ParametrizedAttribute):
            raise TypeError("Attributes should only be Data or ParametrizedAttribute")

    def _verify(self):
        self.verify()

    def verify(self) -> None:
        """
        Check that the attribute parameters satisfy the expected invariants.
        Raise a VerifyException otherwise.
        """
        pass

    def __str__(self) -> str:
        from xdsl.printer import Printer

        res = StringIO()
        printer = Printer(stream=res)
        printer.print_attribute(self)
        return res.getvalue()

@dataclass(frozen=True, init=False)
class BuiltinAttribute(Attribute, ABC):
    """
    This class is used to mark builtin attributes.
    Unlike other attributes in MLIR, parsing of *Builtin* attributes
    is handled directly by the parser.
    Printing of these attributes is handled by the `print_builtin` function, which must
    be implemented by all *Builtin* attributes.
    Attributes outside of the `builtin` dialect should not inherit from `BuiltinAttribute`.
    """

    @abstractmethod
    def print_builtin(self, printer: Printer) -> None:
        """
        Prints the attribute using the supplied printer.
        `BuiltinAttribute`s need not follow the same rules as other attributes, for example
        they do not need to be prefixed by `!` or `#` and do not need to print their name.
        """
        ...

class TypeAttribute(Attribute):
    """
    This class should only be inherited by classes inheriting Attribute.
    This class is only used for printing attributes in the MLIR format,
    inheriting this class prefix the attribute by `!` instead of `#`.
    """

    pass

class OpaqueSyntaxAttribute(Attribute):
    """
    This class should only be inherited by classes inheriting Attribute.
    This class is only used for printing attributes in the opaque form.

    See external [documentation](https://mlir.llvm.org/docs/LangRef/#dialect-attribute-values.).
    """

    pass

class SpacedOpaqueSyntaxAttribute(OpaqueSyntaxAttribute):
    """
    This class should only be inherited by classes inheriting Attribute.
    This class is only used for printing attributes in the opaque form.

    See external [documentation](https://mlir.llvm.org/docs/LangRef/#dialect-attribute-values.).
    """

    pass

@dataclass(frozen=True)
class Data(Attribute, ABC, Generic[DataElement]):
    """An attribute represented by a Python structure."""

    data: DataElement

    @classmethod
    def new(cls: type[Self], params: Any) -> Self:
        """
        Create a new `Data` given its parameter.

        The `params` argument should be of the same type as the `Data` generic
        argument.

        This function should be preferred over `__init__` when instantiating
        attributes in a generic way (i.e., without knowing their concrete type
        statically).
        """
        # Create the new attribute object, without calling its __init__.
        # We do this to allow users to redefine their own __init__.
        attr = cls.__new__(cls)

        # Call the __init__ of Data, which will set the parameters field.
        Data.__init__(attr, params)  # pyright: ignore[reportUnknownMemberType]
        return attr

    @classmethod
    def get(cls, attr: DataElement | Self) -> Self:
        """
        Creates an element of this class from `DataElement`,
        or returns the input when it is already an instance of this class.

        This function is useful for `__init__` methods, for example when a we
        would like to accept either a `StringAttr` or a `str`.
        """
        if not isinstance(attr, cls):
            attr = cls.new(attr)
        return attr

    @classmethod
    @abstractmethod
    def parse_parameter(cls, parser: AttrParser) -> DataElement:
        """Parse the attribute parameter."""

    @abstractmethod
    def print_parameter(self, printer: Printer) -> None:
        """Print the attribute parameter."""

@dataclass(frozen=True, init=False)
class ParametrizedAttribute(Attribute):
    """An attribute parametrized by other attributes."""

    def __init__(self, *parameters: Any):
        irdl_def = self.get_irdl_definition()
        for (f, d), param in zip(irdl_def.parameters, parameters, strict=True):
            if d.converter is not None:
                param = d.converter(param)
            object.__setattr__(self, f, param)
        super().__init__()

    @property
    def parameters(self) -> tuple[Attribute, ...]:
        return (
            *(
                self.__getattribute__(field)
                for field, _ in self.get_irdl_definition().parameters
            ),
        )

    @classmethod
    def new(cls: type[Self], params: Sequence[Attribute]) -> Self:
        """
        Create a new `ParametrizedAttribute` given its parameters.

        This function should be preferred over `__init__` when instantiating
        attributes in a generic way (i.e., without knowing their concrete type
        statically).
        """
        # Create the new attribute object, without calling its __init__.
        # We do this to allow users to redefine their own __init__.
        attr = cls.__new__(cls)

        # Set the parameters based on the definition
        for (f, _), param in zip(
            cls.get_irdl_definition().parameters, params, strict=True
        ):
            object.__setattr__(attr, f, param)
        attr.__post_init__()

        return attr

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> Sequence[Attribute]:
        """Parse the attribute parameters."""
        return parser.parse_paramattr_parameters()

    def print_parameters(self, printer: Printer) -> None:
        """Print the attribute parameters."""
        printer.print_paramattr_parameters(self.parameters)

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

    def _verify(self):
        # Verifier generated by irdl_attr_def
        t: type[ParametrizedAttribute] = type(self)
        attr_def = t.get_irdl_definition()
        attr_def.verify(self)
        super()._verify()

class TypedAttribute(ParametrizedAttribute, ABC):
    """
    An attribute with a type.
    """

    @classmethod
    def get_type_index(cls) -> int: ...

    def get_type(self) -> Attribute:
        return self.parameters[self.get_type_index()]

    @staticmethod
    def parse_with_type(
        parser: AttrParser,
        type: Attribute,
    ) -> TypedAttribute:
        """
        Parse the attribute with the given type.
        """
        ...

    @abstractmethod
    def print_without_type(self, printer: Printer): ...

@dataclass(repr=False)
class Use:
    """The use of a SSA value."""

    _operation: Operation
    """The operation using the value."""

    _index: int
    """The index of the operand using the value in the operation."""

    _prev_use: Use | None = None
    """The previous use of the value in the use list."""

    _next_use: Use | None = None
    """The next use of the value in the use list."""

    @property
    def operation(self) -> Operation:
        """The operation using the value."""
        return self._operation

    @property
    def index(self) -> int:
        """The index of the operand using the value in the operation."""
        return self._index

    def __repr__(self) -> str:
        return (
            f"<Use {id(self)}(_operation={_short_repr(self.operation)}, "
            f"_index={self.index}, _prev_use={_short_repr(self._prev_use)}, "
            f"_next_use={_short_repr(self._next_use)})>"
        )

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

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

@dataclass
class IRUses(Iterable[Use]):
    """
    Multi-pass iterable of the uses of an IR value (SSAValue or Block).
    """

    ir: IRWithUses

    def __iter__(self):
        use = self.ir.first_use
        while use is not None:
            yield use
            use = use._next_use  # pyright: ignore[reportPrivateUsage]

    def __bool__(self) -> bool:
        """Returns `True` if there are operations in this block."""
        return self.ir.first_use is not None

    def get_length(self) -> int:
        """
        Returns the number of uses.
        We do not expose it as `__len__` as it is expensive to compute `O(len)`.
        """
        use = self.ir.first_use
        count = 0
        while use is not None:
            count += 1
            use = use._next_use  # pyright: ignore[reportPrivateUsage]
        return count

@dataclass(eq=False)
class IRWithUses(ABC):
    """IRNode which stores a list of its uses."""

    first_use: Use | None = field(init=False, default=None, repr=False)
    """The first use of the value in the use list."""

    @property
    def uses(self) -> IRUses:
        """Returns an iterable of all uses of the value."""
        return IRUses(self)

    def add_use(self, use: Use):
        """Add a new use of the value."""
        first_use = self.first_use
        use._next_use = first_use  # pyright: ignore[reportPrivateUsage]
        use._prev_use = None  # pyright: ignore[reportPrivateUsage]
        if first_use is not None:
            first_use._prev_use = use  # pyright: ignore[reportPrivateUsage]
        self.first_use = use

    def remove_use(self, use: Use):
        """Remove a use of the value."""
        prev_use = use._prev_use  # pyright: ignore[reportPrivateUsage]
        next_use = use._next_use  # pyright: ignore[reportPrivateUsage]
        if prev_use is not None:
            prev_use._next_use = next_use  # pyright: ignore[reportPrivateUsage]
        if next_use is not None:
            next_use._prev_use = prev_use  # pyright: ignore[reportPrivateUsage]

        if prev_use is None:
            self.first_use = next_use

    def has_one_use(self) -> bool:
        """Returns true if the value has exactly one use."""
        first_use = self.first_use
        return (
            first_use is not None and first_use._next_use is None  # pyright: ignore[reportPrivateUsage]
        )

    def has_more_than_one_use(self) -> bool:
        """Returns true if the value has more than one use."""
        if (first_use := self.first_use) is None:
            return False
        return first_use._next_use is not None  # pyright: ignore[reportPrivateUsage]

    def get_unique_use(self) -> Use | None:
        """
        Returns the single use of the value, or None if there are no uses or
        more than one use.
        """
        if (first_use := self.first_use) is None:
            return None
        if first_use._next_use is not None:  # pyright: ignore[reportPrivateUsage]
            return None
        return first_use

    def get_user_of_unique_use(self) -> Operation | None:
        """
        Returns the user of the single use of the value.
        If there are no uses or more than one use, returns None.
        """
        if (use := self.get_unique_use()) is not None:
            return use.operation
        return None

@dataclass(eq=False)
class IRWithName(ABC):
    _name: str | None = field(init=False, default=None)

    @property
    def name_hint(self) -> str | None:
        return self._name

    @name_hint.setter
    def name_hint(self, name: str | None):
        self._name = self.extract_valid_name(name)

    @classmethod
    def is_valid_name(cls, name: str | None):
        """
        Returns `True` if `name` is None or a valid name for an identifier.
        """
        return name is None or _VALUE_NAME_PATTERN.fullmatch(name)

    @overload
    @classmethod
    def extract_valid_name(cls, name: str) -> str: ...

    @overload
    @classmethod
    def extract_valid_name(cls, name: None) -> None: ...

    @classmethod
    def extract_valid_name(cls, name: str | None) -> str | None:
        """
        If the name is valid, extracts the name before an optional `_\\d+` suffix.
        Raises ValueError otherwise.
        """
        if name is None:
            return
        if _VALUE_NAME_PATTERN.fullmatch(name) is None:
            raise ValueError(
                f"Invalid {cls.__name__} name format `{name}`.",
                r"Make sure names contain only characters of [A-Za-z0-9_$.-] and don't start with a number.",
            )

        if match := _VALUE_NAME_SUFFIX_PATTERN.search(name):
            # Remove `_` followed by numbers at the end of the name
            return name[: match.start()]

        return name