revise prelude dialect groups (#148)

this PR revise the prelude dialect groups as following:

`basic` and `basic_no_opt` are the basic language with immutable list
and anything that is common in eDSLs in the future.
`python_basic` and `python_no_opt` are the exact python language with
mutable list and anything that is Python oriented in the future.

the main issue here is to have a separation where those hard to compile
but easy to interpret Python features should go.

Author: Roger-luo

src/kirin/dialects/ilist/__init__.py

@@ -21,5 +21,6 @@
     ForEach as ForEach,
     IListType as IListType,
 )
+from .passes import IListDesugar as IListDesugar
 from .runtime import IList as IList
 from ._dialect import dialect as dialect

src/kirin/dialects/ilist/passes.py

@@ -0,0 +1,32 @@
+from kirin.ir import Method, types
+from kirin.rewrite import Walk, Fixpoint
+from kirin.passes.abc import Pass
+from kirin.rewrite.result import RewriteResult
+from kirin.dialects.ilist.rewrite import List2IList
+
+
+class IListDesugar(Pass):
+    """This pass desugars the Python list dialect
+    to the immutable list dialect by rewriting all
+    constant `list` type into `IList` type.
+    """
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        for arg in mt.args:
+            _check_list(arg.type, arg.type)
+        return Fixpoint(Walk(List2IList())).rewrite(mt.code)
+
+
+def _check_list(total: types.TypeAttribute, type_: types.TypeAttribute):
+    if isinstance(type_, types.Generic):
+        _check_list(total, type_.body)
+        for var in type_.vars:
+            _check_list(total, var)
+        if type_.vararg:
+            _check_list(total, type_.vararg.typ)
+    elif isinstance(type_, types.PyClass):
+        if issubclass(type_.typ, list):
+            raise TypeError(
+                f"Invalid type {total} for this kernel, use IList instead of {type_}."
+            )
+    return

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

@@ -19,10 +19,15 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         return RewriteResult(has_done_something=has_done_something)
 
     def _rewrite_SSAValue_type(self, value: ir.SSAValue):
-        if value.type.is_subseteq(ir.types.List):
-            if isinstance(value.type, ir.types.Generic):
-                value.type = IListType[value.type.vars[0], ir.types.Any]
-            else:
-                value.type = IListType[ir.types.Any, ir.types.Any]
+        # NOTE: cannot use issubseteq here because type can be Bottom
+        if isinstance(value.type, ir.types.Generic) and issubclass(
+            value.type.body.typ, list
+        ):
+            value.type = IListType[value.type.vars[0], ir.types.Any]
+            return True
+        elif isinstance(value.type, ir.types.PyClass) and issubclass(
+            value.type.typ, list
+        ):
+            value.type = IListType[ir.types.Any, ir.types.Any]
             return True
         return False

src/kirin/dialects/ilist/runtime.py

@@ -40,9 +40,17 @@ def __str__(self) -> str:
         return f"IList({self.data})"
 
     def __iter__(self):
-        raise NotImplementedError("Cannot use IList outside kernel.")
+        return iter(self.data)
 
-    def __getitem__(self, index: int) -> T:
+    @overload
+    def __getitem__(self, index: slice) -> "IList[T, Any]": ...
+
+    @overload
+    def __getitem__(self, index: int) -> T: ...
+
+    def __getitem__(self, index: int | slice) -> T | "IList[T, Any]":
+        if isinstance(index, slice):
+            return IList(self.data[index])
         return self.data[index]
 
     def __eq__(self, value: object) -> bool:

src/kirin/ir/attrs/types.py

@@ -169,6 +169,9 @@ def is_subseteq_TypeVar(self, other: "TypeVar") -> bool:
     def __hash__(self) -> int:
         return hash((PyClass, self.typ))
 
+    def __repr__(self) -> str:
+        return self.typ.__name__
+
     def print_impl(self, printer: Printer) -> None:
         printer.plain_print("!py.", self.typ.__name__)
 
@@ -389,6 +392,12 @@ def is_subseteq_Generic(self, other: "Generic") -> bool:
     def __hash__(self) -> int:
         return hash((Generic, self.body, self.vars, self.vararg))
 
+    def __repr__(self) -> str:
+        if self.vararg is None:
+            return f"{self.body}[{', '.join(map(repr, self.vars))}]"
+        else:
+            return f"{self.body}[{', '.join(map(repr, self.vars))}, {self.vararg}, ...]"
+
     def print_impl(self, printer: Printer) -> None:
         printer.print(self.body)
         printer.plain_print("[")

src/kirin/prelude.py

@@ -2,7 +2,7 @@
 
 from kirin.ir import Method, dialect_group
 from kirin.passes import aggressive
