try to fix some scf bugs (#184)

SCF interpreter is still a bit buggy. Trying to switch to scf to see
what break here.

UPDATE: Ok I see the problem here, we cannot take the similar route as
xDSL because they don't aim to deal with python semantics. It works if
we follow MLIR convention, but if we have to support Python. We need to
allow `scf.IfElse` statement accept other terminators, e.g
`func.Return`, I don't see an easy way transforming to `Yield` in this
case because we are doing more dynamic semantics (e.g the return value
can be a union of different sized tuple), making the following code
seems not possible to have an `yield` equivalent form

```mlir
%1, %2 = scf.if a > b {
   func.return a
} else {
   y = b + 1
   scf.yield a,y
}
```

we cannot lift the `return` statement outside the if statement here...
so directly supporting this would require `run_ssacfg_region` to forward
the `ReturnValue` result until we pop the current function frame.

Author: Roger-luo

src/kirin/analysis/__init__.py

@@ -15,6 +15,10 @@
 
 from kirin.analysis import const as const
 from kirin.analysis.cfg import CFG as CFG
-from kirin.analysis.forward import Forward as Forward, ForwardExtra as ForwardExtra
+from kirin.analysis.forward import (
+    Forward as Forward,
+    ForwardExtra as ForwardExtra,
+    ForwardFrame as ForwardFrame,
+)
 from kirin.analysis.callgraph import CallGraph as CallGraph
 from kirin.analysis.typeinfer import TypeInference as TypeInference

src/kirin/analysis/const/__init__.py

@@ -8,7 +8,7 @@
 the IR.
 """
 
-from .prop import Propagate as Propagate
+from .prop import Propagate as Propagate, ExtraFrameInfo as ExtraFrameInfo
 from .lattice import (
     Pure as Pure,
     Value as Value,

src/kirin/analysis/const/prop.py

@@ -59,19 +59,20 @@ def _try_eval_const_pure(
             _frame = self._interp.new_frame(frame.code)
             _frame.set_values(stmt.args, tuple(x.data for x in values))
             value = self._interp.eval_stmt(_frame, stmt)
-            if isinstance(value, tuple):
-                return tuple(JointResult(Value(each), Pure()) for each in value)
-            elif isinstance(value, interp.ReturnValue):
-                return interp.ReturnValue(
-                    *tuple(JointResult(Value(x), Pure()) for x in value.results)
-                )
-            elif isinstance(value, interp.Successor):
-                return interp.Successor(
-                    value.block,
-                    *tuple(
-                        JointResult(Value(each), Pure()) for each in value.block_args
-                    ),
-                )
+            match value:
+                case tuple():
+                    return tuple(JointResult(Value(each), Pure()) for each in value)
+                case interp.ReturnValue(ret):
+                    return interp.ReturnValue(JointResult(Value(ret), Pure()))
+                case interp.YieldValue(yields):
+                    return interp.YieldValue(
+                        tuple(JointResult(Value(each), Pure()) for each in yields)
+                    )
+                case interp.Successor(block, args):
+                    return interp.Successor(
+                        block,
+                        *tuple(JointResult(Value(each), Pure()) for each in args),
+                    )
         except interp.InterpreterError:
             pass
         return (self.void,)
@@ -104,8 +105,11 @@ def _set_frame_not_pure(self, result: interp.StatementResult[JointResult]):
         if isinstance(result, tuple) and all(x.purity is Pure() for x in result):
             return
 
-        if isinstance(result, interp.ReturnValue) and all(
-            x.purity is Pure() for x in result.results
+        if isinstance(result, interp.ReturnValue) and isinstance(result.value, Pure):
+            return
+
+        if isinstance(result, interp.YieldValue) and all(
+            isinstance(x, Pure) for x in result
         ):
             return
 

src/kirin/dialects/cf/__init__.py

@@ -5,7 +5,6 @@
 )
 from kirin.dialects.cf.lower import CfLowering as CfLowering
 from kirin.dialects.cf.stmts import (
-    Assert as Assert,
     Branch as Branch,
     ConditionalBranch as ConditionalBranch,
 )

src/kirin/dialects/cf/constprop.py

@@ -1,16 +1,12 @@
 from kirin.interp import FrameABC, Successor, MethodTable, impl
 from kirin.analysis import const
-from kirin.dialects.cf.stmts import Assert, Branch, ConditionalBranch
+from kirin.dialects.cf.stmts import Branch, ConditionalBranch
 from kirin.dialects.cf.dialect import dialect
 
 
 @dialect.register(key="constprop")
 class DialectConstProp(MethodTable):
 
-    @impl(Assert)
-    def assert_stmt(self, interp: const.Propagate, frame, stmt: Assert):
-        return ()
-
     @impl(Branch)
     def branch(
         self, interp: const.Propagate, frame: FrameABC[const.JointResult], stmt: Branch

src/kirin/dialects/cf/emit.py

@@ -4,7 +4,7 @@
 from kirin.interp import Successor, MethodTable, impl
 from kirin.emit.julia import EmitJulia
 
-from .stmts import Assert, Branch, ConditionalBranch
+from .stmts import Branch, ConditionalBranch
 from .dialect import dialect
 
 IO_t = TypeVar("IO_t", bound=IO)
@@ -13,15 +13,6 @@
 @dialect.register(key="emit.julia")
 class JuliaMethodTable(MethodTable):
 
-    @impl(Assert)
-    def emit_assert(
-        self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Assert
-    ):
-        interp.writeln(
-            frame, f"@assert {frame.get(stmt.condition)} {frame.get(stmt.message)}"
-        )
-        return ()
-
     @impl(Branch)
     def emit_branch(
         self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Branch

src/kirin/dialects/cf/interp.py

@@ -1,21 +1,11 @@
-from kirin.interp import Frame, Successor, Interpreter, MethodTable, WrapException, impl
-from kirin.dialects.cf.stmts import Assert, Branch, ConditionalBranch
+from kirin.interp import Frame, Successor, Interpreter, MethodTable, impl
+from kirin.dialects.cf.stmts import Branch, ConditionalBranch
 from kirin.dialects.cf.dialect import dialect
 
 
 @dialect.register
 class CfInterpreter(MethodTable):
 
-    @impl(Assert)
-    def assert_stmt(self, interp: Interpreter, frame: Frame, stmt: Assert):
-        if frame.get(stmt.condition) is True:
-            return ()
-
-        if stmt.message:
-            raise WrapException(AssertionError(frame.get(stmt.message)))
-        else:
-            raise WrapException(AssertionError("Assertion failed"))
-
     @impl(Branch)
     def branch(self, interp: Interpreter, frame: Frame, stmt: Branch):
         return Successor(stmt.successor, *frame.get_values(stmt.arguments))

src/kirin/dialects/cf/lower.py

@@ -17,18 +17,6 @@ def lower_Pass(self, state: LoweringState, node: ast.Pass) -> Result:
         state.append_stmt(cf.Branch(arguments=(), successor=next))
         return Result()
 
-    def lower_Assert(self, state: LoweringState, node: ast.Assert) -> Result:
-        from kirin.dialects.py.constant import Constant
-
-        cond = state.visit(node.test).expect_one()
-        if node.msg:
-            message = state.visit(node.msg).expect_one()
-            state.append_stmt(cf.Assert(condition=cond, message=message))
-        else:
-            message_stmt = state.append_stmt(Constant(""))
-            state.append_stmt(cf.Assert(condition=cond, message=message_stmt.result))
-        return Result()
-
     def lower_If(self, state: LoweringState, node: ast.If) -> Result:
         cond = state.visit(node.test).expect_one()
 

src/kirin/dialects/cf/stmts.py

@@ -1,30 +1,10 @@
 from kirin.ir import Block, SSAValue, Statement, IsTerminator
 from kirin.decl import info, statement
 from kirin.print.printer import Printer
-from kirin.ir.attrs.types import Bool, String
+from kirin.ir.attrs.types import Bool
 from kirin.dialects.cf.dialect import dialect
 
 
-@statement(dialect=dialect)
-class Assert(Statement):
-    name = "assert"
-    traits = frozenset({})
-
-    condition: SSAValue
-    message: SSAValue = info.argument(String)
-
-    def print_impl(self, printer: Printer) -> None:
-        with printer.rich(style="keyword"):
-            printer.print_name(self)
-
-        printer.plain_print(" ")
-        printer.print(self.condition)
-
-        if self.message:
-            printer.plain_print(", ")
-            printer.print(self.message)
-
-
 @statement(dialect=dialect)
 class Branch(Statement):
     name = "br"

src/kirin/dialects/cf/typeinfer.py

@@ -1,17 +1,12 @@
-from kirin.ir import types
 from kirin.interp import Successor, MethodTable, AbstractFrame, impl
-from kirin.dialects.cf.stmts import Assert, Branch, ConditionalBranch
+from kirin.dialects.cf.stmts import Branch, ConditionalBranch
 from kirin.analysis.typeinfer import TypeInference
 from kirin.dialects.cf.dialect import dialect
 
 
 @dialect.register(key="typeinfer")
 class TypeInfer(MethodTable):
 
-    @impl(Assert)
-    def assert_stmt(self, interp: TypeInference, frame: AbstractFrame, stmt: Assert):
-        return (types.Bottom,)
-
     @impl(Branch)
     def branch(self, interp: TypeInference, frame: AbstractFrame, stmt: Branch):
         frame.worklist.append(