Constraints

@dataclass
class InferenceContext(ABC):
    """
    Stores the assignments of constraint variables without allowing assignment.
    Used for inference, at which point variables are read-only.
    """

    _variables: dict[str, Attribute] = field(default_factory=dict[str, Attribute])
    """The assignment of constraint variables."""

    _range_variables: dict[str, tuple[Attribute, ...]] = field(
        default_factory=dict[str, tuple[Attribute, ...]]
    )
    """The assignment of constraint range variables."""

    _int_variables: dict[str, int] = field(default_factory=dict[str, int])
    """The assignment of constraint int variables."""

    def get_variable(self, key: str) -> Attribute | None:
        return self._variables.get(key)

    def get_range_variable(self, key: str) -> tuple[Attribute, ...] | None:
        return self._range_variables.get(key)

    def get_int_variable(self, key: str) -> int | None:
        return self._int_variables.get(key)

    @property
    def attr_variables(self) -> AbstractSet[str]:
        return self._variables.keys()

    @property
    def range_variables(self) -> AbstractSet[str]:
        return self._range_variables.keys()

    @property
    def int_variables(self) -> AbstractSet[str]:
        return self._int_variables.keys()

@dataclass
class VerificationContext(InferenceContext):
    """
    Contains the assignment of constraint variables, allowing assignment.
    """

    def set_attr_variable(self, key: str, attr: Attribute) -> bool:
        """
        Tries to assign the attribute 'attr' to key 'key'.
        If a different attribute is already assigned to this key, then this function
        throws a VerifyError.

        The return value indicates whether this key was set by this call.
        """
        if (existing := self._variables.get(key)) is not None:
            if existing != attr:
                raise VerifyException(
                    f"attribute {existing} expected from variable "
                    f"'{key}', but got {attr}"
                )
            return False
        self._variables[key] = attr
        return True

    def set_range_variable(self, key: str, attrs: tuple[Attribute, ...]) -> bool:
        """
        Tries to assign the attributes 'attrs' to key 'key'.
        If a different attribute range is already assigned to this key, then this
        function throws a VerifyError.

        The return value indicates whether this key was set by this call.
        """
        if (existing := self._range_variables.get(key)) is not None:
            if existing != attrs:
                raise VerifyException(
                    f"attributes {tuple(str(x) for x in existing)} expected from range variable "
                    f"'{key}', but got {tuple(str(x) for x in attrs)}"
                )
            return False
        self._range_variables[key] = attrs
        return True

    def set_int_variable(self, key: str, i: int) -> bool:
        """
        Tries to assign the integer 'i' to key 'key'.
        If a different integer is already assigned to this key, then this function throws
        a VerifyError.

        The return value indicates whether this key was set by this call.
        """
        if (existing := self._int_variables.get(key)) is not None:
            if existing != i:
                raise VerifyException(
                    f"integer {existing} expected from int variable "
                    f"'{key}', but got {i}"
                )
            return False
        self._int_variables[key] = i
        return True

@deprecated("Please use either VerificationContext or InferenceContext")
class ConstraintContext(VerificationContext):
    pass

