add StmtTrait and AttrTrait (#359)

Author: Roger-luo

src/kirin/ir/__init__.py

@@ -24,8 +24,10 @@
 from kirin.ir.traits import (
     Pure as Pure,
     Trait as Trait,
+    AttrTrait as AttrTrait,
     HasParent as HasParent,
     MaybePure as MaybePure,
+    StmtTrait as StmtTrait,
     RegionTrait as RegionTrait,
     SymbolTable as SymbolTable,
     ConstantLike as ConstantLike,

src/kirin/ir/traits/__init__.py

@@ -9,7 +9,12 @@
 `IsolatedFromAbove`.
 """
 
-from .abc import Trait as Trait, RegionTrait as RegionTrait
+from .abc import (
+    Trait as Trait,
+    AttrTrait as AttrTrait,
+    StmtTrait as StmtTrait,
+    RegionTrait as RegionTrait,
+)
 from .basic import (
     Pure as Pure,
     HasParent as HasParent,

src/kirin/ir/traits/abc.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 
 if TYPE_CHECKING:
-    from kirin.ir import Block, Region, Statement  # noqa: F401
+    from kirin.ir import Block, Region, Attribute, Statement
     from kirin.graph import Graph
 
 
@@ -20,11 +20,27 @@ def verify(self, node: IRNodeType):
         pass
 
 
+@dataclass(frozen=True)
+class AttrTrait(Trait["Attribute"]):
+    """Base class for all attribute traits."""
+
+    def verify(self, node: "Attribute"):
+        pass
+
+
+@dataclass(frozen=True)
+class StmtTrait(Trait["Statement"], ABC):
+    """Base class for all statement traits."""
+
+    def verify(self, node: "Statement"):
+        pass
+
+
 GraphType = TypeVar("GraphType", bound="Graph[Block]")
 
 
 @dataclass(frozen=True)
-class RegionTrait(Trait["Statement"], Generic[GraphType]):
+class RegionTrait(StmtTrait, Generic[GraphType]):
     """A trait that indicates the properties of the statement's region."""
 
     @abstractmethod

src/kirin/ir/traits/basic.py

@@ -3,14 +3,14 @@
 
 from kirin.ir.nodes.stmt import Statement
 
-from .abc import Trait
+from .abc import StmtTrait
 
 if TYPE_CHECKING:
     from kirin.ir import Statement
 
 
 @dataclass(frozen=True)
-class Pure(Trait["Statement"]):
+class Pure(StmtTrait):
     """A trait that indicates that a statement is pure, i.e., it has no side
     effects.
     """
@@ -19,7 +19,7 @@ class Pure(Trait["Statement"]):
 
 
 @dataclass(frozen=True)
-class MaybePure(Trait["Statement"]):
+class MaybePure(StmtTrait):
     """A trait that indicates the statement may be pure,
     i.e., a call statement can be pure if the callee is pure.
     """
@@ -48,7 +48,7 @@ def set_pure(cls, stmt: "Statement") -> None:
 
 
 @dataclass(frozen=True)
-class ConstantLike(Trait["Statement"]):
+class ConstantLike(StmtTrait):
     """A trait that indicates that a statement is constant-like, i.e., it
     represents a constant value.
     """
@@ -57,7 +57,7 @@ class ConstantLike(Trait["Statement"]):
 
 
 @dataclass(frozen=True)
-class IsTerminator(Trait["Statement"]):
+class IsTerminator(StmtTrait):
     """A trait that indicates that a statement is a terminator, i.e., it
     terminates a block.
     """
@@ -66,19 +66,19 @@ class IsTerminator(Trait["Statement"]):
 
 
 @dataclass(frozen=True)
-class NoTerminator(Trait["Statement"]):
+class NoTerminator(StmtTrait):
     """A trait that indicates that the region of a statement has no terminator."""
 
     pass
 
 
 @dataclass(frozen=True)
-class IsolatedFromAbove(Trait["Statement"]):
+class IsolatedFromAbove(StmtTrait):
     pass
 
 
 @dataclass(frozen=True)
-class HasParent(Trait["Statement"]):
+class HasParent(StmtTrait):
     """A trait that indicates that a statement has a parent
     statement.
     """

src/kirin/ir/traits/callable.py

@@ -2,7 +2,7 @@
 from typing import TYPE_CHECKING, Generic, TypeVar
 from dataclasses import dataclass
 
-from kirin.ir.traits.abc import Trait
+from kirin.ir.traits.abc import StmtTrait
 
 if TYPE_CHECKING:
     from kirin.ir import Region, Statement
@@ -12,7 +12,7 @@
 
 
 @dataclass(frozen=True)
-class CallableStmtInterface(Trait["Statement"], ABC, Generic[StmtType]):
+class CallableStmtInterface(StmtTrait, ABC, Generic[StmtType]):
     """A trait that indicates that a statement is a callable statement.
 
     A callable statement is a statement that can be called as a function.
@@ -26,7 +26,7 @@ def get_callable_region(cls, stmt: "Statement") -> "Region":
 
 
 @dataclass(frozen=True)
-class HasSignature(Trait["Statement"], ABC):
+class HasSignature(StmtTrait, ABC):
     """A trait that indicates that a statement has a function signature
     attribute.
     """

src/kirin/ir/traits/symbol.py

@@ -5,14 +5,14 @@
 
 from kirin.ir.attrs.py import PyAttr
 from kirin.ir.exception import ValidationError
-from kirin.ir.traits.abc import Trait
+from kirin.ir.traits.abc import StmtTrait
 
 if TYPE_CHECKING:
     from kirin.ir import Statement
 
 
 @dataclass(frozen=True)
-class SymbolOpInterface(Trait["Statement"]):
+class SymbolOpInterface(StmtTrait):
     """A trait that indicates that a statement is a symbol operation.
 
     A symbol operation is a statement that has a symbol name attribute.
@@ -34,7 +34,7 @@ def verify(self, stmt: Statement):
 
 
 @dataclass(frozen=True)
-class SymbolTable(Trait["Statement"]):
+class SymbolTable(StmtTrait):
     """
     Statement with SymbolTable trait can only have one region with one block.
     """