Adding inline getitem rewrite for ilist (#446)

This is currently missing, it is definitely nice to have.

Author: weinbe58

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

@@ -2,3 +2,4 @@
 from .const import ConstList2IList as ConstList2IList
 from .unroll import Unroll as Unroll
 from .hint_len import HintLen as HintLen
+from .inline_getitem import InlineGetItem as InlineGetItem

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

@@ -0,0 +1,42 @@
+from kirin import ir
+from kirin.rewrite import abc
+from kirin.analysis import const
+from kirin.dialects import py
+
+from ..stmts import New
+
+
+class InlineGetItem(abc.RewriteRule):
+    """Rewrite rule to inline GetItem statements for IList.
+
+    For example if we have an `ilist.New` statement with a list of items,
+    and we can infer that the index used in `py.GetItem` is constant and within bounds,
+    we replace the `py.GetItem` with the ssa value in the list when the index is an integer
+    or with a new `ilist.New` statement containing the sliced items when the index is a slice.
+
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
+        if not isinstance(node, py.GetItem) or not isinstance(
+            stmt := node.obj.owner, New
+        ):
+            return abc.RewriteResult()
+
+        if not isinstance(index_const := node.index.hints.get("const"), const.Value):
+            return abc.RewriteResult()
+
+        index = index_const.data
+        if isinstance(index, int) and (
+            0 <= index < len(stmt.args) or -len(stmt.args) <= index < 0
+        ):
+            node.result.replace_by(stmt.args[index])
+            return abc.RewriteResult(has_done_something=True)
+        elif isinstance(index, slice):
+            start, stop, step = index.indices(len(stmt.args))
+            new_tuple = New(
+                tuple(stmt.args[start:stop:step]),
+            )
+            node.replace_by(new_tuple)
+            return abc.RewriteResult(has_done_something=True)
+        else:
+            return abc.RewriteResult()

test/dialects/test_ilist.py

@@ -3,6 +3,7 @@
 from kirin import ir, types, rewrite
 from kirin.passes import aggressive
 from kirin.prelude import basic_no_opt, python_basic
+from kirin.analysis import const
 from kirin.dialects import py, func, ilist, lowering
 from kirin.passes.typeinfer import TypeInfer
 
@@ -163,6 +164,41 @@ def const_range():
     assert const_range() == ilist.IList(range(0, 3))
 
 
+def test_inline_get_item():
+    items = tuple(ir.TestValue() for _ in range(2))
+
+    test_block = ir.Block(
+        [
+            qreg := ilist.New(values=items),
+            idx := py.Constant(0),
+            qubit_stmt := py.GetItem(obj=qreg.result, index=idx.result),
+            ilist.New(values=(qubit_stmt.result,)),
+            idx1 := py.Constant(10),
+            qubit_stmt := py.GetItem(obj=qreg.result, index=idx1.result),
+            ilist.New(values=(qubit_stmt.result,)),
+        ]
+    )
+
+    idx.result.hints["const"] = const.Value(0)
+    idx1.result.hints["const"] = const.Value(10)
+    rule = rewrite.Walk(ilist.rewrite.InlineGetItem())
+    rule.rewrite(test_block)
+
+    expected_block = ir.Block(
+        [
+            qreg := ilist.New(values=items),
+            idx := py.Constant(0),
+            qubit_stmt := py.GetItem(obj=qreg.result, index=idx.result),
+            ilist.New(values=(items[0],)),
+            idx1 := py.Constant(10),
+            qubit_stmt := py.GetItem(obj=qreg.result, index=idx1.result),
+            ilist.New(values=(qubit_stmt.result,)),
+        ]
+    )
+
+    assert test_block.is_equal(expected_block)
+
+
 def test_ilist_constprop():
     from kirin.analysis import const