Revert "fixing bug in which purity wasn't propagated properly" (#462)

Reverts #460

Turns out that semantically you have to make the IfElse body not pure if
there is a return value because of a weird interaction with
ConstantFold. Namely the fact that the return value of the if statement
has no uses, and so it will be deleted.

Author: weinbe58

src/kirin/dialects/scf/constprop.py

@@ -2,6 +2,7 @@
 
 from kirin import interp
 from kirin.analysis import const
+from kirin.dialects import func
 
 from .stmts import For, Yield, IfElse
 from ._dialect import dialect
@@ -40,9 +41,12 @@ 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)
-            frame.frame_is_not_pure = (
-                frame.frame_is_not_pure or body_frame.frame_is_not_pure
-            )
+
+            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
         else:
             with interp_.new_frame(stmt, has_parent_access=True) as then_frame:
                 then_results = interp_.frame_call_region(
@@ -58,28 +62,22 @@ 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
             ):
-                ret = interp.ReturnValue(then_results.value.join(else_results.value))
+                return 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,8 +1,6 @@
-from typing import Any
-
 from kirin.prelude import structural_no_opt
 from kirin.analysis import const
-from kirin.dialects import scf, func, ilist
+from kirin.dialects import scf, func
 
 prop = const.Propagate(structural_no_opt)
 
@@ -102,155 +100,3 @@ 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"