structural control flow (#176)

This PR implements a Python-like SCF dialect that can be a lowering
target from Python. To support Python-like multiple loop variables, a
new dialect is also added as `unpack`. The support of assignment is
still not implemented in this PR, we can add that in the future (so you
can only unpack in the loop head in this PR)

Unlike MLIR SCF, this dialect has less restrictions on the iterable (it
is typed `Any`), this will have issues in type inference where we need
to implicitly call an `eltype` method to query the element type of a
iterable, but I think this what Python loops expect anyways.

To use a `affine.for`-like semantics, one can always rewrite this
Python-based `for` statement into `affine.for` when constant fold finds
the loop iterable is a `range`.

I realize to support these statements, we still need to let region
evaluation return multiple values, so the API of
`BaseInterpreter.run_ssa_cfg_region` (I hope we can move this into a
trait instead of inside interpreter so we can support other ways of
execution) is changed into returning `tuple[ValueType, ...]` and in the
case of empty return, it returns `()` instead of `self.void`.

The `run_callable_region` method now checks if region gives `()` and
returns `void` to match Python semantics (given we don't allow functions
return a tuple directly). I think this is mostly non-breaking but let me
know if you are using this API (@kaihsin @weinbe58 )

example:

```python
from kirin.prelude import python_basic
from kirin.dialects import py, scf, func, ilist

xs = ilist.IList([(1, 2), (3, 4)])

@python_basic.union([func, scf, py.range, py.unpack, ilist])
def main(x):
    for a, b in xs:
        x = x + a
    return x

main.print()
```

<img width="680" alt="image"
src="https://github.com/user-attachments/assets/6b9a9e5d-d561-40da-a6f6-3cf7607df138"
/>


supercedes #101

Author: Roger-luo

example/beer/rewrite.py

@@ -3,8 +3,8 @@
 from stmts import RandomBranch
 
 from kirin.dialects import cf
+from kirin.rewrite.abc import RewriteRule, RewriteResult
 from kirin.ir.nodes.stmt import Statement
-from kirin.rewrite.abc import RewriteResult, RewriteRule
 
 
 @dataclass

scripts/release.py

@@ -1,6 +1,6 @@
-import argparse
-import logging
 import os
+import logging
+import argparse
 import subprocess
 
 import tomlkit

src/kirin/analysis/const/prop.py

@@ -62,7 +62,9 @@ def _try_eval_const_pure(
             if isinstance(value, tuple):
                 return tuple(JointResult(Value(each), Pure()) for each in value)
             elif isinstance(value, interp.ReturnValue):
-                return interp.ReturnValue(JointResult(Value(value.result), Pure()))
+                return interp.ReturnValue(
+                    *tuple(JointResult(Value(x), Pure()) for x in value.results)
+                )
             elif isinstance(value, interp.Successor):
                 return interp.Successor(
                     value.block,
@@ -102,7 +104,9 @@ 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 result.result.purity is Pure():
+        if isinstance(result, interp.ReturnValue) and all(
+            x.purity is Pure() for x in result.results
+        ):
             return
 
         if isinstance(result, interp.Successor) and all(

src/kirin/dialects/py/__init__.py

@@ -11,6 +11,7 @@
     unary as unary,
     assign as assign,
     boolop as boolop,
+    unpack as unpack,
     builtin as builtin,
     constant as constant,
     indexing as indexing,

src/kirin/dialects/py/unpack.py

@@ -0,0 +1,79 @@
+import ast
+
+from kirin import ir, interp, lowering
+from kirin.decl import info, statement
+from kirin.print import Printer
+from kirin.exceptions import DialectLoweringError
+
+dialect = ir.Dialect("py.unpack")
+
+
+@statement(dialect=dialect, init=False)
+class Unpack(ir.Statement):
+    value: ir.SSAValue = info.argument(ir.types.Any)
+    names: tuple[str | None, ...] = info.attribute(property=True)
+
+    def __init__(self, value: ir.SSAValue, names: tuple[str | None, ...]):
+        result_types = [ir.types.Any] * len(names)
+        super().__init__(
+            args=(value,),
+            result_types=result_types,
+            args_slice={"value": 0},
+            properties={"names": ir.PyAttr(names)},
+        )
+        for result, name in zip(self.results, names):
+            result.name = name
+
+    def print_impl(self, printer: Printer) -> None:
+        printer.print_name(self)
+        printer.plain_print(" ")
+        printer.print(self.value)
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Unpack)
+    def unpack(self, interp: interp.Interpreter, frame: interp.Frame, stmt: Unpack):
+        return tuple(frame.get(stmt.value))
+
+
+@dialect.register(key="typeinfer")
+class TypeInfer(interp.MethodTable):
+
+    @interp.impl(Unpack)
+    def unpack(self, interp, frame: interp.Frame[ir.types.TypeAttribute], stmt: Unpack):
+        value = frame.get(stmt.value)
+        if isinstance(value, ir.types.Generic) and value.is_subseteq(ir.types.Tuple):
+            if value.vararg:
+                rest = tuple(value.vararg.typ for _ in stmt.names[len(value.vars) :])
+                return tuple(value.vars) + rest
+            else:
+                return value.vars
+        # TODO: support unpacking other types
+        return tuple(ir.types.Any for _ in stmt.names)
+
+
+def unpackable(state: lowering.LoweringState, node: ast.expr, value: ir.SSAValue):
+    if isinstance(node, ast.Name):
+        state.current_frame.defs[node.id] = value
+        value.name = node.id
+        return
+    elif not isinstance(node, ast.Tuple):
+        raise DialectLoweringError(f"unsupported unpack node {node}")
+
+    names: list[str | None] = []
+    continue_unpack: list[int] = []
+    for idx, item in enumerate(node.elts):
+        if isinstance(item, ast.Name):
+            names.append(item.id)
+        else:
+            names.append(None)
+            continue_unpack.append(idx)
+    stmt = state.append_stmt(Unpack(value, tuple(names)))
+    for name, result in zip(names, stmt.results):
+        if name is not None:
+            state.current_frame.defs[name] = result
+
+    for idx in continue_unpack:
+        unpackable(state, node.elts[idx], stmt.results[idx])

src/kirin/dialects/scf/__init__.py

@@ -0,0 +1,3 @@
+from . import interp as interp, lowering as lowering
+from .stmts import For as For, Yield as Yield, IfElse as IfElse
+from ._dialect import dialect as dialect

src/kirin/dialects/scf/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("scf")

src/kirin/dialects/scf/interp.py

@@ -0,0 +1,32 @@
+from kirin import interp
+
+from .stmts import For, Yield, IfElse
+from ._dialect import dialect
+
+
+@dialect.register
+class Concrete(interp.MethodTable):
+
+    @interp.impl(Yield)
+    def yield_stmt(self, interp_: interp.Interpreter, frame: interp.Frame, stmt: Yield):
+        return interp.ReturnValue(*frame.get_values(stmt.values))
+
+    @interp.impl(IfElse)
+    def if_else(self, interp: interp.Interpreter, frame: interp.Frame, stmt: IfElse):
+        cond = frame.get(stmt.cond)
+        if cond:
+            body = stmt.then_body
+        else:
+            body = stmt.else_body
+        return interp.run_ssacfg_region(frame, body)
+
+    @interp.impl(For)
+    def for_loop(self, interpreter: interp.Interpreter, frame: interp.Frame, stmt: For):
+        iterable = frame.get(stmt.iterable)
+        loop_vars = frame.get_values(stmt.initializers)
+        block_args = stmt.body.blocks[0].args
+        for value in iterable:
+            frame.set_values(block_args, (value,) + loop_vars)
+            loop_vars = interpreter.run_ssacfg_region(frame, stmt.body)
+
+        return loop_vars

src/kirin/dialects/scf/lowering.py

@@ -0,0 +1,97 @@
+import ast
+
+from kirin import ir, lowering
+from kirin.exceptions import DialectLoweringError
+from kirin.dialects.py.unpack import unpackable
+
+from .stmts import For, Yield, IfElse
+from ._dialect import dialect
+
+
+@dialect.register
+class Lowering(lowering.FromPythonAST):
+
+    def lower_If(self, state: lowering.LoweringState, node: ast.If) -> lowering.Result:
+        cond = state.visit(node.test).expect_one()
+        frame = state.current_frame
+        body_frame = lowering.Frame.from_stmts(node.body, state)
+        state.push_frame(body_frame)
+        state.exhaust(body_frame)
+        state.pop_frame()
+
+        else_frame = lowering.Frame.from_stmts(node.orelse, state)
+        state.push_frame(else_frame)
+        state.exhaust(else_frame)
+        state.pop_frame()
+
+        yield_names: list[str] = []
+        body_yields: list[ir.SSAValue] = []
+        else_yields: list[ir.SSAValue] = []
+        if node.orelse:
+            for name in body_frame.defs.keys():
+                if name in else_frame.defs:
+                    yield_names.append(name)
+                    body_yields.append(body_frame.get_scope(name))
+                    else_yields.append(else_frame.get_scope(name))
+        else:
+            for name in body_frame.defs.keys():
+                if name in frame.defs:
+                    yield_names.append(name)
+                    body_yields.append(body_frame.get_scope(name))
+                    value = frame.get(name)
+                    if value is None:
+                        raise DialectLoweringError(f"expected value for {name}")
+                    else_yields.append(value)
+
+        body_frame.append_stmt(Yield(*body_yields))
+        else_frame.append_stmt(Yield(*else_yields))
+        stmt = IfElse(
+            cond,
+            then_body=body_frame.current_region,
+            else_body=else_frame.current_region,
+        )
+        for result, name, body, else_ in zip(
+            stmt.results, yield_names, body_yields, else_yields
+        ):
+            result.name = name
+            result.type = body.type.join(else_.type)
+            frame.defs[name] = result
+        state.append_stmt(stmt)
+        return lowering.Result()
+
+    def lower_For(
+        self, state: lowering.LoweringState, node: ast.For
+    ) -> lowering.Result:
+        iter_ = state.visit(node.iter).expect_one()
+
+        yields: list[str] = []
+
+        def new_block_arg_if_inside_loop(frame: lowering.Frame, capture: ir.SSAValue):
+            if not capture.name:
+                raise DialectLoweringError("unexpected loop variable captured")
+            yields.append(capture.name)
+            return frame.current_block.args.append_from(capture.type, capture.name)
+
+        body_frame = state.push_frame(
+            lowering.Frame.from_stmts(
+                node.body, state, capture_callback=new_block_arg_if_inside_loop
+            )
+        )
+        loop_var = body_frame.current_block.args.append_from(ir.types.Any)
+        unpackable(state, node.target, loop_var)
+        state.exhaust(body_frame)
+        # NOTE: this frame won't have phi nodes
+        body_frame.append_stmt(Yield(*[body_frame.defs[name] for name in yields]))  # type: ignore
+        state.pop_frame()
+
+        initializers: list[ir.SSAValue] = []
+        for name in yields:
+            value = state.current_frame.get(name)
+            if value is None:
+                raise DialectLoweringError(f"expected value for {name}")
+            initializers.append(value)
+        stmt = For(iter_, body_frame.current_region, *initializers)
+        for name, result in zip(yields, stmt.results):
+            state.current_frame.defs[name] = result
+        state.append_stmt(stmt)
+        return lowering.Result()