Roger/backport ilist sorted (#497)

cc: @cduck

---------

Co-authored-by: David Plankensteiner <[email protected]>
Co-authored-by: Casey Duckering <[email protected]>

Author: 3 people

src/kirin/dialects/ilist/__init__.py

@@ -31,5 +31,6 @@
     foldl as foldl,
     foldr as foldr,
     range as range,
+    sorted as sorted,
     for_each as for_each,
 )

src/kirin/dialects/ilist/_wrapper.py

@@ -65,3 +65,36 @@ def for_each(
     fn: typing.Callable[[ElemT], typing.Any],
     collection: IList[ElemT, LenT] | list[ElemT],
 ) -> None: ...
+
+
+@typing.overload
+def sorted(collection: IList[ElemT, LenT] | list[ElemT]) -> IList[ElemT, LenT]: ...
+
+
+@typing.overload
+def sorted(
+    collection: IList[ElemT, LenT] | list[ElemT], reverse: bool
+) -> IList[ElemT, LenT]: ...
+
+
+@typing.overload
+def sorted(
+    collection: IList[ElemT, LenT] | list[ElemT],
+    key: typing.Callable[[ElemT], OutElemT],
+) -> IList[ElemT, LenT]: ...
+
+
+@typing.overload
+def sorted(
+    collection: IList[ElemT, LenT] | list[ElemT],
+    key: typing.Callable[[ElemT], OutElemT],
+    reverse: bool,
+) -> IList[ElemT, LenT]: ...
+
+
+@lowering.wraps(stmts.Sorted)
+def sorted(
+    collection: IList[ElemT, LenT] | list[ElemT],
+    key: typing.Callable[[ElemT], OutElemT],
+    reverse: bool,
+) -> IList[ElemT, LenT]: ...

src/kirin/dialects/ilist/interp.py

@@ -3,7 +3,7 @@
 from kirin.dialects.py.len import Len
 from kirin.dialects.py.binop import Add
 
-from .stmts import Map, New, Push, Scan, Foldl, Foldr, Range, ForEach
+from .stmts import Map, New, Push, Scan, Foldl, Foldr, Range, Sorted, ForEach
 from .runtime import IList
 from ._dialect import dialect
 
@@ -91,4 +91,12 @@ def for_each(self, interp: Interpreter, frame: Frame, stmt: ForEach):
         for elem in coll.data:
             # NOTE: assume fn has been type checked
             interp.run_method(fn, (elem,))
-        return (None,)
+        return
+
+    @impl(Sorted)
+    def sorted(self, inter: Interpreter, frame: Frame, stmt: Sorted):
+        key = frame.get(stmt.key)
+        reverse: bool = frame.get(stmt.reverse)
+        coll: IList = frame.get(stmt.collection)
+
+        return (IList(data=sorted(coll.data, key=key, reverse=reverse)),)

src/kirin/dialects/ilist/lowering.py

@@ -1,6 +1,8 @@
 import ast
+from dataclasses import dataclass
 
 from kirin import types, lowering
+from kirin.dialects import py
 
 from . import stmts as ilist
 from ._dialect import dialect
@@ -48,3 +50,43 @@ def lower_Call_IList(
             return state.current_frame.push(stmt)
         else:
             return self.lower_List(state, value)
+
+
+@dataclass(frozen=True)
+class SortedLowering(lowering.FromPythonCall["ilist.Sorted"]):
+    """
+    Custom lowering for Sorted to deal with optional arguments `key=None` and `reverse=False`
+    """
+
+    def lower(
+        self, stmt: type["ilist.Sorted"], state: lowering.State[ast.AST], node: ast.Call
+    ) -> lowering.Result:
+        args = node.args
+
+        if len(args) != 1:
+            raise lowering.BuildError(
+                f"Expected single argument in sorted, got {len(args)}"
+            )
+        collection = state.lower(args[0]).expect_one()
+
+        key = None
+        reverse = None
+        for kwarg in node.keywords:
+            if kwarg.arg == "key":
+                key = state.lower(kwarg.value).expect_one()
+            elif kwarg.arg == "reverse":
+                reverse = state.lower(kwarg.value).expect_one()
+            else:
+                raise lowering.BuildError(
+                    f"Got unexpected keyword argument in sorted {kwarg.arg}"
+                )
+
+        if key is None:
+            key = state.current_frame.push(py.Constant(None)).result
+
+        if reverse is None:
+            reverse = state.current_frame.push(py.Constant(False)).result
+
+        return state.current_frame.push(
+            stmt(collection=collection, key=key, reverse=reverse)
+        )

src/kirin/dialects/ilist/stmts.py

@@ -5,6 +5,7 @@
 
 from .runtime import IList
 from ._dialect import dialect
+from .lowering import SortedLowering
 
 ElemT = types.TypeVar("ElemT")
 ListLen = types.TypeVar("ListLen")
@@ -112,3 +113,15 @@ class ForEach(ir.Statement):
     purity: bool = info.attribute(default=False)
     fn: ir.SSAValue = info.argument(types.Generic(ir.Method, [ElemT], types.NoneType))
     collection: ir.SSAValue = info.argument(IListType[ElemT])
+
+
+@statement(dialect=dialect)
+class Sorted(ir.Statement):
+    traits = frozenset({ir.MaybePure(), SortedLowering()})
+    purity: bool = info.attribute(default=False)
+    collection: ir.SSAValue = info.argument(IListType[ElemT, ListLen])
+    key: ir.SSAValue = info.argument(
+        types.Union((types.Generic(ir.Method, [ElemT], ElemT), types.NoneType))
+    )
+    reverse: ir.SSAValue = info.argument(types.Bool)
+    result: ir.ResultValue = info.result(IListType[ElemT, ListLen])

test/dialects/test_ilist_wrapper.py

@@ -75,3 +75,46 @@ def scan_wrap():
         10 + 1 + 1 + 1 + 3,
         10 + 1 + 1 + 1 + 1 + 4,
     ]
+
+
+def test_sorted():
+    def key_test(a: int) -> int:
+        return a
+
+    @basic
+    def main():
+        ls = [2, 3, 1, 5, 4]
+        return (
+            ilist.sorted(ls),
+            ilist.sorted(ls, key=key_test),
+            ilist.sorted(ls, reverse=True),
+        )
+
+    main.print()
+
+    ls1, ls2, ls3 = main()
+    assert ls1.data == [1, 2, 3, 4, 5]
+    assert ls2.data == ls1.data
+    assert ls3.data == [5, 4, 3, 2, 1]
+
+    def first(x: tuple[str, int]) -> str:
+        return x[0]
+
+    def second(x: tuple[str, int]) -> int:
+        return x[1]
+
+    @basic
+    def main2():
+        ls = [("a", 4), ("b", 3), ("c", 1)]
+        return (
+            ilist.sorted(ls, key=first),
+            ilist.sorted(ls, key=second),
+            ilist.sorted(ls, key=second, reverse=True),
+        )
+
+    main2.print()
+
+    ls1, ls2, ls3 = main2()
+    assert ls1.data == [("a", 4), ("b", 3), ("c", 1)]
+    assert ls3.data == ls1.data
+    assert ls2.data == [("c", 1), ("b", 3), ("a", 4)]