kirin-toolchain 0.13.0__py3-none-any.whl

kirin/ir/nodes/view.py

@@ -0,0 +1,142 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar, Iterator, Sequence, overload
+from dataclasses import dataclass
+
+from typing_extensions import Self
+
+from kirin.ir.ssa import SSAValue
+from kirin.ir.nodes.base import IRNode
+
+ElemType = TypeVar("ElemType", bound=IRNode | SSAValue)
+FieldType = TypeVar("FieldType", bound=Sequence)
+NodeType = TypeVar("NodeType", bound=IRNode | SSAValue)
+
+
+@dataclass
+class View(ABC, Generic[NodeType, ElemType]):
+    node: NodeType
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[ElemType]: ...
+
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+    def __bool__(self) -> bool:
+        return bool(len(self))
+
+    def append(self, value: ElemType) -> None:
+        raise NotImplementedError
+
+    def extend(self, values: Sequence[ElemType]) -> None:
+        for value in values:
+            self.append(value)
+
+    def __reversed__(self) -> Iterator[ElemType]:
+        raise NotImplementedError
+
+
+@dataclass
+class SequenceView(Generic[FieldType, NodeType, ElemType], View[NodeType, ElemType]):
+    field: FieldType
+
+    @classmethod
+    def similar(cls, node: NodeType, field: FieldType) -> Self:
+        return cls(node, field)
+
+    def __iter__(self) -> Iterator[ElemType]:
+        return iter(self.field)
+
+    def __len__(self) -> int:
+        return len(self.field)
+
+    def __reversed__(self) -> Iterator[ElemType]:
+        return reversed(self.field)
+
+    def isempty(self) -> bool:
+        return len(self) == 0
+
+    def __bool__(self) -> bool:
+        return not self.isempty()
+
+    # optional interface
+    @overload
+    def __getitem__(self, idx: int) -> ElemType: ...
+
+    @overload
+    def __getitem__(self, idx: slice) -> Self: ...
+
+    def __getitem__(self, idx: int | slice) -> ElemType | Self:
+        if isinstance(idx, slice):
+            x: FieldType = self.field[idx]  # type: ignore
+            return self.similar(self.node, x)
+        else:
+            return self.field[idx]
+
+
+@dataclass
+class MutableSequenceView(SequenceView[FieldType, NodeType, ElemType]):
+    @overload
+    def __setitem__(self, idx: int, value: ElemType) -> None: ...
+
+    @overload
+    def __setitem__(self, idx: slice, value: Sequence[ElemType]) -> None: ...
+
+    def __setitem__(
+        self, idx: int | slice, value: ElemType | Sequence[ElemType]
+    ) -> None:
+        if isinstance(idx, int) and not isinstance(value, Sequence):
+            return self.set_item(idx, value)
+        elif isinstance(idx, slice):
+            assert isinstance(value, Sequence), "Expected sequence of values"
+            if idx.step is not None:  # no need to support step
+                raise ValueError("Slice step is not supported")
+            return self.set_item_slice(idx, value)
+        else:
+            raise TypeError("Expected int or slice")
+
+    def set_item(self, idx: int, value: ElemType) -> None:
+        raise NotImplementedError
+
+    def set_item_slice(self, s: slice, value: Sequence[ElemType]) -> None:
+        # replace the view of slice
+        for idx in range(s.start, s.stop):
+            if idx < len(value):
+                self.set_item(idx, value[idx])
+            else:
+                del self[idx]
+
+        # insert the rest of the values
+        for idx, v in enumerate(value[s.stop - s.start :]):
+            self.insert(idx + s.stop, v)
+
+    def __delitem__(self, idx: int) -> None:
+        raise NotImplementedError
+
+    def insert(self, idx: int, value: ElemType) -> None:
+        raise NotImplementedError
+
+    def pop(self, idx: int = -1) -> ElemType:
+        item = self.field[idx]
+        del self[idx]
+        return item
+
+    def poplast(self) -> ElemType | None:
+        """Pop the last element from the view.
+
+        Returns:
+            The last element in the view.
+        """
+        if self:
+            return self.pop(-1)
+        return None
+
+    def popfirst(self) -> ElemType | None:
+        """Pop the first element from the view.
+
+        Returns:
+            The first element in the view.
+        """
+        if self:
+            return self.pop(0)
+        return None