@dataclass(frozen=True)
class AttrConstraint(ABC, Generic[AttributeCovT]):
    """Constrain an attribute to a certain value."""

    @abstractmethod
    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        """
        Check if the attribute satisfies the constraint,
        or raise an exception otherwise.
        """
        ...

    def verifies(self, attr: Attribute) -> TypeGuard[AttributeCovT]:
        """
        A helper method to check whether a given attribute matches `self`.
        """
        try:
            self.verify(attr, VerificationContext())
            return True
        except VerifyException:
            return False

    def variables(self) -> set[str]:
        """
        Returns a set of the variables that can be extracted by this constraint.
        These variables are always expected to be set after running `verify`.
        """
        return set()

    def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
        """
        Check if there is enough information to infer the attribute given the
        constraint variables that are already set.
        """
        # By default, we cannot infer anything.
        return False

    def infer(self, context: InferenceContext) -> AttributeCovT:
        """
        Infer the attribute given the the values for all variables.

        Raises an exception if the attribute cannot be inferred. If `can_infer`
        returns `True` with the given constraint variables, this method should
        not raise an exception.
        """
        raise ValueError(f"Cannot infer attribute from constraint {self}")

    def get_bases(self) -> set[type[Attribute]] | None:
        """
        Get a set of base types that can satisfy this constraint, if there exists
        a finite collection, or None otherwise.
        """
        return None

    def relax_constraint(
        self, other: AttrConstraint[AttributeCovT]
    ) -> AttrConstraint[AttributeCovT] | None:
        """
        Attempt to relax this constraint by merging it with `other`,
        returning the result if successful.
        """
        return self if self == other else None

    def __or__(
        self, value: AttrConstraint[_AttributeCovT], /
    ) -> AttrConstraint[AttributeCovT | _AttributeCovT]:
        if isinstance(value, AnyAttr) or self == value:
            return value  # pyright: ignore[reportReturnType]
        return AnyOf.get(self, value)

    def __and__(self, value: AttrConstraint, /) -> AttrConstraint[AttributeCovT]:
        if isinstance(value, AnyAttr) or self == value:
            return self
        return AllOf((self, value))  # pyright: ignore[reportReturnType]

    @abstractmethod
    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> AttrConstraint[AttributeCovT]:
        """
        A helper function to make type vars used in attribute definitions concrete when
        creating constraints for new attributes or operations.
        """
        raise NotImplementedError(
            "Custom constraints must map type vars in nested constraints, if any."
        )

@dataclass(frozen=True)
class AnyAttr(AttrConstraint):
    """Constraint that is verified by all attributes."""

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        pass

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

    def __and__(self, value: AttrConstraint[AttributeCovT], /):
        return value

@dataclass(frozen=True)
class VarConstraint(AttrConstraint[AttributeCovT]):
    """
    Constrain an attribute with the given constraint, and constrain all occurences
    of this constraint (i.e, sharing the same name) to be equal.
    """

    name: str
    """The variable name. All uses of that name refer to the same variable."""

    constraint: AttrConstraint[AttributeCovT]
    """The constraint that the variable must satisfy."""

    @staticmethod
    def get(name: str, constraint: IRDLAttrConstraint[AttributeCovT] = AnyAttr()):
        from xdsl.irdl import irdl_to_attr_constraint

        return VarConstraint(name, irdl_to_attr_constraint(constraint))

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        if constraint_context.set_attr_variable(self.name, attr):
            self.constraint.verify(attr, constraint_context)

    def variables(self) -> set[str]:
        return self.constraint.variables() | {self.name}

    def infer(self, context: InferenceContext) -> AttributeCovT:
        v = context.get_variable(self.name)
        if v is None:
            return self.constraint.infer(context)
        return cast(AttributeCovT, v)

    def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
        return self.name in var_constraint_names or self.constraint.can_infer(
            var_constraint_names
        )

    def get_bases(self) -> set[type[Attribute]] | None:
        return self.constraint.get_bases()

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> VarConstraint[AttributeCovT]:
        return VarConstraint(
            self.name, self.constraint.mapping_type_vars(type_var_mapping)
        )

@dataclass(frozen=True)
class TypeVarConstraint(AttrConstraint):
    """
    Stores the TypeVar instance used to define a generic ParametrizedAttribute.
    """

    type_var: TypeVar
    """The instance of the TypeVar used in the definition."""

    base_constraint: AttrConstraint
    """Constraint inferred from the base of the TypeVar."""

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        self.base_constraint.verify(attr, constraint_context)

    def get_bases(self) -> set[type[Attribute]] | None:
        return self.base_constraint.get_bases()

    def mapping_type_vars(
        self, type_var_mapping: Mapping[TypeVar, AttrConstraint | IntConstraint]
    ) -> AttrConstraint:
        res = type_var_mapping.get(self.type_var)
        if res is None:
            raise KeyError(f"Mapping value missing for type var {self.type_var}")
        if not isinstance(res, AttrConstraint):
            raise ValueError(f"Unexpected constraint {res} for TypeVar {self.type_var}")
        return res

