kirin-toolchain 0.13.0__py3-none-any.whl

kirin/decl/emit/init.py

@@ -0,0 +1,277 @@
+from typing import Any
+from dataclasses import MISSING
+
+from typing_extensions import Unpack
+
+from kirin import ir, types
+from kirin.decl import info
+from kirin.decl.base import BaseModifier, StatementOptions
+
+from ._create_fn import create_fn
+from ._set_new_attribute import set_new_attribute
+
+
+class _HAS_DEFAULT_FACTORY_CLASS:
+    def __repr__(self):
+        return "<factory>"
+
+
+_HAS_DEFAULT_FACTORY = _HAS_DEFAULT_FACTORY_CLASS()
+
+
+class EmitInit(BaseModifier):
+    _POST_INIT_NAME = "__post_init__"
+    _RESULT_VALUE_NAME = "_kirin_ResultValue"
+    _KIRIN_PYATTR = "_kirin_PyAttr"
+    _ATTR_FACTORY_PREFIX = "_kirin_attr_factory_"
+    _SELF_CLASS = "_kirin_self_class"
+    _KIRIN_STMT = "_kirin_Statement"
+
+    def __init__(self, cls: type, **kwargs: Unpack[StatementOptions]) -> None:
+        super().__init__(cls, **kwargs)
+        from kirin.ir.attrs.py import PyAttr
+
+        self._init_params: list[str] = []
+        self._init_body: list[str] = []
+        self._init_locals: dict[str, Any] = {}
+        self.has_post_init = hasattr(self.cls, self._POST_INIT_NAME)
+        self.globals.update(
+            {
+                "_PY_ANY": types.Any,
+                self._KIRIN_STMT: ir.Statement,
+                self._SELF_CLASS: self.cls,
+                self._RESULT_VALUE_NAME: ir.ResultValue,
+                self._KIRIN_PYATTR: PyAttr,
+            }
+        )
+
+    @classmethod
+    def _hint_name(cls, f: info.Field):
+        return f"_hint_{f.name}"
+
+    def emit_init(self):
+        _args_groups = []
+        for f in self.fields.args.values():
+            if f.group:
+                _args_groups.append(f.name)
+        setattr(self.cls, "_arg_groups", frozenset(_args_groups))
+
+        if "__init__" in self.cls.__dict__:
+            return True
+
+        self._init_locals.update(
+            {self._hint_name(f): f.annotation for f in self.fields}
+        )
+        self._init_locals.update(
+            {
+                "MISSING": MISSING,
+                "_HAS_DEFAULT_FACTORY": _HAS_DEFAULT_FACTORY,
+            }
+        )
+
+        self._emit_params()
+        self._emit_field_init(len(_args_groups) > 0)
+
+        fn = create_fn(
+            "__init__",
+            self._init_params,
+            self._init_body,
+            globals=self.globals,
+            locals=self._init_locals,
+            return_type=None,
+        )
+        set_new_attribute(self.cls, "__init__", fn)
+
+    def _emit_params(self):
+        self._init_params.append(self._self_name)
+        for f in self.fields.std_args.values():
+            self._init_params.append(self._init_param_value(f))
+
+        kw_params: list[str] = []
+        for f in self.fields.kw_args.values():
+            kw_params.append(self._init_param_value(f))
+
+        for f in self.fields.attributes.values():
+            if f.init:
+                kw_params.append(self._init_param_value(f))
+
+        for f in self.fields.regions.values():
+            if f.init:
+                kw_params.append(self._init_param_value(f))
+
+        for f in self.fields.blocks.values():
+            if f.init:
+                kw_params.append(self._init_param_value(f))
+
+        if kw_params:
+            self._init_params.append("*")
+            self._init_params.extend(kw_params)
+
+    def _emit_field_init(self, has_args_groups: bool):
+        self._init_body.append(
+            (
+                f"{self._KIRIN_STMT}.__init__({self._self_name},"
+                f"args={self._arg_seq()},"
+                f"regions={self._regions_seq()},"
+                f"successors={self._blocks_seq()},"
+                f"result_types={self._result_types_seq()},"
+                f"attributes={self._attribute_seq()},"
+                f"args_slice={self._args_slice(has_args_groups)}"
+                ")"
+            )
+        )
+
+        # TODO: support InitVar?
+        if self.has_post_init:
+            self._init_body.append(f"{self._self_name}.{self._POST_INIT_NAME}()")
+
+    def _init_param_value(self, f: info.Field):
+        if isinstance(f, info.AttributeField):  # only attribute can have defaults
+            if f.default is MISSING and f.default_factory is None:
+                default = ""
+            elif f.default is not MISSING:
+                default = f"={self._ATTR_FACTORY_PREFIX}{f.name}"
+            elif f.default_factory is not None:
+                default = "=_HAS_DEFAULT_FACTORY"
+            else:
+                raise ValueError("unexpected default")
+            return f"{f.name}:{self._hint_name(f)}{default}"
+        elif isinstance(f, info.ArgumentField):  # arguments are always required
+            return f"{f.name}:{self._hint_name(f)}"
+        elif isinstance(f, info.ResultField):
+            raise ValueError("result fields are not allowed in __init__")
+        elif isinstance(f, info.RegionField):
+            return f"{f.name}:{self._hint_name(f)}=_HAS_DEFAULT_FACTORY"
+        elif isinstance(f, info.BlockField):
+            return f"{f.name}:{self._hint_name(f)}=_HAS_DEFAULT_FACTORY"
+        else:  # the rest fields are all required
+            raise ValueError(f"unexpected field type {type(f)}")
+
+    def _arg_seq(self):
+        args: list[str] = [self._field_param(f) for f in self.fields.args.values()]
+        return self._tuple_str(args)
+
+    def _field_param(self, f: info.ArgumentField):
+        if f.group:
+            return f"*{f.name}"
+        else:
+            return f"{f.name}"
+
+    def _result_types_seq(self):
+        result_types = [
+            self._result_type_value(f) for f in self.fields.results.values()
+        ]
+        return self._tuple_str(result_types)
+
+    def _result_type_value(self, f: info.ResultField):
+        name = f"_result_type_{f.name}"
+        self._init_locals[name] = f.type
+        return name
+
+    def _attribute_seq(self):
+        attrs = ",".join(
+            [
+                f'"{name}": {self._attribute_value(f)}'
+                for name, f in self.fields.attributes.items()
+            ]
+        )
+        return "{" + attrs + "}"
+
+    def _attribute_value(self, f: info.AttributeField):
+        default_name = f"{self._ATTR_FACTORY_PREFIX}{f.name}"
+        if f.default_factory is not None:
+            value = self._field_with_default_factory(
+                default_name, f.name, f.init, f.default_factory
+            )
+        else:
+            # no default factory
+            if f.init:
+                if f.default is not MISSING:
+                    self._init_locals[default_name] = f.default
+                value: str = f.name  # type: ignore
+            else:  # no default factory, no init, no default
+                raise ValueError(
+                    "attribute must have a default or default factory or be initialized (init=True)"
+                )
+
+        # declared via python type, optionally check if we can
+        # convert the value to data.PyAttr
+        if f.pytype:
+            attr_type = f"_kirin_attr_type_{f.name}"
+            self._init_locals[attr_type] = f.type
+            value = (
+                f"{self._KIRIN_PYATTR}({value}, {attr_type}) "
+                f"if not isinstance({value}, {self._KIRIN_PYATTR}) "
+                f"else {value}"
+            )
+
+        return value
+
+    def _regions_seq(self):
+        regions = [self._regions_value(f) for f in self.fields.regions.values()]
+        return self._tuple_str(regions)
+
+    def _regions_value(self, f: info.RegionField):
+        return self._field_with_default_factory(
+            f"_kirin_region_{f.name}", f.name, f.init, f.default_factory
+        )
+
+    def _blocks_seq(self):
+        blocks = [self._blocks_value(f) for f in self.fields.blocks.values()]
+        return self._tuple_str(blocks)
+
+    def _blocks_value(self, f: info.BlockField):
+        return self._field_with_default_factory(
+            f"_kirin_block_{f.name}", f.name, f.init, f.default_factory
+        )
+
+    def _args_slice(self, has_args_groups: bool):
+        if not has_args_groups:
+            return (
+                "{"
+                + ", ".join(
+                    [
+                        f"'{f.name}': {index}"
+                        for index, f in enumerate(self.fields.args.values())
+                    ]
+                )
+                + "}"
+            )
+
+        # NOTE: SSAValue fields do not have default/default_factory
+        # so we can just count the input arguments
+        ret: list[str] = []
+        self._init_body.append("__args_slice_count = 0")
+        for f in self.fields.args.values():
+            if f.group:
+                self._init_body.append(
+                    f"__{f.name}_slice = slice(__args_slice_count, __args_slice_count + len({f.name}))"
+                )
+                self._init_body.append(f"__args_slice_count += len({f.name})")
+            else:
+                self._init_body.append(f"__{f.name}_slice = __args_slice_count")
+                self._init_body.append("__args_slice_count += 1")
+            ret.append(f'"{f.name}": __{f.name}_slice')
+        return "{" + ",".join(ret) + "}"
+
+    def _field_with_default_factory(
+        self, factory_name: str, name: str | None, init: bool, default_factory
+    ):
+        # block always has default_factory
+        if init:
+            self._init_locals[factory_name] = default_factory
+            value = (
+                f"{factory_name}() "
+                f"if {name} is _HAS_DEFAULT_FACTORY "
+                f"else {name}"
+            )
+        else:
+            self._init_locals[factory_name] = default_factory
+            value = f"{factory_name}()"
+        return value
+
+    @staticmethod
+    def _tuple_str(seq: list[str]):
+        if not seq:
+            return "()"
+        return f"({', '.join(seq)}, )"