kirin/ir/ssa.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, ClassVar
+from dataclasses import field, dataclass
+
+from typing_extensions import Self
+
+from kirin.print import Printer, Printable
+from kirin.ir.attrs.abc import Attribute
+from kirin.ir.attrs.types import AnyType, TypeAttribute
+
+if TYPE_CHECKING:
+    from kirin.ir.use import Use
+    from kirin.ir.nodes.stmt import Statement
+    from kirin.ir.nodes.block import Block
+
+
+@dataclass
+class SSAValue(ABC, Printable):
+    """Base class for all SSA values in the IR."""
+
+    type: TypeAttribute = field(default_factory=AnyType, init=False, repr=True)
+    """The type of this SSA value."""
+    hints: dict[str, Attribute] = field(default_factory=dict, init=False, repr=False)
+    """Hints for this SSA value."""
+    uses: set[Use] = field(init=False, default_factory=set, repr=False)
+    """The uses of this SSA value."""
+    _name: str | None = field(init=False, default=None, repr=True)
+    """The name of this SSA value."""
+    name_pattern: ClassVar[re.Pattern[str]] = re.compile(r"([A-Za-z_$.-][\w$.-]*)")
+    """The pattern that the name of this SSA value must match."""
+
+    @property
+    @abstractmethod
+    def owner(self) -> Statement | Block:
+        """The object that owns this SSA value."""
+        ...
+
+    @property
+    def name(self) -> str | None:
+        """The name of this SSA value."""
+        return self._name
+
+    @name.setter
+    def name(self, name: str | None) -> None:
+        if name and not self.name_pattern.fullmatch(name):
+            raise ValueError(f"Invalid name: {name}")
+        self._name = name
+
+    def __repr__(self) -> str:
+        if self.name:
+            return f"{type(self).__name__}({self.name})"
+        return f"{type(self).__name__}({id(self)})"
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def add_use(self, use: Use) -> Self:
+        """Add a use to this SSA value."""
+        self.uses.add(use)
+        return self
+
+    def remove_use(self, use: Use) -> Self:
+        """Remove a use from this SSA value."""
+        # print(use)
+        # assert use in self.uses, "Use not found"
+        if use in self.uses:
+            self.uses.remove(use)
+        return self
+
+    def replace_by(self, other: SSAValue) -> None:
+        """Replace this SSA value with another SSA value. Update all uses."""
+        for use in self.uses.copy():
+            use.stmt.args[use.index] = other
+
+        if other.name is None and self.name is not None:
+            other.name = self.name
+
+        assert len(self.uses) == 0, "Uses not empty"
+
+    # TODO: also delete BlockArgument from arglist
+    def delete(self, safe: bool = True) -> None:
+        """Delete this SSA value. If `safe` is `True`, raise an error if there are uses."""
+        if safe and len(self.uses) > 0:
+            raise ValueError("Cannot delete SSA value with uses")
+        self.replace_by(DeletedSSAValue(self))
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.plain_print(printer.state.ssa_id[self])
+
+
+@dataclass
+class ResultValue(SSAValue):
+    """SSAValue that is a result of a [`Statement`][kirin.ir.nodes.stmt.Statement]."""
+
+    stmt: Statement = field(init=False)
+    """The statement that this value is a result of."""
+    index: int = field(init=False)
+    """The index of this value in the statement's result list."""
+
+    # NOTE: we will assign AnyType unless specified.
+    #       when SSAValue is a ResultValue, the type is inferred
+    #       later in the compilation process.
+    def __init__(
+        self, stmt: Statement, index: int, type: TypeAttribute | None = None
+    ) -> None:
+        super().__init__()
+        self.type = type or AnyType()
+        self.stmt = stmt
+        self.index = index
+
+    @property
+    def owner(self) -> Statement:
+        return self.stmt
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def __repr__(self) -> str:
+        if self.type is self.type.top():
+            type_str = ""
+        else:
+            type_str = f"[{self.type}]"
+
+        if self.name:
+            return (
+                f"<{type(self).__name__}{type_str} {self.name}, uses: {len(self.uses)}>"
+            )
+        return f"<{type(self).__name__}{type_str} stmt: {self.stmt.name}, uses: {len(self.uses)}>"
+
+
+@dataclass
+class BlockArgument(SSAValue):
+    """SSAValue that is an argument to a [`Block`][kirin.ir.Block]."""
+
+    block: Block = field(init=False)
+    """The block that this argument belongs to."""
+    index: int = field(init=False)
+    """The index of this argument in the block's argument list."""
+
+    def __init__(
+        self, block: Block, index: int, type: TypeAttribute = AnyType()
+    ) -> None:
+        super().__init__()
+        self.type = type
+        self.block = block
+        self.index = index
+
+    @property
+    def owner(self) -> Block:
+        return self.block
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def __repr__(self) -> str:
+        if self.name:
+            return f"<{type(self).__name__}[{self.type}] {self.name}, uses: {len(self.uses)}>"
+        return f"<{type(self).__name__}[{self.type}] index: {self.index}, uses: {len(self.uses)}>"
+
+    def print_impl(self, printer: Printer) -> None:
+        super().print_impl(printer)
+        if not isinstance(self.type, AnyType):
+            with printer.rich(style="comment"):
+                printer.plain_print(" : ")
+                printer.print(self.type)
+
+
+@dataclass
+class DeletedSSAValue(SSAValue):
+    value: SSAValue = field(init=False)
+
+    def __init__(self, value: SSAValue) -> None:
+        super().__init__()
+        self.value = value
+        self.type = value.type
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__}[{self.type}] value: {self.value}, uses: {len(self.uses)}>"
+
+    @property
+    def owner(self) -> Statement | Block:
+        return self.value.owner
+
+
+@dataclass
+class TestValue(SSAValue):
+    """Test SSAValue for testing IR construction."""
+
+    def __init__(self, type: TypeAttribute = AnyType()) -> None:
+        super().__init__()
+        self.type = type
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    @property
+    def owner(self) -> Statement | Block:
+        raise NotImplementedError

