fixing bug in which purity wasn't propagated properly (#460)

The following function is definitely pure:

```python
    @structural_no_opt
    def test_func(src: ilist.IList[float, Any]):

        def inner(i: int):
            if src[i] < 0:
                return 0.0
            elif src[i] < 1.0:
                return 1.0
            else:
                return 2.0

        return ilist.map(inner, ilist.range(len(src)))
```

but there is a bug in const prop when you have return values. The
propagation of frame purity was not occurring, making the constant
property of this function impure, which created other issues regarding
folding and DCE.

Author: weinbe58

src/kirin/dialects/scf/constprop.py

@@ -2,7 +2,6 @@
 
 from kirin import interp
 from kirin.analysis import const
-from kirin.dialects import func
 
 from .stmts import For, Yield, IfElse
 from ._dialect import dialect
@@ -41,12 +40,9 @@ def if_else(
             with interp_.new_frame(stmt, has_parent_access=True) as body_frame:
                 ret = interp_.frame_call_region(body_frame, stmt, body, cond)
             frame.entries.update(body_frame.entries)
-
-            if not body_frame.frame_is_not_pure and not isinstance(
-                body.blocks[0].last_stmt, func.Return
-            ):
-                frame.should_be_pure.add(stmt)
-            return ret
+            frame.frame_is_not_pure = (
+                frame.frame_is_not_pure or body_frame.frame_is_not_pure
+            )
         else:
             with interp_.new_frame(stmt, has_parent_access=True) as then_frame:
                 then_results = interp_.frame_call_region(
@@ -62,22 +58,28 @@ def if_else(
             # parent frame variables value except cond
             frame.entries.update(then_frame.entries)
             frame.entries.update(else_frame.entries)
+            # update frame purity
+            # if either frame is not pure, then the whole if-else is not pure
+            frame.frame_is_not_pure = (
+                frame.frame_is_not_pure
+                or then_frame.frame_is_not_pure
+                or else_frame.frame_is_not_pure
+            )
             # TODO: pick the non-return value
             if isinstance(then_results, interp.ReturnValue) and isinstance(
                 else_results, interp.ReturnValue
             ):
-                return interp.ReturnValue(then_results.value.join(else_results.value))
+                ret = interp.ReturnValue(then_results.value.join(else_results.value))
             elif isinstance(then_results, interp.ReturnValue):
                 ret = else_results
             elif isinstance(else_results, interp.ReturnValue):
                 ret = then_results
             else:
-                if not (
-                    then_frame.frame_is_not_pure is True
-                    or else_frame.frame_is_not_pure is True
-                ):
-                    frame.should_be_pure.add(stmt)
                 ret = interp_.join_results(then_results, else_results)
+
+        if not frame.frame_is_not_pure:
+            frame.should_be_pure.add(stmt)
+
         return ret
 
     @interp.impl(For)

test/dialects/scf/test_constprop.py

@@ -1,6 +1,8 @@
+from typing import Any
+
 from kirin.prelude import structural_no_opt
 from kirin.analysis import const
-from kirin.dialects import scf, func
+from kirin.dialects import scf, func, ilist
 
 prop = const.Propagate(structural_no_opt)
 
@@ -100,3 +102,155 @@ def simple_ifelse(x: int):
     assert isinstance(terminator, func.Return)
     assert isinstance(value := frame.entries[terminator.value], const.Value)
     assert value.data == 0
+
+
+def test_purity_1():
+
+    @structural_no_opt
+    def test_func(src: ilist.IList[float, Any]):
+
+        def inner(i: int):
+            if src[i] < 0:
+                return 0.0
+            elif src[i] < 1.0:
+                return 1.0
+            else:
+                return 2.0
+
+        return ilist.map(inner, ilist.range(len(src)))
+
+    frame, ret = prop.run(test_func)
+
+    assert not frame.frame_is_not_pure, "function should be pure"
+
+
+def test_purity_2():
+
+    @structural_no_opt
+    def test_func(src: ilist.IList[float, Any]):
+
+        def inner(i: int):
+            value = 0.0
+            if src[i] < 0:
+                value = 0.0
+            elif src[i] < 1.0:
+                return 1.0
+            else:
+                value = 2.0
+
+            return value
+
+        return ilist.map(inner, ilist.range(len(src)))
+
+    frame, ret = prop.run(test_func)
+
+    assert not frame.frame_is_not_pure, "function should be pure"
+
+
+def test_purity_3():
+
+    @structural_no_opt
+    def test_func(src: ilist.IList[float, Any]):
+
+        def inner(i: int):
+            value = 0.0
+            if src[i] < 0:
+                value = 0.0
+            elif src[i] < 1.0:
+                return 1.0
+            else:
+                return 2.0
+
+            return value
+
+        return ilist.map(inner, ilist.range(len(src)))
+
+    frame, ret = prop.run(test_func)
+
+    assert not frame.frame_is_not_pure, "function should be pure"
+
+
+def test_purity_4():
+
+    @structural_no_opt
+    def test_func(src: list[float]):
+
+        if True:
+            return src
+        else:
+            src.append(2.0)
+            return src
+
+    frame, ret = prop.run(test_func)
+
+    assert not frame.frame_is_not_pure, "function should be pure"
+
+
+def test_purity_5():
+
+    @structural_no_opt
+    def test_func(src: list[float]):
+
+        if False:
+            src.append(2.0)
+
+        return src
+
+    frame, ret = prop.run(test_func)
+
+    assert not frame.frame_is_not_pure, "function should be pure"
+
+
+def test_purity_6():
+
+    @structural_no_opt
+    def test_func(src: list[float]):
+
+        if True:
+            return src
+        else:
+            src.append(2.0)
+
+        return src
+
+    frame, ret = prop.run(test_func)
+
+    assert not frame.frame_is_not_pure, "function should be pure"
+
+
+def test_purity_7():
+
+    @structural_no_opt
+    def test_func(src: list[float], cond: bool):
+
+        if cond:
+            src.append(2.0)
+            return src
+        else:
+            return src
+
+    frame, ret = prop.run(test_func)
+
+    assert frame.frame_is_not_pure, "function should not be pure"
+
+
+def test_purity_8():
+
+    @structural_no_opt
+    def test_func(src: ilist.IList[float, Any], dst: ilist.IList[float, Any]):
+        assert len(src) == len(dst), "src and dst must have the same length"
+
+        def inner(i: int):
+            value = src[i]
+            if src[i] < dst[i]:
+                value = dst[i] - 3.0
+            elif src[i] > dst[i]:
+                return dst[i] + 3.0
+
+            return value
+
+        return ilist.map(inner, ilist.range(len(src)))
+
+    frame, ret = prop.run(test_func)
+
+    assert frame.frame_is_not_pure, "function should be pure"