kirin/decl/emit/name.py

@@ -0,0 +1,10 @@
+from kirin.decl.base import BaseModifier
+from kirin.decl.camel2snake import camel2snake
+
+from ._set_new_attribute import set_new_attribute
+
+
+class EmitName(BaseModifier):
+
+    def emit_name(self):
+        set_new_attribute(self.cls, "name", camel2snake(self.cls.__name__))

kirin/decl/emit/property.py

@@ -0,0 +1,182 @@
+from kirin import ir
+from kirin.decl import info
+from kirin.decl.emit.init import BaseModifier
+
+from ._create_fn import create_fn
+from ._set_new_attribute import set_new_attribute
+
+
+class EmitProperty(BaseModifier):
+    _KIRIN_PYATTR = "_kirin_PyAttr"
+
+    def emit_property(self):
+        for name, f in self.fields.args.items():
+            getter, setter = self._emit_arg_property(f)
+            set_new_attribute(self.cls, name, property(getter, setter))
+
+        for name, f in self.fields.kw_args.items():
+            getter, setter = self._emit_arg_property(f)
+            set_new_attribute(self.cls, name, property(getter, setter))
+
+        for i, (name, f) in enumerate(self.fields.results.items()):
+            getter, setter = self._emit_result_property(i, f)
+            set_new_attribute(self.cls, name, property(getter, setter))
+
+        for name, f in self.fields.attributes.items():
+            getter, setter = self._emit_attribute_property(f)
+            set_new_attribute(self.cls, name, property(getter, setter))
+
+        for i, (name, f) in enumerate(self.fields.blocks.items()):
+            getter, setter = self._emit_successor_property(i, f)
+            set_new_attribute(self.cls, name, property(getter, setter))
+
+        for i, (name, f) in enumerate(self.fields.regions.items()):
+            getter, setter = self._emit_region_property(i, f)
+            set_new_attribute(self.cls, name, property(getter, setter))
+
+    def _emit_arg_property(self, f: info.ArgumentField):
+        getter = create_fn(
+            f"_get_{f.name}",
+            args=[self._self_name],
+            body=[
+                f'return {self._self_name}._args[{self._self_name}._name_args_slice["{f.name}"]]'
+            ],
+            globals=self.globals,
+            return_type=tuple[ir.SSAValue, ...] if f.group else ir.SSAValue,
+        )
+
+        if f.group:
+            return getter, self._emit_arg_property_group_setter(f)
+        return getter, self._emit_arg_property_std_setter(f)
+
+    def _emit_arg_property_std_setter(self, f: info.ArgumentField):
+        return create_fn(
+            f"_set_{f.name}",
+            args=[self._self_name, "value: _value_hint"],
+            body=[
+                f"s = {self._self_name}._name_args_slice['{f.name}']",
+                f"old = {self._self_name}.args[s]",
+                f"old.remove_use(Use({self._self_name}, s))",
+                f"value.add_use(Use({self._self_name}, s))",
+                "self.args = (*self.args[:s], value, *self.args[s + 1:])",
+            ],
+            globals=self.globals,
+            locals={"_value_hint": ir.SSAValue},
+            return_type=None,
+        )
+
+    def _emit_arg_property_group_setter(self, f: info.ArgumentField):
+        return create_fn(
+            f"_set_{f.name}",
+            args=[self._self_name, "value: _value_hint"],
+            body=[
+                f"s = {self._self_name}._name_args_slice['{f.name}']",
+                "start, stop = s.start, s.stop",
+                f"old = {self._self_name}.args[s]",
+                "a_range = range(start, stop, s.step or 1)",
+                "for i, arg in zip(a_range, old):",
+                f"    arg.remove_use(Use({self._self_name}, i))",
+                "for i, arg in zip(a_range, value):",
+                f"    arg.add_use(Use({self._self_name}, i))",
+                f"{self._self_name}.args = (*{self._self_name}.args[:start], *value, *{self._self_name}.args[stop + 1:])",
+            ],
+            globals=self.globals,
+            locals={"_value_hint": tuple[ir.SSAValue, ...]},
+            return_type=None,
+        )
+
+    def _emit_result_property(self, index: int, f: info.ResultField):
+        getter = create_fn(
+            f"_get_{f.name}",
+            args=[self._self_name],
+            body=[f"return {self._self_name}._results[{index}]"],
+            globals=self.globals,
+            return_type=ir.ResultValue,
+        )
+
+        # well you cannot delete what already happened, can you?
+        setter = create_fn(
+            f"_set_{f.name}",
+            args=[self._self_name, "value: _value_hint"],
+            body=[f"raise AttributeError('result property {f.name} is read-only')"],
+            globals=self.globals,
+            locals={"_value_hint": ir.ResultValue},
+            return_type=None,
+        )
+
+        return getter, setter
+
+    def _emit_attribute_property(self, f: info.AttributeField):
+        from kirin.ir.attrs.py import PyAttr
+
+        storage = "attributes"
+        attr = f"{self._self_name}.{storage}['{f.name}']"
+        getter = create_fn(
+            f"_get_{f.name}",
+            args=[self._self_name],
+            body=[f"return {attr}.data" if f.pytype else f"return {attr}"],
+            globals=self.globals,
+            return_type=f.annotation,
+        )
+
+        setter = create_fn(
+            f"_set_{f.name}",
+            args=[self._self_name, "value: _value_hint"],
+            body=[
+                (
+                    f"{attr} = value if isinstance(value, {self._KIRIN_PYATTR}) else {self._KIRIN_PYATTR}(value)"
+                    if f.pytype
+                    else f"{attr} = value"
+                )
+            ],
+            globals=self.globals,
+            locals={"_value_hint": PyAttr if f.pytype else f.annotation},
+            return_type=None,
+        )
+
+        return getter, setter
+
+    def _emit_successor_property(self, index: int, f: info.BlockField):
+        getter = create_fn(
+            f"_get_{f.name}",
+            args=[self._self_name],
+            body=[f"return {self._self_name}.successors[{index}]"],
+            globals=self.globals,
+            return_type=ir.Block,
+        )
+
+        setter = create_fn(
+            f"_set_{f.name}",
+            args=[self._self_name, "value: _value_hint"],
+            body=[f"{self._self_name}.successors[{index}] = value"],
+            globals=self.globals,
+            locals={"_value_hint": ir.Block},
+            return_type=None,
+        )
+
+        return getter, setter
+
+    def _emit_region_property(self, index: int, f: info.RegionField):
+        getter = create_fn(
+            f"_get_{f.name}",
+            args=[self._self_name],
+            body=[f"return {self._self_name}.regions[{index}]"],
+            globals=self.globals,
+            return_type=ir.Region,
+        )
+
+        setter = create_fn(
+            f"_set_{f.name}",
+            args=[self._self_name, "value: _value_hint"],
+            body=[
+                f"old = {self._self_name}.regions[{index}]",
+                "old.parent = None",
+                f"value.parent = {self._self_name}",
+                f"{self._self_name}.regions[{index}] = value",
+            ],
+            globals=self.globals,
+            locals={"_value_hint": ir.Region},
+            return_type=None,
+        )
+
+        return getter, setter

