refine scf.for pretty printing (#183)

fix #177

Author: Roger-luo

src/kirin/dialects/scf/stmts.py

@@ -148,11 +148,27 @@ def verify(self) -> None:
     def print_impl(self, printer: Printer) -> None:
         printer.print_name(self)
         printer.plain_print(" ")
+        block = self.body.blocks[0]
+        printer.print(block.args[0])
+        printer.plain_print(" in ", style="keyword")
         printer.print(self.iterable)
-        printer.plain_print(" init ")
-        printer.print_seq(self.initializers, delim=", ")
-        printer.plain_print(" ")
-        printer.print(self.body)
+        with printer.indent():
+            printer.print_newline()
+            printer.plain_print("iter_args(")
+            for idx, (arg, val) in enumerate(zip(block.args[1:], self.initializers)):
+                printer.print(arg)
+                printer.plain_print(" = ")
+                printer.print(val)
+                if idx < len(self.initializers) - 1:
+                    printer.plain_print(", ")
+            printer.plain_print(") {")
+
+            with printer.align(printer.result_width(block.stmts)):
+                for stmt in block.stmts:
+                    printer.print_newline()
+                    printer.print_stmt(stmt)
+        printer.print_newline()
+        printer.plain_print("}")
 
 
 @statement(dialect=dialect)

src/kirin/ir/nodes/block.py

@@ -439,26 +439,7 @@ def print_impl(self, printer: Printer) -> None:
         with printer.indent(increase=2, mark=False):
             for stmt in self.stmts:
                 printer.print_newline()
-                if stmt._results:
-                    result_str = printer.result_str(stmt._results)
-                    printer.plain_print(
-                        result_str.rjust(printer.state.result_width), " = "
-                    )
-                elif printer.state.result_width:
-                    printer.plain_print(" " * printer.state.result_width, "   ")
-                with printer.indent(printer.state.result_width + 3, mark=True):
-                    printer.print(stmt)
-                    if printer.analysis and any(
-                        result in printer.analysis for result in stmt._results
-                    ):
-                        with printer.rich(style="warning"):
-                            printer.plain_print(" # ---> ")
-                            printer.plain_print(
-                                ", ".join(
-                                    repr(printer.analysis[result])
-                                    for result in stmt._results
-                                )
-                            )
+                printer.print_stmt(stmt)
 
     def typecheck(self) -> None:
         """Checking the types of the Statments in the Block."""

src/kirin/ir/nodes/region.py

@@ -282,6 +282,15 @@ def walk(
         for block in reversed(self.blocks) if reverse else self.blocks:
             yield from block.walk(reverse=reverse, region_first=region_first)
 
+    def stmts(self) -> Iterator[Statement]:
+        """Iterate over all the Statements in the Region. This does not walk into nested Regions.
+
+        Yields:
+            Iterator[Statement]: An iterator that yield Statements of Blocks in the Region.
+        """
+        for block in self.blocks:
+            yield from block.stmts
+
     def print_impl(self, printer: Printer) -> None:
         # populate block ids
         for block in self.blocks:
@@ -293,12 +302,7 @@ def print_impl(self, printer: Printer) -> None:
             printer.plain_print("}")
             return
 
-        result_width = 0
-        for bb in self.blocks:
-            for stmt in bb.stmts:
-                result_width = max(result_width, len(printer.result_str(stmt._results)))
-
-        with printer.align(result_width):
+        with printer.align(printer.result_width(self.stmts())):
             with printer.indent(increase=2, mark=True):
                 printer.print_newline()
                 for idx, bb in enumerate(self.blocks):

src/kirin/print/printer.py

@@ -163,6 +163,25 @@ def print_name(
             self.plain_print(".")
         self.plain_print(node.name)
 
+    def print_stmt(self, node: "ir.Statement"):
+        if node._results:
+            result_str = self.result_str(node._results)
+            self.plain_print(result_str.rjust(self.state.result_width), " = ")
+        elif self.state.result_width:
+            self.plain_print(" " * self.state.result_width, "   ")
+        with self.indent(self.state.result_width + 3, mark=True):
+            self.print(node)
+            if self.analysis and any(
+                result in self.analysis for result in node._results
+            ):
+                with self.rich(style="warning"):
+                    self.plain_print(" # ---> ")
+                    self.plain_print(
+                        ", ".join(
+                            repr(self.analysis[result]) for result in node._results
+                        )
+                    )
+
     def print_dialect_path(
         self, node: Union["ir.Attribute", "ir.Statement"], prefix: str = ""
     ) -> None:
@@ -286,6 +305,20 @@ def print_mapping(
             self.plain_print(f"{key}=")
             emit(value)
 
+    def result_width(self, stmts: Iterable["ir.Statement"]) -> int:
+        """return the maximum width of the result column for a sequence of statements.
+
+        Args:
+            stmts(Iterable[ir.Statement]): sequence of statements
+
+        Returns:
+            int: maximum width of the result column
+        """
+        result_width = 0
+        for stmt in stmts:
+            result_width = max(result_width, len(self.result_str(stmt._results)))
+        return result_width
+
     @contextmanager
     def align(self, width: int):
         """align the result column width, and restore it after the context.