jaclang 0.8.1__py3-none-any.whl → 0.8.2__py3-none-any.whl

jaclang/compiler/passes/tool/doc_ir_gen_pass.py

@@ -78,26 +78,6 @@ class DocIRGenPass(UniPass):
 
         return self.concat(result)
 
-    def _strip_trailing_ws(self, part: doc.DocType) -> doc.DocType:
-        """Recursively strip trailing whitespace from a Doc node."""
-        if isinstance(part, doc.Concat) and part.parts:
-            while (
-                part.parts
-                and isinstance(part.parts[-1], doc.Text)
-                and getattr(part.parts[-1], "text", "") == " "
-            ):
-                part.parts.pop()
-            if part.parts:
-                part.parts[-1] = self._strip_trailing_ws(part.parts[-1])
-        elif isinstance(part, (doc.Group, doc.Indent, doc.Align)):
-            part.contents = self._strip_trailing_ws(part.contents)
-        return part
-
-    def finalize(self, parts: List[doc.DocType], group: bool = True) -> doc.DocType:
-        """Concat parts and remove trailing whitespace before grouping."""
-        result = self._strip_trailing_ws(self.concat(parts))
-        return self.group(result) if group else result
-
     def is_one_line(self, node: uni.UniNode) -> bool:
         """Check if the node is a one line node."""
         kid = [i for i in node.kid if not isinstance(i, uni.CommentToken)]
@@ -139,9 +119,15 @@ class DocIRGenPass(UniPass):
         """Exit import node."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            if isinstance(i, uni.Token) and i.name == Tok.SEMI:
+            if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                parts.pop()
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            elif isinstance(i, uni.Token) and i.name == Tok.SEMI:
                 parts.pop()
                 parts.append(i.gen.doc_ir)
+            elif isinstance(i, uni.Token) and i.name == Tok.RBRACE:
+                parts.append(i.gen.doc_ir)
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
@@ -188,6 +174,13 @@ class DocIRGenPass(UniPass):
                 parts.append(self.space())
             elif isinstance(i, uni.Token) and i.name == Tok.LBRACE:
                 parts.append(i.gen.doc_ir)
+            elif isinstance(i, uni.Token) and i.name == Tok.LPAREN:
+                parts.pop()
+                parts.append(i.gen.doc_ir)
+            elif isinstance(i, uni.Token) and i.name == Tok.RPAREN:
+                parts.pop()
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
             elif isinstance(node.body, Sequence) and i in node.body:
                 if not in_body:
                     body_parts.append(self.hard_line())
@@ -214,7 +207,7 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
 
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_ability(self, node: uni.Ability) -> None:
         """Generate DocIR for abilities."""
@@ -250,19 +243,21 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_func_signature(self, node: uni.FuncSignature) -> None:
         """Generate DocIR for function signatures."""
         parts: list[doc.DocType] = []
         indent_parts: list[doc.DocType] = []
         in_params = False
+        has_parens = False
         for i in node.kid:
             if isinstance(i, uni.Token) and i.name == Tok.LPAREN and node.params:
                 in_params = True
                 parts.append(i.gen.doc_ir)
             elif isinstance(i, uni.Token) and i.name == Tok.RPAREN and node.params:
                 in_params = False
+                has_parens = True
                 parts.append(
                     self.indent(self.concat([self.tight_line(), *indent_parts]))
                 )
@@ -280,9 +275,16 @@ class DocIRGenPass(UniPass):
                 else:
                     indent_parts.append(i.gen.doc_ir)
             else:
+                if (
+                    isinstance(i, uni.Token)
+                    and i.name == Tok.RETURN_HINT
+                    and not has_parens
+                ):
+                    parts.append(self.space())
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_param_var(self, node: uni.ParamVar) -> None:
         """Generate DocIR for parameter variables."""
@@ -294,27 +296,52 @@ class DocIRGenPass(UniPass):
                 parts.append(self.space())
             else:
                 parts.append(i.gen.doc_ir)
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_type_ref(self, node: uni.TypeRef) -> None:
         """Generate DocIR for type references."""
         parts: list[doc.DocType] = []
         for i in node.kid:
             parts.append(i.gen.doc_ir)
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_assignment(self, node: uni.Assignment) -> None:
         """Generate DocIR for assignments."""
-        parts: list[doc.DocType] = []
-        for i in node.kid:
-            if i == node.type_tag or (isinstance(i, uni.Token) and i.name == Tok.SEMI):
-                parts.pop()
-                parts.append(i.gen.doc_ir)
-                parts.append(self.space())
+        lhs_parts: list[doc.DocType] = []
+        rhs_parts: list[doc.DocType] = []
+        eq_tok: Optional[doc.DocType] = None
+        seen_eq = False
+
+        for i in node.kid:
+            if isinstance(i, uni.Token) and i.name == Tok.KW_LET:
+                lhs_parts.append(i.gen.doc_ir)
+                lhs_parts.append(self.space())
+            elif isinstance(i, uni.Token) and i.name == Tok.EQ and not seen_eq:
+                eq_tok = i.gen.doc_ir
+                seen_eq = True
+            elif seen_eq:
+                rhs_parts.append(i.gen.doc_ir)
             else:
-                parts.append(i.gen.doc_ir)
-                parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+                if i == node.aug_op:
+                    lhs_parts.append(self.space())
+                lhs_parts.append(i.gen.doc_ir)
+                if i == node.aug_op:
+                    lhs_parts.append(self.space())
+
+        if eq_tok is not None:
+            rhs_concat = self.concat(rhs_parts)
+            node.gen.doc_ir = self.group(
+                self.concat(
+                    [
+                        *lhs_parts,
+                        self.space(),
+                        eq_tok,
+                        self.indent(self.concat([self.line(), rhs_concat])),
+                    ]
+                )
+            )
+        else:
+            node.gen.doc_ir = self.group(self.concat(lhs_parts + rhs_parts))
 
     def exit_if_stmt(self, node: uni.IfStmt) -> None:
         """Generate DocIR for if statements."""
@@ -343,7 +370,8 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_else_if(self, node: uni.ElseIf) -> None:
         """Generate DocIR for else if statements."""
@@ -372,7 +400,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_else_stmt(self, node: uni.ElseStmt) -> None:
         """Generate DocIR for else statements."""
@@ -401,7 +429,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_binary_expr(self, node: uni.BinaryExpr) -> None:
         """Generate DocIR for binary expressions."""
@@ -409,7 +437,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_expr_stmt(self, node: uni.ExprStmt) -> None:
         """Generate DocIR for expression statements."""
@@ -424,7 +453,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_return_stmt(self, node: uni.ReturnStmt) -> None:
         """Generate DocIR for return statements."""
@@ -436,7 +465,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_func_call(self, node: uni.FuncCall) -> None:
         """Generate DocIR for function calls."""
@@ -475,8 +504,24 @@ class DocIRGenPass(UniPass):
         """Generate DocIR for list values."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