kirin/ir/traits/__init__.py

@@ -0,0 +1,36 @@
+"""Kirin IR Traits.
+
+This module defines the traits that can be used to define the behavior of
+Kirin IR nodes. The base trait is `StmtTrait`, which is a `dataclass` that
+implements the `__hash__` and `__eq__` methods.
+
+There are also some basic traits that are provided for convenience, such as
+`Pure`, `HasParent`, `ConstantLike`, `IsTerminator`, `NoTerminator`, and
+`IsolatedFromAbove`.
+"""
+
+from .abc import (
+    StmtTrait as StmtTrait,
+    RegionTrait as RegionTrait,
+    PythonLoweringTrait as PythonLoweringTrait,
+)
+from .basic import (
+    Pure as Pure,
+    HasParent as HasParent,
+    MaybePure as MaybePure,
+    ConstantLike as ConstantLike,
+    IsTerminator as IsTerminator,
+    NoTerminator as NoTerminator,
+    IsolatedFromAbove as IsolatedFromAbove,
+)
+from .symbol import SymbolTable as SymbolTable, SymbolOpInterface as SymbolOpInterface
+from .callable import (
+    HasSignature as HasSignature,
+    CallableStmtInterface as CallableStmtInterface,
+)
+from .lowering.call import FromPythonCall as FromPythonCall
+from .region.ssacfg import SSACFGRegion as SSACFGRegion
+from .lowering.context import (
+    FromPythonWith as FromPythonWith,
+    FromPythonWithSingleItem as FromPythonWithSingleItem,
+)

kirin/ir/traits/abc.py

@@ -0,0 +1,42 @@
+import ast
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, TypeVar
+from dataclasses import dataclass
+
+if TYPE_CHECKING:
+    from kirin import lowering
+    from kirin.ir import Block, Region, Statement
+    from kirin.graph import Graph
+
+
+@dataclass(frozen=True)
+class StmtTrait(ABC):
+    """Base class for all statement traits."""
+
+    def verify(self, stmt: "Statement"):
+        pass
+
+
+GraphType = TypeVar("GraphType", bound="Graph[Block]")
+
+
+@dataclass(frozen=True)
+class RegionTrait(StmtTrait, Generic[GraphType]):
+    """A trait that indicates the properties of the statement's region."""
+
+    @abstractmethod
+    def get_graph(self, region: "Region") -> GraphType: ...
+
+
+ASTNode = TypeVar("ASTNode", bound=ast.AST)
+StatementType = TypeVar("StatementType", bound="Statement")
+
+
+@dataclass(frozen=True)
+class PythonLoweringTrait(StmtTrait, Generic[StatementType, ASTNode]):
+    """A trait that indicates that a statement can be lowered from Python AST."""
+
+    @abstractmethod
+    def lower(
+        self, stmt: type[StatementType], state: "lowering.LoweringState", node: ASTNode
+    ) -> "lowering.Result": ...

kirin/ir/traits/basic.py

