Merge branch 'rafaelha/fix-slicing-reverse-edge-case' of https://github.com/QuEraComputing/kirin into rafaelha/inline_getitem_has_done_something_fix

Author: rafaelha

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

@@ -2,7 +2,6 @@
 from kirin.rewrite import abc
 from kirin.analysis import const
 from kirin.dialects import py
-from kirin.dialects.py.slice import SliceAttribute
 
 from ..stmts import New
 
@@ -35,10 +34,8 @@ def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         ):
             node.result.replace_by(stmt.args[index])
             return abc.RewriteResult(has_done_something=True)
-        elif isinstance(index, (slice, SliceAttribute)):
-            new_tuple = New(
-                tuple(stmt.args[index.start : index.stop : index.step]),
-            )
+        elif isinstance(index, slice):
+            new_tuple = New(tuple(stmt.args[index]))
             node.replace_by(new_tuple)
             return abc.RewriteResult(has_done_something=True)
         else:

src/kirin/dialects/py/indexing.py

@@ -98,17 +98,7 @@ class Concrete(interp.MethodTable):
 
     @interp.impl(GetItem)
     def getindex(self, interp, frame: interp.Frame, stmt: GetItem):
-        from kirin.dialects.py.slice import SliceAttribute
-
-        index = frame.get(stmt.index)
-
-        # need to handle special case of slice attribute
-        if isinstance(index, SliceAttribute):
-            index_value = index.unwrap()
-        else:
-            index_value = index
-
-        return (frame.get(stmt.obj)[index_value],)
+        return (frame.get(stmt.obj)[frame.get(stmt.index)],)
 
 
 @dialect.register(key="typeinfer")
@@ -218,23 +208,13 @@ def getitem(
             return (const.Unknown(),)
 
         if isinstance(obj, const.Value):
-            from kirin.dialects.py.slice import SliceAttribute
-
-            # need to handle special case of slice attribute
-            if isinstance(index.data, SliceAttribute):
-                index_value = index.data.unwrap()
-            else:
-                index_value = index.data
-
-            return (const.Value(obj.data[index_value]),)
-
+            return (const.Value(obj.data[index.data]),)
         elif isinstance(obj, const.PartialTuple):
             obj = obj.data
             if isinstance(index.data, int) and 0 <= index.data < len(obj):
                 return (obj[index.data],)
             elif isinstance(index.data, slice):
-                sl = index.data
-                return (const.PartialTuple(obj[sl.start : sl.stop : sl.step]),)
+                return (const.PartialTuple(obj[index.data]),)
         return (const.Unknown(),)
 
 

src/kirin/dialects/py/slice.py

@@ -13,7 +13,6 @@
 
 from kirin import ir, types, interp, lowering
 from kirin.decl import info, statement
-from kirin.print.printer import Printer
 from kirin.dialects.py.constant import Constant
 
 dialect = ir.Dialect("py.slice")
@@ -63,55 +62,18 @@ def __init__(
         )
 
 
-@dataclass
-class SliceAttribute(ir.Data[slice]):
-
-    start: int | None
-    stop: int | None
-    step: int | None
-
-    def __post_init__(self) -> None:
-        if self.start is None and self.step is None:
-            self.type = types.Slice[types.Literal(self.stop)]
-        else:
-            self.type = types.Slice3[
-                types.Literal(self.start),
-                types.Literal(self.stop),
-                types.Literal(self.step),
-            ]
-
-    def unwrap(self):
-        return slice(self.start, self.stop, self.step)
-
-    def __hash__(self):
-        return hash((type(self), self.start, self.stop, self.step))
-
-    def print_impl(self, printer: Printer) -> None:
-        return printer.plain_print(f"slice({self.start}, {self.stop}, {self.step})")
-
-    def is_structurally_equal(
-        self, other: ir.Attribute, context: dict | None = None
-    ) -> bool:
-        return (
-            isinstance(other, SliceAttribute)
-            and self.start == other.start
-            and self.stop == other.stop
-            and self.step == other.step
-        )
-
-
 @dialect.register
 class Concrete(interp.MethodTable):
 
     @interp.impl(Slice)
     def _slice(self, interp, frame: interp.Frame, stmt: Slice):
         start, stop, step = frame.get_values(stmt.args)
         if start is None and step is None:
-            return (SliceAttribute(None, stop, None),)
+            return (slice(stop),)
         elif step is None:
-            return (SliceAttribute(start, stop, None),)
+            return (slice(start, stop),)
         else:
-            return (SliceAttribute(start, stop, step),)
+            return (slice(start, stop, step),)
 
 
 @dialect.register

src/kirin/rewrite/getitem.py

@@ -4,7 +4,6 @@
 from kirin.analysis import const
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
-from kirin.dialects.py.slice import SliceAttribute
 
 
 @dataclass
@@ -30,10 +29,8 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
         ):
             node.result.replace_by(stmt.args[index])
             return RewriteResult(has_done_something=True)
