api: remove parameters field from Trait

tests/test_traits.py

@@ -118,8 +118,8 @@ def test_has_trait_object():
     """
     assert TestOp.has_trait(LargerOperandTrait)
     assert not TestOp.has_trait(LargerResultTrait)
-    assert not TestOp.has_trait(BitwidthSumLessThanTrait, 0)
-    assert TestOp.has_trait(BitwidthSumLessThanTrait, 64)
+    assert not TestOp.has_trait(BitwidthSumLessThanTrait(0))
+    assert TestOp.has_trait(BitwidthSumLessThanTrait(64))
 
 
 def test_get_traits_of_type():
@@ -445,8 +445,8 @@ def test_lazy_parent():
     """Test the trait infrastructure for an operation that defines a trait "lazily"."""
     op = HasLazyParentOp.create()
     assert len(op.get_traits_of_type(HasParent)) != 0
-    assert op.get_traits_of_type(HasParent)[0].parameters == (TestOp,)
-    assert op.has_trait(HasParent, (TestOp,))
+    assert op.get_traits_of_type(HasParent)[0].op_types == (TestOp,)
+    assert op.has_trait(HasParent(TestOp))
     assert op.traits == frozenset([HasParent(TestOp)])
 
 
@@ -461,7 +461,7 @@ def test_has_ancestor():
     op = AncestorOp()
 
     assert op.get_traits_of_type(HasAncestor) == [HasAncestor(TestOp)]
-    assert op.has_trait(HasAncestor, (TestOp,))
+    assert op.has_trait(HasAncestor(TestOp))
 
     with pytest.raises(
         VerifyException, match="'test.ancestor' expects ancestor op 'test.test'"

xdsl/dialects/csl/csl.py

@@ -11,7 +11,7 @@
 
 from abc import ABC
 from collections.abc import Sequence
-from dataclasses import dataclass
+from dataclasses import KW_ONLY, dataclass, field
 from typing import Annotated, ClassVar, TypeAlias
 
 from xdsl.dialects import builtin
@@ -186,14 +186,15 @@ class InModuleKind(OpTrait):
     Ops with this trait are always allowed inside a csl_wrapper.module
     """
 
-    def __init__(self, kind: ModuleKind, *, direct_child: bool = True):
-        super().__init__((kind, direct_child))
+    kind: ModuleKind = field()
+    _: KW_ONLY
+    direct_child: bool = field(default=True)
 
     def verify(self, op: Operation) -> None:
         from xdsl.dialects.csl import csl_wrapper
 
-        kind: ModuleKind = self.parameters[0]
-        direct_child: bool = self.parameters[1]
+        kind: ModuleKind = self.kind
+        direct_child: bool = self.direct_child
 
         direct = "direct" if direct_child else "indirect"
         parent_module = op.parent_op()

xdsl/dialects/csl/csl_stencil.py

@@ -490,7 +490,7 @@ def verify_(self) -> None:
 
         # As promised by HasAncestor(ApplyOp)
         trait = cast(
-            HasAncestor, AccessOp.get_trait(HasAncestor, (stencil.ApplyOp, ApplyOp))
+            HasAncestor, AccessOp.get_trait(HasAncestor(stencil.ApplyOp, ApplyOp))
         )
         apply = trait.get_ancestor(self)
         assert isinstance(apply, stencil.ApplyOp | ApplyOp)