@dataclass(frozen=True)
class EqAttrConstraint(AttrConstraint[AttributeCovT], Generic[AttributeCovT]):
    """Constrain an attribute to be equal to another attribute."""

    attr: AttributeCovT
    """The attribute we want to check equality with."""

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        if attr != self.attr:
            raise VerifyException(f"Expected attribute {self.attr} but got {attr}")

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

    def infer(self, context: InferenceContext) -> AttributeCovT:
        return self.attr

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

    def relax_constraint(
        self, other: AttrConstraint[AttributeCovT]
    ) -> AttrConstraint[AttributeCovT] | None:
        if isinstance(other, AttrSetConstraint):
            return AttrSetConstraint.get(self.attr, *other.values)
        if isinstance(other, EqAttrConstraint):
            return AttrSetConstraint.get(self.attr, other.attr)

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

@dataclass(frozen=True)
class AttrSetConstraint(AttrConstraint[AttributeCovT], Generic[AttributeCovT]):
    """Constrain an attribute to be one of a set of attributes."""

    values: frozenset[AttributeCovT]

    def __repr__(self) -> str:
        values = ", ".join(sorted(repr(value) for value in self.values))
        return f"AttrSetConstraint(frozenset([{values}]))"

    @staticmethod
    def get(*values: AttributeInvT) -> AttrConstraint[AttributeInvT]:
        s = frozenset(values)
        if len(s) == 1:
            return EqAttrConstraint(values[0])
        return AttrSetConstraint(s)

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        if attr not in self.values:
            raise VerifyException(
                f"Expected one of {', '.join(sorted(str(value) for value in self.values))}, but got {attr}"
            )

    def get_bases(self) -> set[type[Attribute]] | None:
        return {type(attr) for attr in self.values}

    def relax_constraint(
        self, other: AttrConstraint[AttributeCovT]
    ) -> AttrConstraint[AttributeCovT] | None:
        if isinstance(other, AttrSetConstraint):
            return AttrSetConstraint.get(*self.values, *other.values)
        if isinstance(other, EqAttrConstraint):
            return AttrSetConstraint.get(*self.values, other.attr)

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

@dataclass(frozen=True)
class BaseAttr(AttrConstraint[AttributeCovT], Generic[AttributeCovT]):
    """Constrain an attribute to be of a given base type."""

    attr: type[AttributeCovT]
    """The expected attribute base type."""

    def __repr__(self):
        return f"BaseAttr({self.attr.__name__})"

    def verify(
        self,
        attr: Attribute,
        constraint_context: VerificationContext,
    ) -> None:
        if not isinstance(attr, self.attr):
            if hasattr(self.attr, "name"):
                raise VerifyException(
                    f"{attr} should be of base attribute {self.attr.name}"
                )
            else:
                raise VerifyException(
                    f"{attr} should be of attribute subclassing `{self.attr.__name__}`"
                )

    def can_infer(self, var_constraint_names: AbstractSet[str]) -> bool:
        return (
            is_runtime_final(self.attr)
            and issubclass(self.attr, ParametrizedAttribute)
            and not self.attr.get_irdl_definition().parameters
        )

    def infer(self, context: InferenceContext) -> AttributeCovT:
        assert issubclass(self.attr, ParametrizedAttribute)
        attr = self.attr.new(())
        return attr

    def get_bases(self) -> set[type[Attribute]] | None:
        if is_runtime_final(self.attr):
            return {self.attr}
        return None

    def relax_constraint(
        self, other: AttrConstraint[AttributeCovT]
    ) -> AttrConstraint[AttributeCovT] | None:
        if not isinstance(other, BaseAttr):
            # Trying to match on ParamAttrConstraint does strange things to pyright
            # so we just appeal to symmetry instead.
            return other.relax_constraint(self)
        return super().relax_constraint(other)

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