kirin-toolchain 0.13.0__py3-none-any.whl

kirin/decl/info.py

@@ -0,0 +1,346 @@
+from abc import abstractmethod
+from types import GenericAlias
+from typing import Any, Callable, Optional
+from functools import cached_property
+from dataclasses import MISSING, field, dataclass
+
+from kirin import types
+from kirin.ir import Block, Region, Attribute
+
+
+@dataclass
+class Field:
+    name: Optional[str] = field(default=None, init=False)
+    annotation: Any = field(default=None, init=False)
+    kw_only: bool
+    alias: Optional[str]
+
+    __class_getitem__ = classmethod(GenericAlias)
+
+    @abstractmethod
+    def has_no_default(self) -> bool: ...
+
+
+@dataclass
+class AttributeField(Field):
+    default: Any
+    init: bool
+    repr: bool
+    default_factory: Optional[Callable[[], Attribute]]
+    type: types.TypeAttribute
+    pytype: bool = False
+    "if `True`, annotation is a python type hint instead of `TypeAttribute`"
+
+    def has_no_default(self):
+        return self.default is MISSING and self.default_factory is None
+
+
+def attribute(
+    type: types.TypeAttribute = types.Any,
+    *,
+    init: bool = True,
+    repr: bool = True,
+    default: Any = MISSING,
+    default_factory: Optional[Callable[[], Any]] = None,
+    kw_only: bool = True,
+    alias: Optional[str] = None,
+) -> Any:
+    if kw_only is False:
+        raise TypeError("attribute fields must be keyword-only")
+
+    return AttributeField(
+        type=type,
+        init=init,
+        repr=repr,
+        default=default,
+        default_factory=default_factory,
+        kw_only=kw_only,
+        alias=alias,
+    )
+
+
+@dataclass
+class ArgumentField(Field):
+    type: types.TypeAttribute
+    """type of the argument, will be used in validation.
+    """
+    print: bool = True
+    """if `True`, this argument name is printed in the signature.
+    """
+    group: bool = False  # NOTE: this cannot be set by user
+    """if `True`, this argument is annotated with Tuple[SSAValue, ...]
+    """
+
+    def has_no_default(self):
+        return True
+
+
+# NOTE: argument must appear in init and repr
+def argument(
+    type: types.TypeAttribute = types.Any,
+    *,
+    print: bool = True,
+    kw_only: bool = False,
+    alias: Optional[str] = None,
+) -> Any:
+    """Field specifier for arguments.
+
+    Args:
+        type(TypeAttribute): type of the argument, will be used in validation.
+        print(bool): if `True`, this argument name is printed in the signature.
+        kw_only(bool): if `True`, this argument is keyword-only.
+        alias(Optional[str]): an alias for the argument name in the `__init__` method.
+    """
+    return ArgumentField(
+        type=type,
+        print=print,
+        kw_only=kw_only,
+        alias=alias,
+    )
+
+
+@dataclass
+class ResultField(Field):
+    init: bool
+    repr: bool
+    type: types.TypeAttribute = field(default_factory=types.AnyType)
+
+    def has_no_default(self):
+        return True
+
+
+def result(
+    type: types.TypeAttribute = types.Any,
+    *,
+    # NOTE: init is false, use other hooks to set custom results
+    # or just mutate the statement after creation
+    init: bool = False,
+    repr: bool = True,
+    kw_only: bool = True,
+    alias: Optional[str] = None,
+) -> Any:
+    """Field specifier for results.
+
+    Args:
+        type(TypeAttribute): type of the result.
+        init(bool): if `True`, this result field is included in the `__init__` method.
+        repr(bool): if `True`, this result field is included in the `__repr__` and pretty printing.
+        kw_only(bool): if `True`, this result field is keyword-only.
+        alias(Optional[str]): an alias for the result field name in the `__init__` method.
+    """
+    if kw_only is False:  # for linting
+        raise TypeError("result fields must be keyword-only")
+
+    if init is True:
+        raise TypeError("result fields cannot appear in __init__")
+
+    return ResultField(
+        type=type,
+        init=init,
+        repr=repr,
+        kw_only=kw_only,
+        alias=alias,
+    )
+
+
+@dataclass
+class RegionField(Field):
+    init: bool
+    repr: bool
+    multi: bool
+    default_factory: Callable[[], Region]
+
+    def has_no_default(self):
+        return False
+
+
+def region(
+    *,
+    init: bool = True,  # so we can use the default_factory
+    repr: bool = True,
+    kw_only: bool = True,
+    alias: Optional[str] = None,
+    multi: bool = False,
+    default_factory: Callable[[], Region] = Region,
+) -> Any:
+    """Field specifier for regions.
+
+    Args:
+        init(bool): if `True`, this region field is included in the `__init__` method.
+        repr(bool): if `True`, this region field is included in the `__repr__` and pretty printing.
+        kw_only(bool): if `True`, this region field is keyword-only.
+        alias(Optional[str]): an alias for the region field name in the `__init__` method.
+        multi(bool): if `True`, this region can contain multiple blocks.
+        default_factory(Callable[[], Region]): a factory function to create a default region.
+    """
+    if kw_only is False:
+        raise TypeError("region fields must be keyword-only")
+
+    return RegionField(
+        init=init,
+        repr=repr,
+        kw_only=kw_only,
+        alias=alias,
+        multi=multi,
+        default_factory=default_factory,
+    )
+
+
+@dataclass
+class BlockField(Field):
+    init: bool
+    repr: bool
+    default_factory: Callable[[], Block]
+
+    def has_no_default(self):
+        return False
+
+
+def block(
+    *,
+    init: bool = True,
+    repr: bool = True,
+    kw_only: bool = True,
+    alias: Optional[str] = None,
+    default_factory: Callable[[], Block] = Block,
+) -> Any:
+    """Field specifier for blocks.
+
+    Args:
+        init(bool): if `True`, this block field is included in the `__init__` method.
+        repr(bool): if `True`, this block field is included in the `__repr__` and pretty printing.
+        kw_only(bool): if `True`, this block field is keyword-only.
+        alias(Optional[str]): an alias for the block field name in the `__init__` method.
+        default_factory(Callable[[], Block]): a factory function to create a default block.
+    """
+    if kw_only is False:
+        raise TypeError("block fields must be keyword-only")
+
+    return BlockField(
+        init=init,
+        repr=repr,
+        kw_only=kw_only,
+        alias=alias,
+        default_factory=default_factory,
+    )
+
+
+@dataclass
+class StatementFields:
+    std_args: dict[str, ArgumentField] = field(default_factory=dict)
+    """standard arguments of the statement."""
+    kw_args: dict[str, ArgumentField] = field(default_factory=dict)
+    """keyword-only arguments of the statement."""
+    results: dict[str, ResultField] = field(default_factory=dict)
+    """results of the statement."""
+    regions: dict[str, RegionField] = field(default_factory=dict)
+    """regions of the statement."""
+    blocks: dict[str, BlockField] = field(default_factory=dict)
+    """blocks of the statement."""
+    attributes: dict[str, AttributeField] = field(default_factory=dict)
+    """attributes of the statement."""
+
+    class Args:
+        def __init__(self, fields: "StatementFields"):
+            self.fields = fields
+
+        def __len__(self):
+            return len(self.fields.std_args) + len(self.fields.kw_args)
+
+        def __getitem__(self, name):
+            if (value := self.fields.std_args.get(name)) is not None:
+                return value
+            elif (value := self.fields.kw_args.get(name)) is not None:
+                return value
+            raise KeyError(name)
+
+        def __setitem__(self, name: str, value: ArgumentField):
+            if value.kw_only:
+                self.fields.kw_args[name] = value
+            else:
+                self.fields.std_args[name] = value
+
+        def __contains__(self, name):
+            return name in self.fields.std_args or name in self.fields.kw_args
+
+        def values(self):
+            yield from self.fields.std_args.values()
+            yield from self.fields.kw_args.values()
+
+        def items(self):
+            yield from self.fields.std_args.items()
+            yield from self.fields.kw_args.items()
+
+        def keys(self):
+            yield from self.fields.std_args.keys()
+            yield from self.fields.kw_args.keys()
+
+    @property
+    def args(self):
+        """iterable of all argument fields."""
+        return self.Args(self)
+
+    @classmethod
+    def from_fields(cls, fields: dict[str, Field]):
+        ret = cls()
+        for name, f in fields.items():
+            ret[name] = f
+        return ret
+
+    def __contains__(self, name):
+        return (
+            name in self.args
+            or name in self.results
+            or name in self.regions
+            or name in self.blocks
+            or name in self.attributes
+        )
+
+    def __setitem__(self, name, value):
+        if isinstance(value, ArgumentField):
+            self.args[name] = value
+        elif isinstance(value, ResultField):
+            self.results[name] = value
+        elif isinstance(value, RegionField):
+            self.regions[name] = value
+        elif isinstance(value, BlockField):
+            self.blocks[name] = value
+        elif isinstance(value, AttributeField):
+            self.attributes[name] = value
+        else:
+            raise TypeError(f"unknown field type {value}")
+
+    def __iter__(self):
+        yield from self.args.values()
+        yield from self.kw_args.values()
+        yield from self.results.values()
+        yield from self.regions.values()
+        yield from self.blocks.values()
+        yield from self.attributes.values()
+
+    def __len__(self):
+        return (
+            len(self.args)
+            + len(self.results)
+            + len(self.regions)
+            + len(self.blocks)
+            + len(self.attributes)
+        )
+
+    @cached_property
+    def attr_or_props(self):
+        return set(self.attributes.keys())
+
+    @cached_property
+    def required_names(self):
+        """set of all fields that do not have a default value."""
+        return set(
+            list(self.args.keys())
+            + [name for name, f in self.attributes.items() if f.has_no_default()]
+            + [name for name, f in self.blocks.items() if f.has_no_default()]
+            + [name for name, f in self.regions.items() if f.has_no_default()]
+        )
+
+    @cached_property
+    def group_arg_names(self):
+        return set([name for name, f in self.args.items() if f.group])