@@ -0,0 +1,78 @@
+from typing import TYPE_CHECKING
+from dataclasses import dataclass
+
+from .abc import StmtTrait
+
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+
+
+@dataclass(frozen=True)
+class Pure(StmtTrait):
+    """A trait that indicates that a statement is pure, i.e., it has no side
+    effects.
+    """
+
+    pass
+
+
+@dataclass(frozen=True)
+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.
+    """
+
+    @classmethod
+    def is_pure(cls, stmt: "Statement") -> bool:
+        # TODO: simplify this after removing property
+        from kirin.ir.attrs.py import PyAttr
+
+        purity = stmt.attributes.get("purity")
+        if isinstance(purity, PyAttr) and purity.data:
+            return True
+        return False
+
+    @classmethod
+    def set_pure(cls, stmt: "Statement") -> None:
+        from kirin.ir.attrs.py import PyAttr
+
+        stmt.attributes["purity"] = PyAttr(True)
+
+
+@dataclass(frozen=True)
+class ConstantLike(StmtTrait):
+    """A trait that indicates that a statement is constant-like, i.e., it
+    represents a constant value.
+    """
+
+    pass
+
+
+@dataclass(frozen=True)
+class IsTerminator(StmtTrait):
+    """A trait that indicates that a statement is a terminator, i.e., it
+    terminates a block.
+    """
+
+    pass
+
+
+@dataclass(frozen=True)
+class NoTerminator(StmtTrait):
+    """A trait that indicates that the region of a statement has no terminator."""
+
+    pass
+
+
+@dataclass(frozen=True)
+class IsolatedFromAbove(StmtTrait):
+    pass
+
+
+@dataclass(frozen=True)
+class HasParent(StmtTrait):
+    """A trait that indicates that a statement has a parent
+    statement.
+    """
+
+    parents: tuple[type["Statement"]]

kirin/ir/traits/callable.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, TypeVar
+from dataclasses import dataclass
+
+from kirin.ir.traits.abc import StmtTrait
+
+if TYPE_CHECKING:
+    from kirin.ir import Region, Statement
+    from kirin.dialects.func.attrs import Signature
+
+StmtType = TypeVar("StmtType", bound="Statement")
+
+
+@dataclass(frozen=True)
+class CallableStmtInterface(StmtTrait, 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.
+    """
+
+    @classmethod
+    @abstractmethod
+    def get_callable_region(cls, stmt: "StmtType") -> "Region":
+        """Returns the body of the callable region"""
+        ...
+
+
+@dataclass(frozen=True)
+class HasSignature(StmtTrait, ABC):
+    """A trait that indicates that a statement has a function signature
+    attribute.
+    """
+
+    @classmethod
+    def get_signature(cls, stmt: "Statement"):
+        signature: Signature | None = stmt.attributes.get("signature")  # type: ignore
+        if signature is None:
+            raise ValueError(f"Statement {stmt.name} does not have a function type")
+
+        return signature
+
+    @classmethod
+    def set_signature(cls, stmt: "Statement", signature: "Signature"):
+        stmt.attributes["signature"] = signature
+
+    def verify(self, stmt: "Statement"):
+        from kirin.dialects.func.attrs import Signature
+
+        signature = self.get_signature(stmt)
+        if not isinstance(signature, Signature):
+            raise ValueError(f"{signature} is not a Signature attribute")

kirin/ir/traits/lowering/__init__.py

@@ -0,0 +1,2 @@
+"""Lowering traits for the IR.
+"""

kirin/ir/traits/lowering/call.py

@@ -0,0 +1,37 @@
+import ast
+from typing import TYPE_CHECKING, TypeVar
+from dataclasses import dataclass
+
+from ..abc import PythonLoweringTrait
+
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+    from kirin.lowering import Result, LoweringState
+
+
+StatementType = TypeVar("StatementType", bound="Statement")
+
+
+@dataclass(frozen=True)
+class FromPythonCall(PythonLoweringTrait[StatementType, ast.Call]):
+    """Trait for customizing lowering of Python calls to a statement.
+
+    Declared in a statement definition to indicate that the statement can be
+    constructed from a Python call (i.e., a function call `ast.Call` in the
+    Python AST).
+
+    Subclassing this trait allows for customizing the lowering of Python calls
+    to the statement. The `lower` method should be implemented to parse the
+    arguments from the Python call and construct the statement instance.
+    """
+
+    def lower(
+        self, stmt: type[StatementType], state: "LoweringState", node: ast.Call
+    ) -> "Result":
+        return state.default_Call_lower(stmt, node)
+
+    def verify(self, stmt: "Statement"):
+        assert len(stmt.regions) == 0, "FromPythonCall statements cannot have regions"
+        assert (
+            len(stmt.successors) == 0
+        ), "FromPythonCall statements cannot have successors"