-from kirin.dialects import cf, func, math
+from kirin.dialects import cf, func, math, ilist
 from kirin.dialects.py import (
     cmp,
     len,
@@ -49,7 +49,7 @@ def run_pass(mt: Method) -> None:
     return run_pass
 
 
-@dialect_group(python_basic.union([list, range, slice]))
+@dialect_group(python_basic.union([list, range, slice, cf, func, math]))
 def python_no_opt(self):
     """The Python dialect without optimization passes."""
 
@@ -59,19 +59,27 @@ def run_pass(mt: Method) -> None:
     return run_pass
 
 
-@dialect_group(python_no_opt.union([cf, func, math]))
+@dialect_group(python_basic.union([ilist, range, slice, cf, func, math]))
 def basic_no_opt(self):
-    """The basic kernel without optimization passes.
+    """The basic kernel without optimization passes. This is a builtin
+    eDSL that includes the basic dialects that are commonly used in
+    Python-like eDSLs.
 
     This eDSL includes the basic dialects without any optimization passes.
     Other eDSL can usually be built on top of this eDSL by utilizing the
     `basic_no_opt.add` method to add more dialects and optimization passes.
 
+    Note that unlike Python, list in this eDSL is immutable, and the
+    `append` method is not available. Use `+` operator to concatenate lists
+    instead. Immutable list is easier to optimize and reason about.
+
     See also [`basic`][kirin.prelude.basic] for the basic kernel with optimization passes.
+    See also [`ilist`][kirin.dialects.ilist] for the immutable list dialect.
     """
+    ilist_desugar = ilist.IListDesugar(self)
 
     def run_pass(mt: Method) -> None:
-        pass
+        ilist_desugar.fixpoint(mt)
 
     return run_pass
 
@@ -99,6 +107,7 @@ def main(x: int) -> int:
     ```
     """
     fold_pass = Fold(self)
+    ilist_desugar = ilist.IListDesugar(self)
     aggressive_fold_pass = aggressive.Fold(self)
     typeinfer_pass = TypeInfer(self)
 
@@ -122,6 +131,8 @@ def run_pass(
         if verify:
             mt.verify()
 
+        ilist_desugar.fixpoint(mt)
+
         if fold:
             if aggressive:
                 aggressive_fold_pass.fixpoint(mt)

test/analysis/dataflow/test_constprop.py

@@ -3,7 +3,7 @@
 from kirin.decl import info, statement
 from kirin.prelude import basic_no_opt
 from kirin.analysis import const
-from kirin.dialects import py
+from kirin.dialects import ilist
 
 
 class TestLattice:
@@ -279,10 +279,9 @@ def side_effect_true_branch_const(cond: bool):
 
 def test_non_pure_recursion():
     @basic_no_opt
-    def for_loop_append(cntr: int, x: list, n_range: int):
+    def for_loop_append(cntr: int, x: ilist.IList, n_range: int):
         if cntr < n_range:
-            py.Append(x, cntr)  # type: ignore
-            for_loop_append(cntr + 1, x, n_range)
+            for_loop_append(cntr + 1, x + [cntr], n_range)
 
         return x
 

test/dialects/pystmts/test_coll_add.py

@@ -1,5 +1,6 @@
 from kirin.ir import types
 from kirin.prelude import basic
+from kirin.dialects.ilist import IList, IListType
 
 
 @basic(typeinfer=True)
@@ -8,10 +9,10 @@ def tuple_new(x: int, xs: tuple):
 
 
 @basic(typeinfer=True)
-def list_new(x: int, xs: list):
+def list_new(x: int, xs: IList):
     return xs + [1, x]
 
 
 def test_tuple_add():
     assert tuple_new.return_type.is_subseteq(types.Tuple)
-    assert list_new.return_type.is_subseteq(types.List)
+    assert list_new.return_type.is_subseteq(IListType)

test/dialects/pystmts/test_getitem.py

@@ -1,5 +1,8 @@
+from typing import Any
+
 from kirin.ir import types
 from kirin.prelude import basic
+from kirin.dialects.ilist import IList, IListType
 
 
 @basic(typeinfer=True)
@@ -43,12 +46,12 @@ def tuple_const_err(xs: tuple[int, float, str]):
 
 
 @basic(typeinfer=True)
-def list_infer(xs: list[int], i: int):
+def list_infer(xs: IList[int, Any], i: int):
     return xs[i]
 
 
 @basic(typeinfer=True)
-def list_slice(xs: list[int], i: slice):
+def list_slice(xs: IList[int, Any], i: slice):
     return xs[i]
 
 
@@ -73,5 +76,5 @@ def test_getitem_typeinfer():
     assert tuple_err.return_type.is_equal(types.Bottom)
     assert tuple_const_err.return_type.is_equal(types.Bottom)
     assert list_infer.return_type.is_subseteq(types.Int)
-    assert list_slice.return_type.is_subseteq(types.List[types.Int])
+    assert list_slice.return_type.is_subseteq(IListType[types.Int])
     assert unknown.return_type.is_equal(types.Any)

test/program/py/test_list_append.py

@@ -1,11 +1,11 @@
 # type: ignore
-from kirin.prelude import basic
+from kirin.prelude import python_no_opt
 from kirin.dialects import py
 
 
 def test_list_append():
 
-    @basic
+    @python_no_opt
     def test_append():
         x = []
         py.Append(x, 1)
@@ -20,7 +20,7 @@ def test_append():
 
 
 def test_recursive_append():
-    @basic
+    @python_no_opt
     def for_loop_append(cntr: int, x: list, n_range: int):
         if cntr < n_range:
             py.Append(x, cntr)