kirin/decl/scan_fields.py

@@ -0,0 +1,157 @@
+import re
+import sys
+import inspect
+import dataclasses
+from types import ModuleType
+from typing import Callable, get_args
+from dataclasses import KW_ONLY
+
+from beartype.door import is_subhint
+
+from kirin import ir, types
+
+from .base import BaseModifier
+from .info import Field, ArgumentField, AttributeField, StatementFields, argument
+
+
+class ScanFields(BaseModifier):
+    _FIELDS = "__kirin_stmt_fields"
+    # String regex that string annotations for ClassVar or InitVar must match.
+    # Allows "identifier.identifier[" or "identifier[".
+    # https://bugs.python.org/issue33453 for details.
+    _MODULE_IDENTIFIER_RE = re.compile(r"^(?:\s*(\w+)\s*\.)?\s*(\w+)")
+
+    def scan_fields(self):
+        self._scan_base_fields()
+        self._scan_annotations()
+        setattr(self.cls, self._FIELDS, self.fields)
+        return
+
+    def _scan_base_fields(self):
+        # let's just assume dict preserve insertion order in python 3.7+
+        for b in self.cls.__mro__[-1:0:-1]:  # taken from dataclass impl
+            base_fields: StatementFields | None = getattr(b, self._FIELDS, None)
+            if base_fields is not None:
+                self.has_statement_bases = True
+                for f in base_fields:
+                    assert f.name is not None, "field name must be set"
+                    self.fields[f.name] = f
+
+    def _scan_annotations(self):
+        cls_fields: list[Field] = []
+        cls_annotations = inspect.get_annotations(self.cls, eval_str=True)
+        for name, typ in cls_annotations.items():
+            # See if this is a marker to change the value of kw_only.
+            if self._is_kw_only(typ) or (
+                isinstance(typ, str)
+                and self._is_type(typ, dataclasses, KW_ONLY, self._is_kw_only)
+            ):
+                if self.KW_ONLY_seen:
+                    raise TypeError(
+                        f"{name!r} is KW_ONLY, but KW_ONLY "
+                        "has already been specified"
+                    )
+                self.KW_ONLY_seen = True
+                self.kw_only = True
+            else:
+                cls_fields.append(self._get_field(name, typ))
+
+        for f in cls_fields:
+            name: str = f.name  # type: ignore # name has been set
+            self.fields[name] = f
+            if hasattr(self.cls, name):
+                # remove the field from the class
+                # unlike dataclass, we don't actually
+                # store any values in the class, they
+                # are stored inside the IR node.
+                delattr(self.cls, name)
+
+    def _get_field(self, name: str, typ):
+        """Return a Field object for this field name and type."""
+        guess = self._guess_type(typ)
+        default = getattr(self.cls, name, dataclasses.MISSING)
+        if isinstance(default, Field):
+            f = default
+        # no default descriptor, create a new one
+        elif is_subhint(typ, ir.SSAValue) or is_subhint(typ, tuple[ir.SSAValue, ...]):
+            f: Field = (
+                argument()
+            )  # only argument can have default, others must use field specifiers
+        elif typ is ir.ResultValue:
+            raise ValueError(
+                f"expect field specifiers for `ir.ResultValue` fields: {name}"
+            )
+        elif typ is ir.Region:
+            raise ValueError(f"expect field specifiers for `ir.Region` fields: {name}")
+        elif typ is ir.Block:
+            raise ValueError(f"expect field specifiers for `ir.Block` fields: {name}")
+        else:
+            # could be a wrong SSAValue field, e.g list[SSAValue], try check args
+            if typ_args := get_args(typ):
+                if any(is_subhint(arg, ir.SSAValue) for arg in typ_args):
+                    raise ValueError(
+                        f"unsupported SSAValue field: {name},"
+                        f" expect `SSAValue` or `tuple[SSAValue, ...]`, got {typ}"
+                    )
+            raise ValueError(f"expect field specifiers for attribute fields: {name}")
+
+        f.name = name
+        f.annotation = typ
+        self._post_init_field(f, guess)
+        return f
+
+    def _post_init_field(self, f: Field, guess: type | None):
+        if isinstance(f, ArgumentField) and is_subhint(guess, tuple[ir.SSAValue, ...]):
+            f.group = True
+        # try to narrow the type based on the guess
+        elif isinstance(f, AttributeField):
+            if guess and not is_subhint(
+                guess, ir.Attribute
+            ):  # not specified, and using python type
+                if f.type is types.Any:  # not set or too generic
+                    f.type = types.hint2type(guess)
+                f.pytype = True
+
+    @staticmethod
+    def _is_kw_only(a_type):
+        return a_type is KW_ONLY
+
+    def _is_type(
+        self,
+        annotation: type | str,
+        obj_module: ModuleType,
+        obj_type: type,
+        is_type_predicate: Callable[[type], bool],
+    ):
+        """Given a type annotation string, does it refer to `obj_type` in
+        `obj_module`?  For example, when checking that annotation denotes a
+        `ClassVar`, then `obj_module` is typing, and a_type is
+        `typing.ClassVar`.
+
+        Taken from dataclasses.py/is_type.
+        """
+        guess = self._guess_type(annotation)
+        if guess is not None:
+            return guess.__module__ is obj_module and is_type_predicate(guess)
+        return False
+
+    def _guess_type(self, annotation: type | str) -> type | None:
+        """Guess the type/hint object from a string annotation."""
+        if not isinstance(annotation, str):
+            return annotation
+
+        match = self._MODULE_IDENTIFIER_RE.match(annotation)
+        ns = None
+        if match:
+            module_name = match.group(1)
+            if not module_name:
+                # No module name, assume the class's module did
+                # "from dataclasses import InitVar".
+                ns = sys.modules.get(self.cls.__module__).__dict__
+            else:
+                # Look up module_name in the class's module.
+                if self.cls_module:
+                    ns = self.cls_module.__dict__.get(module_name).__dict__
+            if ns:
+                return ns.get(match.group(2))
+        return None

