use a worklist for emit (#128)

Author: Roger-luo

src/kirin/dialects/cf/emit.py

@@ -1,7 +1,7 @@
 from typing import IO, TypeVar
 
 from kirin import emit
-from kirin.interp import MethodTable, impl
+from kirin.interp import Successor, MethodTable, impl
 from kirin.emit.julia import EmitJulia
 
 from .stmts import Assert, Branch, ConditionalBranch
@@ -27,6 +27,9 @@ def emit_branch(
         self, interp: EmitJulia[IO_t], frame: emit.EmitStrFrame, stmt: Branch
     ):
         interp.writeln(frame, f"@goto {interp.block_id[stmt.successor]};")
+        frame.worklist.append(
+            Successor(stmt.successor, frame.get_values(stmt.arguments))
+        )
         return ()
 
     @impl(ConditionalBranch)
@@ -51,4 +54,11 @@ def emit_cbr(
         interp.writeln(frame, f"@goto {interp.block_id[stmt.else_successor]};")
         frame.indent -= 1
         interp.writeln(frame, "end")
+
+        frame.worklist.append(
+            Successor(stmt.then_successor, frame.get_values(stmt.then_arguments))
+        )
+        frame.worklist.append(
+            Successor(stmt.else_successor, frame.get_values(stmt.else_arguments))
+        )
         return ()

src/kirin/emit/abc.py

@@ -2,13 +2,15 @@
 from dataclasses import field, dataclass
 
 from kirin import ir, interp
+from kirin.worklist import WorkList
 from kirin.exceptions import FuelExhaustedError
 
 ValueType = TypeVar("ValueType")
 
 
 @dataclass
 class EmitFrame(interp.Frame[ValueType]):
+    worklist: WorkList[interp.Successor] = field(default_factory=WorkList)
     block_ref: dict[ir.Block, ValueType] = field(default_factory=dict)
 
 
@@ -51,12 +53,14 @@ def run_ssacfg_region(
         self, frame: FrameType, region: ir.Region
     ) -> ValueType | interp.Err[ValueType]:
         result = self.bottom
-        for block in region.blocks:
-            block_header = self.emit_block(frame, block)
+        frame.worklist.append(
+            interp.Successor(region.blocks[0], frame.get_values(region.blocks[0].args))
+        )
+        while (succ := frame.worklist.pop()) is not None:
+            block_header = self.emit_block(frame, succ.block)
             if isinstance(block_header, interp.Err):
                 return block_header
-            frame.block_ref[block] = block_header
-
+            frame.block_ref[succ.block] = block_header
         return result
 
     def emit_attribute(self, attr: ir.Attribute) -> ValueType:

src/kirin/emit/abc.pyi

@@ -2,11 +2,13 @@ from typing import TypeVar, Iterable
 from dataclasses import field, dataclass
 
 from kirin import ir, interp
+from kirin.worklist import WorkList
 
 ValueType = TypeVar("ValueType")
 
 @dataclass
 class EmitFrame(interp.Frame[ValueType]):
+    worklist: WorkList[interp.Successor] = field(default_factory=WorkList)
     block_ref: dict[ir.Block, ValueType] = field(default_factory=dict)
 
 FrameType = TypeVar("FrameType", bound=EmitFrame)

test/emit/test_julia.py

@@ -26,10 +26,7 @@ def foo():
     assert "function foo()" in generated
     assert "@label block_1;" in generated
     assert "return x" in generated
-    assert "@label block_2;" in generated
-    assert "return nothing" in generated
     assert "return foo" in generated
-    assert "return nothing" in generated
 
 
 def test_py_stmts():