Two Fixes to constant prop, IList (#524)

changes in this PR:

1. Bug fixes in `ilist` dialect that gaurd against non-pure execution of
`map` `for_each` `foldl`, etc.
~~2. Adding a try-except inside `try_eval_const_pure` in cases where the
concrete interpreter is missing an implementation of a statement that is
pure, falling back to returning `Unknown` for all results.~~

Author: weinbe58

src/kirin/dialects/ilist/constprop.py

@@ -31,7 +31,7 @@ def one_args(
         # 1. if the function is a constant method, and the method is pure, then the map is pure
         if isinstance(fn, const.Value) and isinstance(method := fn.data, ir.Method):
             self.detect_purity(interp_, frame, stmt, method.code, (fn, const.Unknown()))
-            if isinstance(collection, const.Value):
+            if isinstance(collection, const.Value) and stmt in frame.should_be_pure:
                 return interp_.try_eval_const_pure(frame, stmt, (fn, collection))
         elif isinstance(fn, const.PartialLambda):
             self.detect_purity(interp_, frame, stmt, fn.code, (fn, const.Unknown()))
@@ -57,7 +57,11 @@ def two_args(self, interp_: const.Propagate, frame: const.Frame, stmt: Foldl):
                 method.code,
                 (fn, const.Unknown(), const.Unknown()),
             )
-            if isinstance(collection, const.Value) and isinstance(init, const.Value):
+            if (
+                isinstance(collection, const.Value)
+                and isinstance(init, const.Value)
+                and stmt in frame.should_be_pure
+            ):
                 return interp_.try_eval_const_pure(frame, stmt, (fn, collection, init))
         elif isinstance(fn, const.PartialLambda):
             self.detect_purity(

test/dialects/test_ilist.py

@@ -1,10 +1,12 @@
 from typing import Any, Literal
 
 from kirin import ir, types, rewrite
+from kirin.decl import info, statement
 from kirin.passes import aggressive
 from kirin.prelude import structural, basic_no_opt, python_basic
 from kirin.analysis import const
 from kirin.dialects import py, func, ilist, lowering
+from kirin.lowering import FromPythonCall
 from kirin.passes.typeinfer import TypeInfer
 
 
@@ -426,6 +428,31 @@ def main2():
     assert target.data == (6, 6)
 
 
+def test_ilist_constprop_non_pure():
+
+    new_dialect = ir.Dialect("test")
+
+    @statement(dialect=new_dialect)
+    class DefaultInit(ir.Statement):
+        name = "test"
+        traits = frozenset({FromPythonCall()})
+        result: ir.ResultValue = info.result(types.Float)
+
+    dialect_group = basic_no_opt.add(new_dialect)
+
+    @dialect_group
+    def test():
+
+        def inner(_: int):
+            return DefaultInit()
+
+        return ilist.map(inner, ilist.range(10))
+
+    _, res = const.Propagate(dialect_group).run(test)
+
+    assert isinstance(res, const.Unknown)
+
+
 rule = rewrite.Fixpoint(rewrite.Walk(ilist.rewrite.Unroll()))
 xs = ilist.IList([1, 2, 3])