kirin/decl/verify.py

@@ -0,0 +1,69 @@
+import inspect
+import textwrap
+
+from beartype.door import is_subhint
+
+from kirin import ir
+
+from .base import BaseModifier
+from .info import Field, ResultField, ArgumentField
+
+
+class Verify(BaseModifier):
+    RESERVED_NAMES = {
+        "name",
+        "traits",
+        "dialect",
+        "args",
+        "parent",
+        "parent_node",
+        "parent_region",
+        "parent_block",
+        "next_stmt",
+        "prev_stmt",
+        "results",
+        "regions",
+        "successors",
+        "attributes",
+    }
+
+    def verify_mro(self):
+        if not issubclass(self.cls, ir.Statement):
+            raise TypeError(
+                f"expected {self.cls.__name__} to be a subclass of ir.Statement"
+            )
+
+    def verify_fields(self):
+        # Do we have any Field members that don't also have annotations?
+        cls_annotations = inspect.get_annotations(self.cls)
+        for name, value in self.cls.__dict__.items():
+            if isinstance(value, Field) and name not in cls_annotations:
+                raise ValueError(f"{name!r} is a field but has no type annotation")
+
+        for f in self.fields:
+            if f.name in self.RESERVED_NAMES:
+                raise ValueError(f"{f.name!r} is a reserved name")
+
+            if isinstance(f, ArgumentField):
+                if not (
+                    is_subhint(f.annotation, ir.SSAValue)
+                    or is_subhint(f.annotation, tuple[ir.SSAValue, ...])
+                ):
+                    raise ValueError(
+                        f"{f.name!r} is an argument field but has an invalid type annotation"
+                    )
+
+                if is_subhint(f.annotation, ir.ResultValue):
+                    raise ValueError(
+                        textwrap.dedent(
+                            f"{f.name!r} is an argument field but has an invalid type annotation {f.annotation}"
+                            " (did you mean to use `info.result` instead?)"
+                        )
+                    )
+
+            if isinstance(f, ResultField) and not is_subhint(
+                f.annotation, ir.ResultValue
+            ):
+                raise ValueError(
+                    f"{f.name!r} is a result field but has an invalid type annotation"
+                )