kirin/decl/emit/repr.py

@@ -0,0 +1,31 @@
+from kirin.decl.base import BaseModifier
+
+from ._create_fn import create_fn
+from ._set_new_attribute import set_new_attribute
+
+
+class EmitRepr(BaseModifier):
+
+    def emit_repr(self):
+        if "repr" not in self.params or not self.params["repr"]:
+            return
+
+        body = [f'ret = "{self.cls.__name__}("']
+        for idx, field in enumerate(self.fields):
+            if idx > 0:
+                body.append('ret += ", "')
+            body.append(f'ret += f"{field.name}={{{self._self_name}.{field.name}}}"')
+        body.append('ret += ")"')
+        body.append("return ret")
+
+        set_new_attribute(
+            self.cls,
+            "__repr__",
+            create_fn(
+                "__repr__",
+                args=[self._self_name],
+                body=body,
+                globals=self.globals,
+                return_type=str,
+            ),
+        )

kirin/decl/emit/traits.py

@@ -0,0 +1,13 @@
+from kirin.decl.base import BaseModifier
+
+from ._set_new_attribute import set_new_attribute
+
+
+class EmitTraits(BaseModifier):
+
+    def emit_traits(self):
+        # if no parent defines traits, set it to empty set
+        for base in self.cls.__mro__[-1:0:-1]:
+            if hasattr(base, "traits"):
+                return
+        set_new_attribute(self.cls, "traits", frozenset({}))