-        elif isinstance(index, (slice, SliceAttribute)):
-            new_tuple = py.tuple.New(
-                tuple(stmt.args[index.start : index.stop : index.step]),
-            )
+        elif isinstance(index, slice):
+            new_tuple = py.tuple.New(tuple(stmt.args[index]))
             node.replace_by(new_tuple)
             return RewriteResult(has_done_something=True)
         else:

src/kirin/types.py

@@ -28,7 +28,6 @@
 NoneType = PyClass(type(None))
 List = Generic(list, TypeVar("T"))
 Slice = Generic(slice, TypeVar("T"))
-Slice3 = Generic(slice, TypeVar("T1"), TypeVar("T2"), TypeVar("T3"))
 Tuple = Generic(tuple, Vararg(TypeVar("T")))
 Dict = Generic(dict, TypeVar("K"), TypeVar("V"))
 Set = Generic(set, TypeVar("T"))

test/dialects/ilist/test_inline_getitem.py

@@ -33,7 +33,6 @@ def func(x: int):
 
     assert before == after
     assert len(func.callable_region.blocks[0].stmts) == 1
-    print(func.code.print())
 
 
 @pytest.mark.parametrize(
@@ -54,8 +53,6 @@ def func():
         ylist = ilist.New(values=(0, 1, 2, 3, 4), elem_type=types.PyClass(int))
         return ylist[sl]
 
-    func.code.print()
-
     stmt_types = [type(stmt) for stmt in func.callable_region.blocks[0].stmts]
     assert GetItem in stmt_types
 
@@ -90,16 +87,13 @@ def func():
     assert GetItem in stmt_types
 
     before = func()
-    func.code.print()
 
     apply_getitem_optimization(func)
 
     stmt_types = [type(stmt) for stmt in func.callable_region.blocks[0].stmts]
     assert GetItem not in stmt_types
     after = func()
 
-    func.code.print()
-
     assert before == after
 
 

test/dialects/pystmts/test_slice.py

@@ -1,7 +1,6 @@
 from kirin import types
 from kirin.prelude import basic_no_opt
-from kirin.dialects import py, ilist
-from kirin.dialects.py.slice import SliceAttribute
+from kirin.dialects import py
 
 
 @basic_no_opt
@@ -45,32 +44,3 @@ def test_wrong_slice():
 
     stmt: py.slice.Slice = wrong_slice.code.body.blocks[0].stmts.at(7)
     assert stmt.result.type.is_subseteq(types.Bottom)
-
-
-def test_slice_attr():
-
-    @basic_no_opt
-    def test():
-
-        return (slice(0, 20), slice(30), slice(1, 40, 5))
-
-    result = test()
-    assert result == (
-        SliceAttribute(0, 20, None),
-        SliceAttribute(None, 30, None),
-        SliceAttribute(1, 40, 5),
-    )
-
-
-def test_slice_attr_hash():
-    assert hash(SliceAttribute(0, 20, None)) == hash((SliceAttribute, 0, 20, None))
-
-
-def test_slice_get_index():
-    @basic_no_opt
-    def test():
-        x = slice(0, 20, None)
-        y = range(40)
-        return y[x]
-
-    assert test() == ilist.IList(range(0, 20, 1))

test/rules/test_getitem.py

@@ -30,7 +30,6 @@ def func(x: int):
 
     assert before == after
     assert len(func.callable_region.blocks[0].stmts) == 1
-    print(func.code.print())
 
 
 @pytest.mark.parametrize(
@@ -51,8 +50,6 @@ def func():
         ylist = (0, 1, 2, 3, 4)
         return ylist[sl]
 
-    func.code.print()
-
     stmt_types = [type(stmt) for stmt in func.callable_region.blocks[0].stmts]
     assert GetItem in stmt_types
 
@@ -87,16 +84,13 @@ def func():
     assert GetItem in stmt_types
 
     before = func()
-    func.code.print()
 
     apply_getitem_optimization(func)
 
     stmt_types = [type(stmt) for stmt in func.callable_region.blocks[0].stmts]
     assert GetItem not in stmt_types
     after = func()
 
-    func.code.print()
-
     assert before == after