@@ -529,7 +529,8 @@ def get_apply(self) -> stencil.ApplyOp | ApplyOp:
         Simple helper to get the parent apply and raise otherwise.
         """
         trait = cast(
-            HasAncestor, self.get_trait(HasAncestor, (stencil.ApplyOp, ApplyOp))
+            HasAncestor,
+            self.get_trait(HasAncestor(stencil.ApplyOp, ApplyOp)),
         )
         ancestor = trait.get_ancestor(self)
         if ancestor is None:

xdsl/dialects/stencil.py

@@ -802,7 +802,8 @@ def get_apply(self):
         """
         Simple helper to get the parent apply and raise otherwise.
         """
-        trait = cast(HasAncestor, self.get_trait(HasAncestor, (ApplyOp,)))
+        trait = self.get_trait(HasAncestor(ApplyOp))
+        assert trait is not None
         ancestor = trait.get_ancestor(self)
         if ancestor is None:
             raise ValueError(
@@ -861,8 +862,9 @@ def print(self, printer: Printer):
             print_keyword=True,
         )
 
-        # IRDL-enforced, not supposed to use custom syntax if not veriied
-        trait = cast(HasAncestor, AccessOp.get_trait(HasAncestor, (ApplyOp,)))
+        # IRDL-enforced, not supposed to use custom syntax if not verified
+        trait = AccessOp.get_trait(HasAncestor(ApplyOp))
+        assert trait is not None
         apply = cast(ApplyOp, trait.get_ancestor(self))
 
         mapping = self.offset_mapping
@@ -949,7 +951,8 @@ def get(
 
     def verify_(self) -> None:
         # As promised by HasAncestor(ApplyOp)
-        trait = cast(HasAncestor, AccessOp.get_trait(HasAncestor, (ApplyOp,)))
+        trait = AccessOp.get_trait(HasAncestor(ApplyOp))
+        assert trait is not None
         apply = trait.get_ancestor(self)
         assert isinstance(apply, ApplyOp)
 
@@ -1002,7 +1005,8 @@ def get_apply(self):
         """
         Simple helper to get the parent apply and raise otherwise.
         """
-        trait = cast(HasAncestor, self.get_trait(HasAncestor, (ApplyOp,)))
+        trait = self.get_trait(HasAncestor(ApplyOp))
+        assert trait is not None
         ancestor = trait.get_ancestor(self)
         if ancestor is None:
             raise ValueError(

xdsl/ir/core.py

@@ -939,8 +939,8 @@ def clone(
     @classmethod
     def has_trait(
         cls,
-        trait: type[OpTrait],
-        parameters: Any = None,
+        trait: type[OpTrait] | OpTrait,
+        *,
         value_if_unregistered: bool = True,
     ) -> bool:
         """
@@ -953,18 +953,21 @@ def has_trait(
         if issubclass(cls, UnregisteredOp):
             return value_if_unregistered
 
-        return cls.get_trait(trait, parameters) is not None
+        return cls.get_trait(trait) is not None
 
     @classmethod
-    def get_trait(
-        cls, trait: type[OpTraitInvT], parameters: Any = None
-    ) -> OpTraitInvT | None:
+    def get_trait(cls, trait: type[OpTraitInvT] | OpTraitInvT) -> OpTraitInvT | None:
         """
         Return a trait with the given type and parameters, if it exists.
         """
-        for t in cls.traits:
-            if isinstance(t, trait) and t.parameters == parameters:
-                return t
+        if isinstance(trait, type):
+            for t in cls.traits:
+                if isinstance(t, trait):
+                    return t
+        else:
+            for t in cls.traits:
+                if t == trait:
+                    return cast(OpTraitInvT, t)
         return None
 
     @classmethod

xdsl/traits.py

@@ -2,9 +2,9 @@
 
 import abc
 from collections.abc import Iterator
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, TypeVar, final
+from typing import TYPE_CHECKING, TypeVar, final
 
 from xdsl.utils.exceptions import VerifyException
 
@@ -20,13 +20,9 @@ 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.
-    Traits have parameters, which by default is just the `None` value. Parameters should
-    always be comparable and hashable.
     Note that traits are the merge of traits and interfaces in MLIR.
     """
 
-    parameters: Any = field(default=None)
-
     def verify(self, op: Operation) -> None:
         """Check that the operation satisfies the trait requirements."""
         pass
@@ -47,22 +43,20 @@ class ConstantLike(OpTrait):
 class HasParent(OpTrait):
     """Constraint the operation to have a specific parent operation."""
 
-    parameters: tuple[type[Operation], ...]
+    op_types: tuple[type[Operation], ...]
 
-    def __init__(self, *parameters: type[Operation]):
-        if not parameters:
-            raise ValueError("parameters must not be empty")
-        super().__init__(parameters)
+    def __init__(self, head_param: type[Operation], *tail_params: type[Operation]):
+        object.__setattr__(self, "op_types", (head_param, *tail_params))
 
     def verify(self, op: Operation) -> None:
         parent = op.parent_op()
-        if isinstance(parent, self.parameters):
+        if isinstance(parent, self.op_types):
             return
-        if len(self.parameters) == 1:
+        if len(self.op_types) == 1:
             raise VerifyException(
-                f"'{op.name}' expects parent op '{self.parameters[0].name}'"
+                f"'{op.name}' expects parent op '{self.op_types[0].name}'"
             )
-        names = ", ".join(f"'{p.name}'" for p in self.parameters)
+        names = ", ".join(f"'{p.name}'" for p in self.op_types)
         raise VerifyException(f"'{op.name}' expects parent op to be one of {names}")
 
 
@@ -73,18 +67,18 @@ class HasAncestor(OpTrait):
     parent.
     """
 
-    parameters: tuple[type[Operation], ...]
+    op_types: tuple[type[Operation], ...]
 
     def __init__(self, head_param: type[Operation], *tail_params: type[Operation]):
-        super().__init__((head_param, *tail_params))
+        object.__setattr__(self, "op_types", (head_param, *tail_params))
 
     def verify(self, op: Operation) -> None:
         if self.get_ancestor(op) is None:
-            if len(self.parameters) == 1:
+            if len(self.op_types) == 1:
                 raise VerifyException(
-                    f"'{op.name}' expects ancestor op '{self.parameters[0].name}'"
+                    f"'{op.name}' expects ancestor op '{self.op_types[0].name}'"
                 )
-            names = ", ".join(f"'{p.name}'" for p in self.parameters)
+            names = ", ".join(f"'{p.name}'" for p in self.op_types)
             raise VerifyException(
                 f"'{op.name}' expects ancestor op to be one of {names}"
             )
@@ -98,7 +92,7 @@ def walk_ancestors(self, op: Operation) -> Iterator[Operation]:
 
     def get_ancestor(self, op: Operation) -> Operation | None:
         ancestors = self.walk_ancestors(op)
-        matching_ancestors = (a for a in ancestors if isinstance(a, self.parameters))
+        matching_ancestors = (a for a in ancestors if isinstance(a, self.op_types))
         return next(matching_ancestors, None)
 
 
@@ -133,6 +127,7 @@ def verify(self, op: Operation) -> None:
                 )
 
 
