Builtin

class ShapedType(Attribute, ABC):
    @abstractmethod
    def get_num_dims(self) -> int: ...

    @abstractmethod
    def get_shape(self) -> tuple[int, ...]: ...

    def element_count(self) -> int:
        return prod(self.get_shape())

    def has_static_shape(self) -> bool:
        """
        Check if all dimensions have static (non-dynamic) shapes.
        """
        return all(dim != DYNAMIC_INDEX for dim in self.get_shape())

    @staticmethod
    def strides_for_shape(
        shape: Sequence[int], factor: int = 1
    ) -> tuple[int | None, ...]:
        """
        Returns a tuple of strides for a given shape, with row-major layout.
        Strides that depend on a dynamic index are returned as `None`.
        The optional `factor` parameter specifies the stride for the innermost
        dimension, defaulting to 1.
        """
        rev_strides: list[int | None] = []
        stride: int | None = factor
        for dim in reversed(shape):
            rev_strides.append(stride)
            if stride is None or dim == DYNAMIC_INDEX:
                stride = None
            else:
                stride *= dim
        return tuple(reversed(rev_strides))

class ContainerType(ABC, Generic[_ContainerElementTypeT]):
    @abstractmethod
    def get_element_type(self) -> _ContainerElementTypeT:
        pass

@irdl_attr_definition
class NoneAttr(ParametrizedAttribute, BuiltinAttribute):
    """An attribute representing the absence of an attribute."""

    name = "none"

    def print_builtin(self, printer: Printer):
        printer.print_string("none")

@irdl_attr_definition
class ArrayAttr(
    _BuiltinGenericData[tuple[AttributeCovT, ...]],
    Iterable[AttributeCovT],
    Generic[AttributeCovT],
):
    name = "array"

    def __init__(self, param: Iterable[AttributeCovT]) -> None:
        super().__init__(tuple(param))

    def print_builtin(self, printer: Printer):
        with printer.in_square_brackets():
            printer.print_list(self.data, printer.print_attribute)

    @staticmethod
    @override
    def constr(
        constr: IRDLAttrConstraint[AttributeInvT]
        | RangeConstraint[AttributeInvT]
        | None = None,
    ) -> ArrayOfConstraint[AttributeInvT]:
        return ArrayOfConstraint[AttributeInvT](
            cast(
                AttrConstraint[AttributeInvT],
                TypeVarConstraint(AttributeCovT, AnyAttr()),
            )
            if constr is None
            else constr
        )

    def __len__(self):
        return len(self.data)

    def __iter__(self) -> Iterator[AttributeCovT]:
        return iter(self.data)

@dataclass(frozen=True)
class ArrayOfConstraint(AttrConstraint[ArrayAttr[AttributeCovT]]):
    elem_range_constraint: RangeConstraint[AttributeCovT]
    """
    A constraint that enforces an ArrayData whose elements satisfy
    the underlying range constraint.
    """

    def __init__(
        self,
        constr: (IRDLAttrConstraint[AttributeCovT] | RangeConstraint[AttributeCovT]),
    ):
        if isinstance(constr, RangeConstraint):
            object.__setattr__(self, "elem_range_constraint", constr)
        else:
            object.__setattr__(self, "elem_range_constraint", RangeOf.get(constr))

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        if not isa(attr, ArrayAttr):
            raise VerifyException(
                f"expected ArrayAttr attribute, but got '{type(attr)}'"
            )
        self.elem_range_constraint.verify(attr.data, constraint_context)

    def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
        return self.elem_range_constraint.can_infer(
            var_constraint_names, length_known=False
        )

    def infer(self, context: InferenceContext) -> ArrayAttr[AttributeCovT]:
        return ArrayAttr(self.elem_range_constraint.infer(context, length=None))

    def get_bases(self) -> set[type[Attribute]] | None:
        return {ArrayAttr}

    def variables(self) -> set[str]:
        return self.elem_range_constraint.variables()

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> AttrConstraint[ArrayAttr[AttributeCovT]]:
        return ArrayOfConstraint(
            self.elem_range_constraint.mapping_type_vars(type_var_mapping)
        )

@irdl_attr_definition
class StringAttr(_BuiltinData[str]):
    name = "string"

    def print_builtin(self, printer: Printer):
        printer.print_string_literal(self.data)

@dataclass(frozen=True)
class SymbolNameConstraint(AttrConstraint[StringAttr]):
    """
    Constrain an attribute to be a StringAttr.
    This constraint has special assembly format support.
    """

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        if not isinstance(attr, StringAttr):
            raise VerifyException(f"{attr} should be a string")

    def get_bases(self) -> set[type[Attribute]] | None:
        return {StringAttr}

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> AttrConstraint[StringAttr]:
        return self

@irdl_attr_definition
class BytesAttr(_BuiltinData[bytes]):
    name = "bytes"

    def print_builtin(self, printer: Printer):
        printer.print_bytes_literal(self.data)

