add wrappers for ilist (#215)

Author: kaihsin

src/kirin/dialects/ilist/__init__.py

@@ -24,3 +24,4 @@
 from .passes import IListDesugar as IListDesugar
 from .runtime import IList as IList
 from ._dialect import dialect as dialect
+from ._wrapper import map as map, scan as scan, foldl as foldl, foldr as foldr

src/kirin/dialects/ilist/_wrapper.py

@@ -0,0 +1,39 @@
+from typing import Any, TypeVar, Iterable
+
+from kirin import ir
+from kirin.lowering import wraps
+
+from . import stmts
+from .runtime import IList
+
+T_Elem = TypeVar("T_Elem")
+T_Out = TypeVar("T_Out")
+T_Result = TypeVar("T_Result")
+
+
+@wraps(stmts.Map)
+def map(fn: ir.Method[[T_Elem], T_Out], collection: Iterable) -> IList | list: ...
+
+
+@wraps(stmts.Foldr)
+def foldr(
+    fn: ir.Method[[T_Elem, T_Out], T_Out], collection: Iterable, init: T_Out
+) -> T_Out: ...
+
+
+@wraps(stmts.Foldl)
+def foldl(
+    fn: ir.Method[[T_Out, T_Elem], T_Out], collection: Iterable, init: T_Out
+) -> T_Out: ...
+
+
+@wraps(stmts.Scan)
+def scan(
+    fn: ir.Method[[T_Out, T_Elem], tuple[T_Out, T_Result]],
+    collection: Iterable,
+    init: T_Out,
+) -> tuple[T_Out, IList[T_Result, Any]]: ...
+
+
+@wraps(stmts.ForEach)
+def for_each(fn: ir.Method[[T_Elem], None], collection: Iterable) -> None: ...

test/dialects/test_ilist_wrapper.py

@@ -0,0 +1,77 @@
+from kirin.prelude import basic
+from kirin.dialects import ilist
+
+
+def test_map_wrapper():
+
+    @basic
+    def add1(x: int):
+        return x + 1
+
+    @basic
+    def map_wrap():
+        return ilist.map(add1, range(5))
+
+    out = map_wrap()
+    assert isinstance(out, ilist.IList)
+    assert out.data == [1, 2, 3, 4, 5]
+
+
+def test_foldr_wrapper():
+
+    @basic
+    def add_fold(x: int, out: int):
+        return out + x
+
+    @basic
+    def map_foldr():
+        return ilist.foldr(add_fold, range(5), init=10)
+
+    out = map_foldr()
+    assert isinstance(out, int)
+    assert out == 10 + 0 + 1 + 2 + 3 + 4
+
+
+def test_foldl_wrapper():
+
+    @basic
+    def add_fold2(out: int, x: int):
+        return out + x
+
+    @basic
+    def map_foldl():
+        return ilist.foldr(add_fold2, range(5), init=10)
+
+    out = map_foldl()
+    assert isinstance(out, int)
+    assert out == 10 + 0 + 1 + 2 + 3 + 4
+
+
+def test_scan_wrapper():
+
+    @basic
+    def add_scan(out: int, x: int):
+        return out + 1, out + x
+
+    @basic
+    def scan_wrap():
+        return ilist.scan(add_scan, range(5), init=10)
+
+    out = scan_wrap()
+    assert isinstance(out, tuple)
+    assert len(out) == 2
+
+    res = out[0]
+    out_list = out[1]
+
+    assert isinstance(res, int)
+    assert res == 10 + 1 * 5
+
+    assert isinstance(out_list, ilist.IList)
+    assert out_list.data == [
+        10 + 0,
+        10 + 1 + 1,
+        10 + 1 + 1 + 2,
+        10 + 1 + 1 + 1 + 3,
+        10 + 1 + 1 + 1 + 1 + 4,
+    ]