kirin/dialects/__init__.py

@@ -0,0 +1,14 @@
+"""Built-in dialects for Kirin.
+
+This module contains the built-in dialects for Kirin. Each dialect is an
+instance of the `Dialect` class. Each submodule contains a `dialect` variable
+that is an instance of the corresponding `Dialect` class.
+
+The modules can be directly used as dialects. For example, you can write
+
+```python
+from kirin.dialects import py, func
+```
+
+to import the Python and function dialects.
+"""

kirin/dialects/_pprint_helper.py

@@ -0,0 +1,53 @@
+from typing import TYPE_CHECKING
+
+from kirin import ir
+
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+    from kirin.print.printer import Printer
+
+
+def pprint_calllike(
+    invoke_or_call: "Statement", callee: str, printer: "Printer"
+) -> None:
+    with printer.rich(style="red"):
+        printer.print_name(invoke_or_call)
+    printer.plain_print(" ")
+
+    n_total = len(invoke_or_call.args)
+    printer.plain_print(callee)
+    if (inputs := getattr(invoke_or_call, "inputs", None)) is None:
+        raise ValueError(f"{invoke_or_call} does not have inputs")
+
+    if not isinstance(inputs, tuple):
+        raise ValueError(f"inputs of {invoke_or_call} is not a tuple")
+
+    if (kwargs := getattr(invoke_or_call, "kwargs", None)) is None:
+        raise ValueError(f"{invoke_or_call} does not have kwargs")
+
+    if not isinstance(kwargs, tuple):
+        raise ValueError(f"kwargs of {invoke_or_call} is not a tuple")
+
+    positional = inputs[: n_total - len(kwargs)]
+    kwargs = dict(
+        zip(
+            kwargs,
+            inputs[n_total - len(kwargs) :],
+        )
+    )
+
+    printer.plain_print("(")
+    printer.print_seq(positional)
+    if kwargs and positional:
+        printer.plain_print(", ")
+    printer.print_mapping(kwargs, delim=", ")
+    printer.plain_print(")")
+
+    with printer.rich(style="comment"):
+        printer.plain_print(" : ")
+        printer.print_seq(
+            [result.type for result in invoke_or_call._results],
+            delim=", ",
+        )
+        if trait := invoke_or_call.get_trait(ir.MaybePure):
+            printer.plain_print(f" maybe_pure={trait.is_pure(invoke_or_call)}")

kirin/dialects/cf/__init__.py

@@ -0,0 +1,20 @@
+"""Control flow dialect.
+
+This dialect provides a low-level control flow representation.
+
+This dialect does not provide any lowering strategies, to lowering
+a Python AST to this dialect, use the `kirin.dialects.lowering.cf` dialect
+with this dialect.
+"""
+
+from kirin.dialects.cf import (
+    emit as emit,
+    constprop as constprop,
+    typeinfer as typeinfer,
+)
+from kirin.dialects.cf.stmts import (
+    Branch as Branch,
+    ConditionalBranch as ConditionalBranch,
+)
+from kirin.dialects.cf.interp import CfInterpreter as CfInterpreter
+from kirin.dialects.cf.dialect import dialect as dialect