kirin/decl/emit/typecheck.py

@@ -0,0 +1,77 @@
+from typing import Any
+
+from kirin.exceptions import VerificationError
+from kirin.decl.emit.init import BaseModifier
+
+from ._create_fn import create_fn
+from ._set_new_attribute import set_new_attribute
+
+
+class EmitTypeCheck(BaseModifier):
+    _VERIFICATION_ERROR = "_kirin_VerificationError"
+
+    def emit_typecheck(self):
+        typecheck_locals: dict[str, Any] = {
+            self._VERIFICATION_ERROR: VerificationError,
+        }
+        body: list[str] = []
+        for name, f in self.fields.args.items():
+            if f.type is f.type.top():
+                continue
+
+            value_type = f"_args_{f.name}_type"
+            typecheck_locals.update({value_type: f.type})
+            if f.group:
+                body.extend(
+                    (
+                        f"for v in {self._self_name}.{f.name}:",
+                        *self._guard_ssa_type("v", f.name, value_type, indent=1),
+                    )
+                )
+            else:
+                body.extend(
+                    self._guard_ssa_type(
+                        f"{self._self_name}.{name}", f.name, value_type
+                    )
+                )
+
+        for name, f in self.fields.results.items():
+            if f.type is f.type.top():
+                continue
+
+            value_type = f"_results_{f.name}_type"
+            typecheck_locals.update({value_type: f.type})
+            body.extend(
+                self._guard_ssa_type(f"{self._self_name}.{name}", name, value_type)
+            )
+
+        for name in self.fields.blocks.keys():
+            body.append(f"{self._self_name}.{name}.typecheck()")
+
+        for name, f in self.fields.regions.items():
+            body.append(f"{self._self_name}.{name}.typecheck()")
+
+        # NOTE: we still need to generate this because it is abstract
+        if not body:
+            body.append("pass")
+
+        set_new_attribute(
+            self.cls,
+            "typecheck",
+            create_fn(
+                name="_kirin_decl_typecheck",
+                args=[self._self_name],
+                body=body,
+                globals=self.globals,
+                locals=typecheck_locals,
+                return_type=None,
+            ),
+        )
+
+    def _guard_ssa_type(self, ssa, name, type, indent: int = 0):
+        space = "  " * indent
+        msg = f"'Invalid type for {name}, expected ' + repr({type}) + ', got ' + repr({ssa}.type)"
+        return (
+            space + f"if not {ssa}.type.is_subseteq({type}):",
+            space + f"    raise {self._VERIFICATION_ERROR}({self._self_name}, {msg})",
+        )

