switch constant attr to explicit + make IList an attribute (#308)

this PR changes `IList` to an attribute allowing it to be a compile-time
value. It also extends the runtime of `IList` so that one can keep track
of the element type information at runtime. This solves the problem of
capturing global values like below

```python
alist = [1, 2, 3] # could be hard to know element type when list is long
alist = IList([1, 2, 3], elem=types.Int) # allow user explicitly type the global value

@basic
def main()
    return alist
```

This PR is breaking becausae now accessing the `value` field of
`Constant` will give either `PyAttr` or some other attribute type (e.g
`IList`) instead of original `T` and store as `PyAttr`. This also
partially fix #270 . A new attribute base class `Data` and method
`unwrap` are provided for convenience to covnert between attribute and
runtime Python value.

so this is gonna be a tricky breaking change for downstream because it
is hard to detect this behaviour change due to the fact that interpreter
won't check its value type. But this will allow us using other
attributes as a constant value as part of the python dialect.

Author: Roger-luo

src/kirin/dialects/func/typeinfer.py

@@ -24,12 +24,14 @@ def const_none(self, interp: TypeInference, frame: Frame, stmt: ConstantNone):
         return (types.NoneType,)
 
     @impl(Return)
-    def return_(self, interp: TypeInference, frame: Frame, stmt: Return) -> ReturnValue:
+    def return_(
+        self, interp: TypeInference, frame: Frame[types.TypeAttribute], stmt: Return
+    ) -> ReturnValue:
         if (
             isinstance(hint := stmt.value.hints.get("const"), const.Value)
             and hint.data is not None
         ):
-            return ReturnValue(types.Literal(hint.data))
+            return ReturnValue(types.Literal(hint.data, frame.get(stmt.value)))
         return ReturnValue(frame.get(stmt.value))
 
     @impl(Call)

src/kirin/dialects/ilist/interp.py

@@ -13,23 +13,30 @@ class IListInterpreter(MethodTable):
 
     @impl(Range)
     def _range(self, interp, frame: Frame, stmt: Range):
-        return (IList(range(*frame.get_values(stmt.args))),)
+        return (IList(range(*frame.get_values(stmt.args)), elem=types.Int),)
 
     @impl(New)
     def new(self, interp, frame: Frame, stmt: New):
-        return (IList(list(frame.get_values(stmt.values))),)
+        elem_type = types.Any
+        if stmt.values:
+            elem_type = stmt.values[0].type
+            for each in stmt.values[1:]:
+                elem_type = elem_type.join(each.type)
+        return (IList(list(frame.get_values(stmt.values)), elem=elem_type),)
 
     @impl(Len, types.PyClass(IList))
     def len(self, interp, frame: Frame, stmt: Len):
         return (len(frame.get(stmt.value).data),)
 
     @impl(Add, types.PyClass(IList), types.PyClass(IList))
     def add(self, interp, frame: Frame, stmt: Add):
-        return (IList(frame.get(stmt.lhs).data + frame.get(stmt.rhs).data),)
+        lhs, rhs = frame.get(stmt.lhs), frame.get(stmt.rhs)
+        return (IList(lhs.data + rhs.data, elem=lhs.elem.join(rhs.elem)),)
 
     @impl(Push)
     def push(self, interp, frame: Frame, stmt: Push):
-        return (IList(frame.get(stmt.lst).data + [frame.get(stmt.value)]),)
+        lst = frame.get(stmt.lst)
+        return (IList(lst.data + [frame.get(stmt.value)], elem=lst.elem),)
 
     @impl(Map)
     def map(self, interp: Interpreter, frame: Frame, stmt: Map):
@@ -40,7 +47,7 @@ def map(self, interp: Interpreter, frame: Frame, stmt: Map):
             # NOTE: assume fn has been type checked
             _, item = interp.run_method(fn, (elem,))
             ret.append(item)
-        return (IList(ret),)
+        return (IList(ret, elem=fn.return_type),)
 
     @impl(Scan)
     def scan(self, interp: Interpreter, frame: Frame, stmt: Scan):
@@ -54,7 +61,15 @@ def scan(self, interp: Interpreter, frame: Frame, stmt: Scan):
             # NOTE: assume fn has been type checked
             _, (carry, y) = interp.run_method(fn, (carry, elem))
             ys.append(y)
-        return ((carry, IList(ys)),)
+
+        if (
+            isinstance(fn.return_type, types.Generic)
+            and fn.return_type.is_subseteq(types.Tuple)
+            and len(fn.return_type.vars) == 2
+        ):
+            return ((carry, IList(ys, fn.return_type.vars[1])),)
+        else:
+            return ((carry, IList(ys, types.Any)),)
 
     @impl(Foldr)
     def foldr(self, interp: Interpreter, frame: Frame, stmt: Foldr):

src/kirin/dialects/ilist/rewrite/const.py

@@ -2,6 +2,7 @@
 from kirin.analysis import const
 from kirin.rewrite.abc import RewriteRule
 from kirin.rewrite.result import RewriteResult
+from kirin.dialects.py.constant import Constant
 
 from ..stmts import IListType
 from ..runtime import IList
@@ -14,6 +15,9 @@ class ConstList2IList(RewriteRule):
     """
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if isinstance(node, Constant):
+            return self.rewrite_Constant(node)
+
         has_done_something = False
         for result in node.results:
             if not isinstance(hint := result.hints.get("const"), const.Value):
@@ -29,6 +33,14 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
                 has_done_something = self._rewrite_IList_type(result, data)
         return RewriteResult(has_done_something=has_done_something)
 
+    def rewrite_Constant(self, node: Constant) -> RewriteResult:
+        if isinstance(node.value, ir.PyAttr) and isinstance(node.value.data, list):
+            stmt = Constant(value=IList(data=node.value.data))
+            node.replace_by(stmt)
+            self._rewrite_IList_type(stmt.result, node.value.data)
+            return RewriteResult(has_done_something=True)
+        return RewriteResult()
+
     def _rewrite_IList_type(self, result: ir.SSAValue, data):
         if not isinstance(data, IList):
             return False

src/kirin/dialects/ilist/rewrite/list.py

@@ -16,12 +16,15 @@ def rewrite_Block(self, node: ir.Block) -> RewriteResult:
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         has_done_something = False
+        if isinstance(node, constant.Constant) and isinstance(node.value, list):
+            eltype = self._eltype(node.result.type)
+            node.replace_by(
+                constant.Constant(value=IList(data=node.value, elem=eltype))
+            )
+
         for result in node.results:
             has_done_something = self._rewrite_SSAValue_type(result)
 
-        if has_done_something and isinstance(node, constant.Constant):
-            node.replace_by(constant.Constant(value=IList(data=node.value)))
-
         return RewriteResult(has_done_something=has_done_something)
 
     def _rewrite_SSAValue_type(self, value: ir.SSAValue):
@@ -36,3 +39,8 @@ def _rewrite_SSAValue_type(self, value: ir.SSAValue):
             value.type = IListType[types.Any, types.Any]
             return True
         return False
+
+    def _eltype(self, type: types.TypeAttribute):
+        if isinstance(type, types.Generic) and issubclass(type.body.typ, (list, IList)):
+            return type.vars[0]
+        return types.Any

src/kirin/dialects/ilist/runtime.py

@@ -3,15 +3,22 @@
 from dataclasses import dataclass
 from collections.abc import Sequence
 
+from kirin import ir, types
+from kirin.print.printer import Printer
+
 T = TypeVar("T")
 L = TypeVar("L")
 
 
 @dataclass
-class IList(Generic[T, L]):
+class IList(ir.Data[Sequence[T]], Generic[T, L]):
     """A simple immutable list."""
 
     data: Sequence[T]
+    elem: types.TypeAttribute = types.Any
+
+    def __post_init__(self):
+        self.type = types.Generic(IList, self.elem, types.Literal(len(self.data)))
 
     def __hash__(self) -> int:
         return id(self)  # do not hash the data
@@ -26,7 +33,16 @@ def __add__(self, other: "IList[T, Any]") -> "IList[T, Any]": ...
     def __add__(self, other: list[T]) -> "IList[T, Any]": ...
 
     def __add__(self, other):
-        return IList(self.data + other)
+        if isinstance(other, list):
+            return IList(list(self.data) + other, elem=self.elem)
+        elif isinstance(other, IList):
+            return IList(
+                list(self.data) + list(other.data), elem=self.elem.join(other.elem)
+            )
+        else:
+            raise TypeError(
+                f"unsupported operand type(s) for +: 'IList' and '{type(other)}'"
+            )
 
     @overload
     def __radd__(self, other: "IList[T, Any]") -> "IList[T, Any]": ...
@@ -61,3 +77,13 @@ def __eq__(self, value: object) -> bool:
         if not isinstance(value, IList):
             return False
         return self.data == value.data
+
+    def unwrap(self) -> "IList[T, L]":
+        return self
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.plain_print("IList(")
+        printer.print_seq(
+            self.data, delim=", ", prefix="[", suffix="]", emit=printer.plain_print
+        )
+        printer.plain_print(")")

src/kirin/dialects/py/constant.py

@@ -27,12 +27,12 @@
 class Constant(ir.Statement, Generic[T]):
     name = "constant"
     traits = frozenset({ir.Pure(), ir.ConstantLike(), ir.FromPythonCall()})
-    value: T = info.attribute()
+    value: ir.Data[T] = info.attribute()
     result: ir.ResultValue = info.result()
 
     # NOTE: we allow py.Constant take data.PyAttr too
-    def __init__(self, value: T | ir.PyAttr[T]) -> None:
-        if not isinstance(value, ir.PyAttr):
+    def __init__(self, value: T | ir.Data[T]) -> None:
+        if not isinstance(value, ir.Data):
             value = ir.PyAttr(value)
         super().__init__(
             attributes={"value": value},
@@ -68,12 +68,12 @@ class Concrete(interp.MethodTable):
 
     @interp.impl(Constant)
     def constant(self, interp, frame: interp.Frame, stmt: Constant):
-        return (stmt.value,)
+        return (stmt.value.unwrap(),)
 
 
 @dialect.register(key="emit.julia")
 class JuliaTable(interp.MethodTable):
 
     @interp.impl(Constant)
     def emit_Constant(self, emit: EmitJulia, frame: EmitStrFrame, stmt: Constant):
-        return (emit.emit_attribute(ir.PyAttr(stmt.value)),)
+        return (emit.emit_attribute(stmt.value),)

src/kirin/ir/__init__.py

@@ -45,3 +45,4 @@
 from kirin.ir.dialect import Dialect as Dialect
 from kirin.ir.attrs.py import PyAttr as PyAttr
 from kirin.ir.attrs.abc import Attribute as Attribute, AttributeMeta as AttributeMeta
+from kirin.ir.attrs.data import Data as Data

src/kirin/ir/attrs/data.py

@@ -0,0 +1,28 @@
+from abc import abstractmethod
+from typing import Generic, TypeVar
+from dataclasses import field, dataclass
+
+from .abc import Attribute
+from .types import TypeAttribute
+
+T = TypeVar("T", covariant=True)
+
+
+@dataclass(eq=False)
+class Data(Attribute, Generic[T]):
+    """Base class for data attributes.
+
+    Data attributes are compile-time constants that can be used to
+    represent runtime data inside the IR.
+
+    This class is meant to be subclassed by specific data attributes.
+    It provides a `type` attribute that should be set to the type of
+    the data.
+    """
+
+    type: TypeAttribute = field(init=False, repr=False)
+
+    @abstractmethod
+    def unwrap(self) -> T:
+        """Returns the underlying data value."""
+        ...

src/kirin/ir/attrs/py.py

@@ -1,16 +1,16 @@
-from typing import Generic, TypeVar
+from typing import TypeVar
 from dataclasses import dataclass
 
 from kirin.print import Printer
 
-from .abc import Attribute
+from .data import Data
 from .types import PyClass, TypeAttribute
 
 T = TypeVar("T")
 
 
 @dataclass
-class PyAttr(Generic[T], Attribute):
+class PyAttr(Data[T]):
     """Python attribute for compile-time values.
     This is a generic attribute that holds a Python value.
 
@@ -25,7 +25,6 @@ class PyAttr(Generic[T], Attribute):
 
     name = "PyAttr"
     data: T
-    type: TypeAttribute
 
     def __init__(self, data: T, pytype: TypeAttribute | None = None):
         self.data = data
@@ -43,3 +42,6 @@ def print_impl(self, printer: Printer) -> None:
         with printer.rich(style="comment"):
             printer.plain_print(" : ")
             printer.print(self.type)
+
+    def unwrap(self) -> T:
+        return self.data

src/kirin/ir/attrs/types.py

@@ -179,16 +179,16 @@ def __init__(self, *args, **kwargs):
         super(LiteralMeta, self).__init__(*args, **kwargs)
         self._cache = {}
 
-    def __call__(self, data):
-        if isinstance(data, Attribute):
-            return data
+    def __call__(self, data, datatype=None):
+        if isinstance(data, TypeAttribute):
+            return data  # already a type
         elif not isinstance(data, Hashable):
-            return PyClass(type(data))
-        elif data in self._cache:
-            return self._cache[data]
+            raise ValueError("Literal data must be hashable")
+        elif (data, datatype) in self._cache:
+            return self._cache[(data, datatype)]
 
-        instance = super(LiteralMeta, self).__call__(data)
-        self._cache[data] = instance
+        instance = super(LiteralMeta, self).__call__(data, datatype)
+        self._cache[(data, datatype)] = instance
         return instance
 
 
@@ -200,6 +200,14 @@ def __call__(self, data):
 class Literal(TypeAttribute, typing.Generic[LiteralType], metaclass=LiteralMeta):
     name = "Literal"
     data: LiteralType
+    type: TypeAttribute
+    """type of the literal, this is useful when the Python type of
+    data does not represent the type in IR, e.g Literal(1, types.Int32)
+    """
+
+    def __init__(self, data: LiteralType, datatype: TypeAttribute | None = None):
+        self.data = data
+        self.type = datatype or PyClass(type(data))
 
     def is_equal(self, other: TypeAttribute) -> bool:
         return self is other
@@ -211,16 +219,16 @@ def is_subseteq_Union(self, other: "Union") -> bool:
         return any(self.is_subseteq(t) for t in other.types)
 
     def is_subseteq_Literal(self, other: "Literal") -> bool:
-        return self.data == other.data
+        return self.data == other.data and self.type.is_subseteq(other.type)
 
     def is_subseteq_fallback(self, other: TypeAttribute) -> bool:
-        return PyClass(type(self.data)).is_subseteq(other)
+        return self.type.is_subseteq(other)
 
     def __hash__(self) -> int:
         return hash((Literal, self.data))
 
     def print_impl(self, printer: Printer) -> None:
-        printer.plain_print(repr(self.data))
+        printer.plain_print("Literal(", repr(self.data), ",", self.type, ")")
 
 
 @typing.final

src/kirin/ir/nodes/stmt.py

@@ -611,7 +611,7 @@ def print_impl(self, printer: Printer) -> None:
             if isinstance(values, SSAValue):
                 printer.print(values)
             else:
-                printer.print_seq(values, delim=", ")
+                printer.print_seq(values, delim=", ", prefix="(", suffix=")")
 
             if idx < len(self._name_args_slice) - 1:
                 printer.plain_print(", ")