-        node.gen.doc_ir = self.group(self.concat(parts))
+            if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            else:
+                parts.append(i.gen.doc_ir)
+        not_broke = self.concat(parts)
+        parts = []
+        for i in node.kid:
+            if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.hard_line())
+            elif isinstance(i, uni.Token) and i.name == Tok.LSQUARE:
+                parts.append(self.hard_line())
+                parts.append(i.gen.doc_ir)
+            else:
+                parts.append(i.gen.doc_ir)
+        broke = self.concat(parts)
+        node.gen.doc_ir = self.group(self.if_break(broke, not_broke))
 
     def exit_dict_val(self, node: uni.DictVal) -> None:
         """Generate DocIR for dictionary values."""
@@ -488,7 +533,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_k_v_pair(self, node: uni.KVPair) -> None:
         """Generate DocIR for key-value pairs."""
@@ -496,7 +541,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_has_var(self, node: uni.HasVar) -> None:
         """Generate DocIR for has variable declarations."""
@@ -527,7 +573,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_while_stmt(self, node: uni.WhileStmt) -> None:
         """Generate DocIR for while statements."""
@@ -535,7 +581,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_in_for_stmt(self, node: uni.InForStmt) -> None:
         """Generate DocIR for for-in statements."""
@@ -543,7 +589,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_iter_for_stmt(self, node: uni.IterForStmt) -> None:
         """Generate DocIR for iterative for statements."""
@@ -551,7 +597,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_try_stmt(self, node: uni.TryStmt) -> None:
         """Generate DocIR for try statements."""
@@ -559,7 +605,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_except(self, node: uni.Except) -> None:
         """Generate DocIR for except clauses."""
@@ -567,7 +613,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_finally_stmt(self, node: uni.FinallyStmt) -> None:
         """Generate DocIR for finally statements."""
