ClosureFieldLowering rewrite: Lower captured fields into constants (#538)

A ClosureFieldLowering pass that lowers closure-captured fields into
`py.Constant`.
e.g. the higher-order function `outer` returns `inner` which closes over
parameter `y`.
```
@basic
def outer(y: int):
    def inner(x: int):
        return x * y + 1
    return outer

inner_ker = outer(y=10)
```
`inner_ker.print()` will go from:
```
func.lambda inner(%y : !py.int) -> !Any {
  ^0(%inner_self, %x):
  │ %y_1 = func.getfield(%inner_self, 0) : !py.int
  │   %0 = py.binop.mult(%x : !py.int, %y_1) : ~T
  │   %1 = py.constant.constant 1 : !py.int
  │   %2 = py.binop.add(%0, %1) : ~T
  │        func.return %2
} // func.lambda inner
```
to:
```
func.lambda inner(%y : !py.int) -> !Any {
  ^0(%inner_self, %x):
  │ %y_1 = py.constant.constant 10 : !py.int
  │   %0 = py.binop.mult(%x : !py.int, %y_1) : ~T
  │   %1 = py.constant.constant 1 : !py.int
  │   %2 = py.binop.add(%0, %1) : ~T
  │        func.return %2
} // func.lambda inner
```

Author: zhenrongliew

src/kirin/ir/nodes/block.py

@@ -391,6 +391,10 @@ def is_structurally_equal(  # pyright: ignore[reportIncompatibleMethodOverride]
         if context is None:
             context = {}
 
+        if self in context:
+            return context[self] is other
+        context[self] = other
+
         if len(self._args) != len(other._args) or len(self.stmts) != len(other.stmts):
             return False
 

src/kirin/passes/fold.py

@@ -10,6 +10,7 @@
     CFGCompactify,
     InlineGetItem,
     DeadCodeElimination,
+    ClosureFieldLowering,
 )
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
@@ -28,6 +29,7 @@ class Fold(Pass):
     - `InlineGetItem`
     - `Call2Invoke`
     - `DeadCodeElimination`
+    - `ClosureFieldLowering`
     """
 
     hint_const: HintConst = field(init=False)
@@ -46,6 +48,7 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
                         InlineGetItem(),
                         Call2Invoke(),
                         DeadCodeElimination(),
+                        ClosureFieldLowering(),
                     )
                 )
             )

src/kirin/rewrite/__init__.py

@@ -13,3 +13,4 @@
 from .wrap_const import WrapConst as WrapConst
 from .call2invoke import Call2Invoke as Call2Invoke
 from .type_assert import InlineTypeAssert as InlineTypeAssert
+from .closurefieldlowering import ClosureFieldLowering as ClosureFieldLowering

src/kirin/rewrite/closurefieldlowering.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.dialects import py, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+
+@dataclass
+class ClosureFieldLowering(RewriteRule):
+    """Lowers captured closure fields into py.Constants.
+    - Trigger on func.Invoke
+    - If the callee Method has non-empty .fields, lower its func.GetField to py.Constant
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, func.Invoke):
+            return RewriteResult(has_done_something=False)
+
+        method = node.callee
+        if not method.fields:
+            return RewriteResult(has_done_something=False)
+        # Replace func.GetField with py.Constant.
+        changed = self._lower_captured_fields(method)
+        if changed:
+            method.fields = ()
+        return RewriteResult(has_done_something=changed)
+
+    def _get_field_index(self, getfield_stmt: func.GetField) -> int | None:
+        fld = getfield_stmt.attributes.get("field")
+        if fld:
+            return getfield_stmt.field
+        else:
+            return None
+
+    def _lower_captured_fields(self, method: ir.Method) -> bool:
+        changed = False
+        fields = method.fields
+        if not fields:
+            return False
+
+        for region in method.code.regions:
+            for block in region.blocks:
+                for stmt in list(block.stmts):
+                    if not isinstance(stmt, func.GetField):
+                        continue
+                    idx = self._get_field_index(stmt)
+                    if idx is None:
+                        continue
+                    captured = fields[idx]
+                    # Skip Methods.
+                    if isinstance(captured, ir.Method):
+                        continue
+                    # Replace GetField with Constant.
+                    const_stmt = py.Constant(captured)
+                    const_stmt.insert_before(stmt)
+                    if stmt.results and const_stmt.results:
+                        stmt.results[0].replace_by(const_stmt.results[0])
+                    stmt.delete()
+                    changed = True
+        return changed

test/serialization/test_jsonserializer.py

@@ -9,7 +9,7 @@
 def foo(x: int, y: float, z: bool):
     c = [[(200.0, 200.0), (210.0, 200.0)]]
     if z:
-        c.append([(222.0, 333.0)])
+        c = [(222.0, 333.0)]
     else:
         return [1, 2, 3, 4]
     return c
@@ -47,6 +47,23 @@ def my_kernel2(y: int):
     return my_kernel1(y) * 10
 
 
+@basic
+def foo2(y: int):
+
+    def inner(x: int):
+        return x * y + 1
+
+    return inner
+
+
+inner_ker = foo2(y=10)
+
+
+@basic
+def main_lambda(z: int):
+    return inner_ker(z)
+
+
 @basic
 def slicing():
     in1 = ("a", "b", "c", "d", "e", "f", "g", "h")
@@ -94,6 +111,10 @@ def test_round_trip5():
     round_trip(slicing)
 
 
+def test_round_trip6():
+    round_trip(main_lambda)
+
+
 def test_deterministic():
     serializer = Serializer()
     s1 = serializer.encode(loop_ilist)