@irdl_attr_definition
class SymbolRefAttr(ParametrizedAttribute, BuiltinAttribute):
    name = "symbol_ref"
    root_reference: StringAttr
    nested_references: ArrayAttr[StringAttr]

    def __init__(
        self,
        root: str | StringAttr,
        nested: Sequence[str] | Sequence[StringAttr] | ArrayAttr[StringAttr] = [],
    ) -> None:
        if isinstance(root, str):
            root = StringAttr(root)
        if not isinstance(nested, ArrayAttr):
            nested = ArrayAttr(
                [StringAttr(x) if isinstance(x, str) else x for x in nested]
            )
        super().__init__(root, nested)

    @staticmethod
    def get(value: str | StringAttr | SymbolRefAttr) -> SymbolRefAttr:
        if isinstance(value, str | StringAttr):
            return SymbolRefAttr(value)
        return value

    def string_value(self):
        root = self.root_reference.data
        for ref in self.nested_references.data:
            root += "." + ref.data
        return root

    def print_builtin(self, printer: Printer):
        printer.print_symbol_name(self.root_reference)
        for ref in self.nested_references.data:
            printer.print_string("::")
            printer.print_symbol_name(ref)

class EmptyArrayAttrConstraint(AttrConstraint):
    """
    Constrain attribute to be empty ArrayData
    """

    def verify(self, attr: Attribute, constraint_context: VerificationContext) -> None:
        if not isa(attr, ArrayAttr):
            raise VerifyException(f"expected ArrayData attribute, but got {attr}")
        if attr.data:
            raise VerifyException(f"expected empty array, but got {attr}")

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> EmptyArrayAttrConstraint:
        return self

@irdl_attr_definition
class IntAttr(GenericData[_IntCovT], Generic[_IntCovT]):
    name = "builtin.int"

    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> _IntCovT:
        with parser.in_angle_brackets():
            data = parser.parse_integer()
            return cast(_IntCovT, data)

    def print_parameter(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string(f"{self.data}")

    def __bool__(self) -> bool:
        """Returns True if value is non-zero."""
        return bool(self.data)

    @staticmethod
    @override
    def constr(
        constr: IntConstraint | int | TypeForm[int] | None = None,
    ) -> AttrConstraint[IntAttr]:
        return IntAttrConstraint.get(
            IntTypeVarConstraint(_IntCovT, AnyInt()) if constr is None else constr
        )

@dataclass(frozen=True)
class IntAttrConstraint(AttrConstraint[IntAttr]):
    """
    Constrains the value of an IntAttr.
    """

    @staticmethod
    def get(
        int_constraint: IntConstraint | int | TypeForm[int] | None = None,
    ) -> AttrConstraint[IntAttr]:
        if int_constraint is None:
            return BaseAttr(IntAttr)
        if not isinstance(int_constraint, IntConstraint):
            int_constraint = get_int_constraint(int_constraint)
        if int_constraint == AnyInt():
            return BaseAttr(IntAttr)
        if isinstance(int_constraint, EqIntConstraint):
            return EqAttrConstraint(IntAttr(int_constraint.value))
        return IntAttrConstraint(int_constraint)

    int_constraint: IntConstraint

    def verify(self, attr: Attribute, constraint_context: VerificationContext) -> None:
        if not isa(attr, IntAttr):
            raise VerifyException(f"attribute {attr} expected to be an IntAttr")
        self.int_constraint.verify(attr.data, constraint_context)

    def variables(self) -> set[str]:
        return self.int_constraint.variables()

    def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
        return self.int_constraint.can_infer(var_constraint_names)

    def infer(self, context: InferenceContext) -> IntAttr:
        return IntAttr(self.int_constraint.infer(context))

    def get_bases(self) -> set[type[Attribute]] | None:
        return {IntAttr}

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ):
        return IntAttrConstraint.get(
            self.int_constraint.mapping_type_vars(type_var_mapping)
        )

class Signedness(ConstraintConvertible, Enum):
    "Signedness semantics for integer"

    SIGNLESS = 0
    "No signedness semantics"

    SIGNED = 1
    UNSIGNED = 2

    def value_range(self, bitwidth: int) -> tuple[int, int]:
        """
        For a given bitwidth, returns (min, max+1), where min and max are the smallest and
        largest representable values.

        Signless integers are bit patterns, so the representable range is the union of the
        signed and unsigned representable ranges.
        """
        match self:
            case Signedness.SIGNLESS:
                return signless_value_range(bitwidth)
            case Signedness.SIGNED:
                return signed_value_range(bitwidth)
            case Signedness.UNSIGNED:
                return unsigned_value_range(bitwidth)

    @staticmethod
    def base_constr() -> AttrConstraint[SignednessAttr]:
        """The constraint for this class."""
        return BaseAttr(SignednessAttr)

    def constr(self) -> AttrConstraint[SignednessAttr]:
        return EqAttrConstraint(SignednessAttr(self))