@@ -575,14 +621,30 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_tuple_val(self, node: uni.TupleVal) -> None:
         """Generate DocIR for tuple values."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
-        node.gen.doc_ir = self.group(self.concat(parts))
+            if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            else:
+                parts.append(i.gen.doc_ir)
+        not_broke = self.concat(parts)
+        parts = []
+        for i in node.kid:
+            if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.hard_line())
+            elif isinstance(i, uni.Token) and i.name == Tok.LPAREN:
+                parts.append(self.hard_line())
+                parts.append(i.gen.doc_ir)
+            else:
+                parts.append(i.gen.doc_ir)
+        broke = self.concat(parts)
+        node.gen.doc_ir = self.group(self.if_break(broke, not_broke))
 
     def exit_multi_string(self, node: uni.MultiString) -> None:
         """Generate DocIR for multiline strings."""
@@ -605,7 +667,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_list_compr(self, node: uni.ListCompr) -> None:
         """Generate DocIR for list comprehensions."""
@@ -613,7 +676,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_inner_compr(self, node: uni.InnerCompr) -> None:
         """Generate DocIR for inner comprehension clauses."""
@@ -621,7 +685,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_f_string(self, node: uni.FString) -> None:
         """Generate DocIR for formatted strings."""
@@ -644,7 +709,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-
+        parts.pop()
         node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_bool_expr(self, node: uni.BoolExpr) -> None:
@@ -657,7 +722,6 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.line())  # Potential break
-
         node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_unary_expr(self, node: uni.UnaryExpr) -> None:
@@ -723,7 +787,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_gen_compr(self, node: uni.GenCompr) -> None:
         """Generate DocIR for generator comprehensions."""
@@ -731,7 +796,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_set_compr(self, node: uni.SetCompr) -> None:
         """Generate DocIR for set comprehensions."""
@@ -739,7 +805,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_dict_compr(self, node: uni.DictCompr) -> None:
         """Generate DocIR for dictionary comprehensions."""
@@ -750,14 +817,15 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_k_w_pair(self, node: uni.KWPair) -> None:
         """Generate DocIR for keyword arguments."""
         parts: list[doc.DocType] = []
         for i in node.kid:
             parts.append(i.gen.doc_ir)
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_await_expr(self, node: uni.AwaitExpr) -> None:
         """Generate DocIR for await expressions."""
@@ -765,7 +833,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_yield_expr(self, node: uni.YieldExpr) -> None:
         """Generate DocIR for yield expressions."""
@@ -773,7 +842,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_ctrl_stmt(self, node: uni.CtrlStmt) -> None:
         """Generate DocIR for control statements (break, continue, skip)."""
@@ -788,7 +858,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_disengage_stmt(self, node: uni.DisengageStmt) -> None:
         """Generate DocIR for disengage statements."""
@@ -803,7 +874,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_assert_stmt(self, node: uni.AssertStmt) -> None:
         """Generate DocIR for assert statements."""
@@ -811,7 +883,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_raise_stmt(self, node: uni.RaiseStmt) -> None:
         """Generate DocIR for raise statements."""
@@ -819,7 +892,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_global_vars(self, node: uni.GlobalVars) -> None:
         """Generate DocIR for global variables."""
@@ -835,7 +909,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_module_code(self, node: uni.ModuleCode) -> None:
         """Generate DocIR for module code."""
@@ -870,7 +944,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_global_stmt(self, node: uni.GlobalStmt) -> None:
         """Generate DocIR for global statements."""
@@ -885,7 +959,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_visit_stmt(self, node: uni.VisitStmt) -> None:
         """Generate DocIR for visit statements."""
@@ -898,7 +972,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_ignore_stmt(self, node: uni.IgnoreStmt) -> None:
         """Generate DocIR for ignore statements."""
@@ -906,7 +980,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_connect_op(self, node: uni.ConnectOp) -> None:
         """Generate DocIR for connect operator."""
@@ -914,7 +988,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_disconnect_op(self, node: uni.DisconnectOp) -> None:
         """Generate DocIR for disconnect operator."""
@@ -922,7 +997,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_compare_expr(self, node: uni.CompareExpr) -> None:
         """Generate DocIR for comparison expressions."""
@@ -930,7 +1006,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_atom_unit(self, node: uni.AtomUnit) -> None:
         """Generate DocIR for atom units (parenthesized expressions)."""
@@ -952,7 +1029,8 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
             prev_item = i
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_expr_as_item(self, node: uni.ExprAsItem) -> None:
         """Generate DocIR for expression as item nodes."""
@@ -960,7 +1038,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_filter_compr(self, node: uni.FilterCompr) -> None:
         """Generate DocIR for filter comprehensions."""
@@ -968,7 +1046,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_assign_compr(self, node: uni.AssignCompr) -> None:
         """Generate DocIR for assignment comprehensions."""
@@ -976,7 +1054,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_py_inline_code(self, node: uni.PyInlineCode) -> None:
         """Generate DocIR for Python inline code blocks."""
@@ -993,7 +1071,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_test(self, node: uni.Test) -> None:
         """Generate DocIR for test nodes."""
@@ -1009,7 +1087,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_check_stmt(self, node: uni.CheckStmt) -> None:
         """Generate DocIR for check statements."""
@@ -1019,7 +1097,7 @@ class DocIRGenPass(UniPass):
                 parts.pop()
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_stmt(self, node: uni.MatchStmt) -> None:
         """Generate DocIR for match statements."""
@@ -1039,7 +1117,7 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_case(self, node: uni.MatchCase) -> None:
         """Generate DocIR for match cases."""
@@ -1056,7 +1134,7 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         parts.append(self.indent(self.concat([self.hard_line()] + indent_parts)))
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_value(self, node: uni.MatchValue) -> None:
         """Generate DocIR for match value patterns."""
@@ -1064,7 +1142,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_singleton(self, node: uni.MatchSingleton) -> None:
         """Generate DocIR for match singleton patterns."""
@@ -1072,7 +1151,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_sequence(self, node: uni.MatchSequence) -> None:
         """Generate DocIR for match sequence patterns."""
@@ -1080,7 +1160,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_mapping(self, node: uni.MatchMapping) -> None:
         """Generate DocIR for match mapping patterns."""
@@ -1088,7 +1169,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_or(self, node: uni.MatchOr) -> None:
         """Generate DocIR for match OR patterns."""
@@ -1096,7 +1178,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_as(self, node: uni.MatchAs) -> None:
         """Generate DocIR for match AS patterns."""
@@ -1104,7 +1187,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_wild(self, node: uni.MatchWild) -> None:
         """Generate DocIR for match wildcard patterns."""
@@ -1112,7 +1196,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_star(self, node: uni.MatchStar) -> None:
         """Generate DocIR for match star patterns (e.g., *args, **kwargs)."""
@@ -1120,7 +1205,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_k_v_pair(self, node: uni.MatchKVPair) -> None:
         """Generate DocIR for match key-value pairs."""
@@ -1131,7 +1217,8 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_match_arch(self, node: uni.MatchArch) -> None:
         """Generate DocIR for match architecture patterns."""
@@ -1142,7 +1229,7 @@ class DocIRGenPass(UniPass):
                 parts.append(self.space())
             else:
                 parts.append(i.gen.doc_ir)
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_enum(self, node: uni.Enum) -> None:
         """Generate DocIR for enum declarations."""
@@ -1160,30 +1247,74 @@ class DocIRGenPass(UniPass):
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_sub_tag(self, node: uni.SubTag) -> None:
         """Generate DocIR for sub-tag nodes."""
-        parts: list[doc.DocType] = []
+        before_colon: list[doc.DocType] = []
+        after_colon: list[doc.DocType] = []
+        seen_colon = False
+
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
-            parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+            if isinstance(i, uni.Token) and i.name == Tok.COLON and not seen_colon:
+                colon_tok = i.gen.doc_ir
+                seen_colon = True
+            elif seen_colon:
+                after_colon.append(i.gen.doc_ir)
+            else:
+                before_colon.append(i.gen.doc_ir)
+
+        if seen_colon:
+            flat = self.concat([*before_colon, colon_tok, self.space(), *after_colon])
+            broke = self.concat(
+                [
+                    *before_colon,
+                    colon_tok,
+                    self.indent(self.concat([self.line(), *after_colon])),
+                ]
+            )
+            node.gen.doc_ir = self.group(self.if_break(broke, flat))
+        else:
+            node.gen.doc_ir = self.concat(before_colon + after_colon)
 
     def exit_impl_def(self, node: uni.ImplDef) -> None:
         """Generate DocIR for implementation definitions."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        in_body = False
         for i in node.kid:
             if i == node.doc or (node.decorators and i in node.decorators):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.hard_line())
-            elif i == node.target:
+            elif i in node.target:
+                parts.append(i.gen.doc_ir)
+            elif (
+                in_body
+                or isinstance(node.body, Sequence)
+                and node.body
+                and i == node.body[0]
+            ):
+                if not in_body:
+                    parts.pop()
+                    body_parts.append(self.hard_line())
+                if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                    body_parts.pop()
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                in_body = True
+                if in_body and isinstance(node.body, Sequence) and i == node.body[-1]:
+                    in_body = False
+                    body_parts.pop()
+                    parts.append(self.indent(self.concat(body_parts)))
+                    parts.append(self.hard_line())
+            elif isinstance(i, uni.Token) and i.name == Tok.SEMI:
+                parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_event_signature(self, node: uni.EventSignature) -> None:
         """Generate DocIR for event signatures."""
@@ -1191,7 +1322,8 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        parts.pop()
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_typed_ctx_block(self, node: uni.TypedCtxBlock) -> None:
         """Generate DocIR for typed context blocks."""