support binding for contextmanager (#346)

this PR supports binding of `with <stmt>` etc. by supporting
`@contextmanager` etc. in the `with` lowering transform.

removed `br` inside with lowering transform by default. One should
manually insert a terminator after the lowering or implement a custom
lowering as needed.

Author: Roger-luo

src/kirin/lowering/python/lowering.py

@@ -167,7 +167,7 @@ def visit_Call(self, state: State[ast.AST], node: ast.Call) -> Result:
                 node.func, state.lineno_offset, state.col_offset
             )
 
-        global_callee_result = self.lower_global_no_raise(state, node.func)
+        global_callee_result = state.get_global(node.func, no_raise=True)
         if global_callee_result is None:
             return self.visit_Call_local(state, node)
 
@@ -253,6 +253,9 @@ def visit_With(self, state: State[ast.AST], node: ast.With) -> Result:
             raise BuildError("expected context expression to be a call")
 
         global_callee = state.get_global(item.context_expr.func).data
+        if isinstance(global_callee, Binding):
+            global_callee = global_callee.parent
+
         if not issubclass(global_callee, ir.Statement):
             raise BuildError(
                 f"expected context expression to be a statement, got {global_callee}"

src/kirin/lowering/python/traits.py

@@ -242,9 +242,6 @@ def lower(
                     raise BuildError(
                         f"Expected exactly one block in region {region_name}"
                     )
-                body_frame.curr_region.blocks[0].stmts.append(
-                    cf.Branch(arguments=(), successor=body_frame.next_block)
-                )
 
                 if len(body_frame.curr_region.blocks) != 1:
                     raise BuildError(

test/lowering/test_with_binding.py

@@ -0,0 +1,53 @@
+from typing import Any, Generator
+from contextlib import contextmanager
+
+from kirin import ir, lowering
+from kirin.decl import info, statement
+from kirin.prelude import structural_no_opt
+from kirin.dialects import ilist
+
+dialect = ir.Dialect("with_binding")
+
+
+@statement(dialect=dialect)
+class ContextStatatement(ir.Statement):
+    traits = frozenset({lowering.FromPythonWithSingleItem()})
+    body: ir.Region = info.region(multi=False)
+
+
+@ir.dialect_group(structural_no_opt.add(dialect))
+def dummy(self):
+
+    def run_pass(mt):
+
+        return mt
+
+    return run_pass
+
+
+@lowering.wraps(ContextStatatement)
+@contextmanager
+def context_statement() -> Generator[Any, None, None]: ...
+
+
+@dummy
+def with_binding():
+    x = 1
+
+    def fn(x):
+        return x**2
+
+    with context_statement():
+        with context_statement():
+            x = ilist.map(fn, ilist.range(10))
+
+    return x
+
+
+def test_with_binding():
+    stmt = with_binding.callable_region.blocks[0].stmts.at(-2)
+    assert isinstance(stmt, ContextStatatement)
+    assert len(stmt.body.blocks) == 1
+    stmt = stmt.body.blocks[0].stmts.at(0)
+    assert isinstance(stmt, ContextStatatement)
+    assert len(stmt.body.blocks[0].stmts) == 5