kirin/decl/emit/verify.py

@@ -0,0 +1,51 @@
+from typing import Any
+
+from kirin.exceptions import VerificationError
+from kirin.decl.emit.init import BaseModifier
+
+from ._create_fn import create_fn
+from ._set_new_attribute import set_new_attribute
+
+
+class EmitVerify(BaseModifier):
+    _VERIFICATION_ERROR = "_kirin_VerificationError"
+
+    def emit_verify(self):
+        verify_locals: dict[str, Any] = {
+            self._VERIFICATION_ERROR: VerificationError,
+        }
+        body: list[str] = []
+        for name in self.fields.blocks.keys():
+            body.append(f"{self._self_name}.{name}.verify()")
+
+        for name, f in self.fields.regions.items():
+            body.append(f"{self._self_name}.{name}.verify()")
+            if not f.multi:
+                body.append(f"if len({self._self_name}.{name}.blocks) != 1:")
+                body.append(
+                    f"    raise {self._VERIFICATION_ERROR}({self._self_name},"
+                    f" 'Invalid number of blocks for {name}')"
+                )
+
+        if (traits := getattr(self.cls, "traits", None)) is not None:
+            for trait in traits:
+                trait_obj = f"_kirin_verify_trait_{trait.__class__.__name__}"
+                verify_locals.update({trait_obj: trait})
+                body.append(f"{trait_obj}.verify({self._self_name})")
+
+        # NOTE: we still need to generate this because it is abstract
+        if not body:
+            body.append("pass")
+
+        set_new_attribute(
+            self.cls,
+            "verify",
+            create_fn(
+                name="_kirin_decl_verify",
+                args=[self._self_name],
+                body=body,
+                globals=self.globals,
+                locals=verify_locals,
+                return_type=None,
+            ),
+        )