+@dataclass(frozen=True)
 class SingleBlockImplicitTerminator(OpTrait):
     """
     Checks the existence of the specified terminator to an operation which has
@@ -144,7 +139,7 @@ class SingleBlockImplicitTerminator(OpTrait):
     https://mlir.llvm.org/docs/Traits/#single-block-with-implicit-terminator
     """
 
-    parameters: type[Operation]
+    op_type: type[Operation]
 
     def verify(self, op: Operation) -> None:
         for region in op.regions:
@@ -156,13 +151,13 @@ def verify(self, op: Operation) -> None:
                 if (last_op := block.last_op) is None:
                     raise VerifyException(
                         f"'{op.name}' contains empty block instead of at least "
-                        f"terminating with {self.parameters.name}"
+                        f"terminating with {self.op_type.name}"
                     )
 
-                if not isinstance(last_op, self.parameters):
+                if not isinstance(last_op, self.op_type):
                     raise VerifyException(
                         f"'{op.name}' terminates with operation {last_op.name} "
-                        f"instead of {self.parameters.name}"
+                        f"instead of {self.op_type.name}"
                     )
 
 
@@ -181,11 +176,11 @@ def ensure_terminator(op: Operation, trait: SingleBlockImplicitTerminator) -> No
             if (
                 (last_op := block.last_op) is not None
                 and last_op.has_trait(IsTerminator)
-                and not isinstance(last_op, trait.parameters)
+                and not isinstance(last_op, trait.op_type)
             ):
                 raise VerifyException(
                     f"'{op.name}' terminates with operation {last_op.name} "
-                    f"instead of {trait.parameters.name}"
+                    f"instead of {trait.op_type.name}"
                 )
 
     from xdsl.builder import ImplicitBuilder
@@ -200,7 +195,7 @@ def ensure_terminator(op: Operation, trait: SingleBlockImplicitTerminator) -> No
                 IsTerminator
             ):
                 with ImplicitBuilder(block):
-                    trait.parameters.create()
+                    trait.op_type.create()
 
 
 class IsolatedFromAbove(OpTrait):