jaclang 0.8.7__py3-none-any.whl → 0.8.9__py3-none-any.whl

jaclang/compiler/passes/tool/doc_ir_gen_pass.py

@@ -135,20 +135,53 @@ class DocIRGenPass(UniPass):
     def exit_import(self, node: uni.Import) -> None:
         """Exit import node."""
         parts: list[doc.DocType] = []
+        mod_items: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
+        is_in_items: bool = False
         for i in node.kid:
             if isinstance(i, uni.Token) and i.name == Tok.COMMA:
-                parts.pop()
-                parts.append(i.gen.doc_ir)
-                parts.append(self.space())
+                if is_in_items:
+                    mod_items.pop()
+                    mod_items.append(i.gen.doc_ir)
+                    mod_items.append(self.line())
+                else:
+                    parts.pop()
+                    parts.append(i.gen.doc_ir)
+                    parts.append(self.line())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             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:
+                is_in_items = False
+                mod_items.pop()
+                parts.append(
+                    self.group(
+                        self.concat(
+                            [
+                                self.indent(self.concat([self.line(), *mod_items])),
+                                self.line(),
+                            ]
+                        )
+                    )
+                )
                 parts.append(i.gen.doc_ir)
-            else:
+            elif isinstance(i, uni.Token) and i.name == Tok.LBRACE:
+                is_in_items = True
                 parts.append(i.gen.doc_ir)
-                parts.append(self.space())
+            else:
+                if is_in_items:
+                    mod_items.append(i.gen.doc_ir)
+                    mod_items.append(self.space())
+                else:
+                    parts.append(i.gen.doc_ir)
+                    parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_module_item(self, node: uni.ModuleItem) -> None:
         """Generate DocIR for module items."""
@@ -178,6 +211,7 @@ class DocIRGenPass(UniPass):
         """Generate DocIR for archetypes."""
         parts: list[doc.DocType] = []
         body_parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         prev_item = None
         in_body = False
         for i in node.kid:
@@ -198,6 +232,10 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             elif isinstance(node.body, Sequence) and i in node.body:
                 if not in_body:
                     body_parts.append(self.hard_line())
@@ -220,16 +258,20 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif not in_body and isinstance(i, uni.Token) and i.name == Tok.DECOR_OP:
+                parts.append(i.gen.doc_ir)
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_ability(self, node: uni.Ability) -> None:
         """Generate DocIR for abilities."""
         parts: list[doc.DocType] = []
         body_parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         in_body = False
         for i in node.kid:
             if i == node.doc or (node.decorators and i in node.decorators):
@@ -257,10 +299,18 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
+            elif not in_body and isinstance(i, uni.Token) and i.name == Tok.DECOR_OP:
+                parts.append(i.gen.doc_ir)
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_func_signature(self, node: uni.FuncSignature) -> None:
         """Generate DocIR for function signatures."""
@@ -275,8 +325,17 @@ class DocIRGenPass(UniPass):
             elif isinstance(i, uni.Token) and i.name == Tok.RPAREN and node.params:
                 in_params = False
                 has_parens = True
+                if isinstance(indent_parts[-1], doc.Line):
+                    indent_parts.pop()
                 parts.append(
-                    self.indent(self.concat([self.tight_line(), *indent_parts]))
+                    self.indent(
+                        self.concat(
+                            [
+                                self.tight_line(),
+                                self.group(self.concat([*indent_parts])),
+                            ]
+                        )
+                    )
                 )
                 parts.append(self.tight_line())
                 parts.append(i.gen.doc_ir)
@@ -346,14 +405,14 @@ class DocIRGenPass(UniPass):
                     lhs_parts.append(self.space())
 
         if eq_tok is not None:
-            rhs_concat = self.concat(rhs_parts)
+            rhs_concat = self.group(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])),
+                        self.concat([self.space(), rhs_concat]),
                     ]
                 )
             )
@@ -375,6 +434,20 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif i == node.condition and isinstance(i, uni.BoolExpr):
+                cond_str = i.gen.doc_ir
+                flat = self.concat([cond_str, self.space()])
+                broken = self.group(
+                    self.concat(
+                        [
+                            self.text("("),
+                            self.indent(self.concat([self.line(), cond_str])),
+                            self.line(),
+                            self.text(")"),
+                        ]
+                    )
+                )
+                parts.append(self.group(self.if_break(broken, flat)))
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
@@ -397,6 +470,21 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif i == node.condition and isinstance(i, uni.BoolExpr):
+                cond_str = i.gen.doc_ir
+                flat = self.concat([cond_str, self.space()])
+                broken = self.group(
+                    self.concat(
+                        [
+                            self.text("("),
+                            self.indent(self.concat([self.line(), cond_str])),
+                            self.line(),
+                            self.text(")"),
+                            self.space(),
+                        ]
+                    )
+                )
+                parts.append(self.if_break(broken, flat))
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
@@ -471,8 +559,18 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
             elif isinstance(i, uni.Token) and i.name == Tok.RPAREN and node.params:
                 in_params = False
+
+                if isinstance(indent_parts[-1], doc.Line):
+                    indent_parts.pop()
                 parts.append(
-                    self.indent(self.concat([self.tight_line(), *indent_parts]))
+                    self.indent(
+                        self.concat(
+                            [
+                                self.tight_line(),
+                                self.group(self.concat([*indent_parts])),
+                            ]
+                        )
+                    )
                 )
                 parts.append(self.tight_line())
                 parts.append(i.gen.doc_ir)
@@ -510,12 +608,16 @@ class DocIRGenPass(UniPass):
             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)
+        broke = self.concat(
+            [
+                parts[0],
+                self.indent(self.concat([self.hard_line(), *parts[1:-1]])),
+                self.hard_line(),
+                parts[-1],
+            ]
+        )
         node.gen.doc_ir = self.group(self.if_break(broke, not_broke))
 
     def exit_dict_val(self, node: uni.DictVal) -> None:
@@ -699,12 +801,16 @@ class DocIRGenPass(UniPass):
             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)
+        broke = self.concat(
+            [
+                parts[0],
+                self.indent(self.concat([self.hard_line(), *parts[1:-1]])),
+                self.hard_line(),
+                parts[-1],
+            ]
+        )
         node.gen.doc_ir = self.group(self.if_break(broke, not_broke))
 
     def exit_multi_string(self, node: uni.MultiString) -> None:
@@ -748,17 +854,34 @@ class DocIRGenPass(UniPass):
         """Generate DocIR for list comprehensions."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
+            if isinstance(i, uni.InnerCompr):
+                parts.append(self.group(self.concat([self.tight_line(), i.gen.doc_ir])))
+            else:
+                parts.append(i.gen.doc_ir)
             parts.append(self.space())
         parts.pop()
-        node.gen.doc_ir = self.group(self.concat(parts))
+        node.gen.doc_ir = self.group(
+            self.concat(
+                [
+                    parts[0],
+                    self.indent(self.concat([self.tight_line(), *parts[2:-2]])),
+                    self.tight_line(),
+                    parts[-1],
+                ]
+            )
+        )
 
     def exit_inner_compr(self, node: uni.InnerCompr) -> None:
         """Generate DocIR for inner comprehension clauses."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
-            parts.append(self.space())
+            if isinstance(i, uni.Token) and i.name == Tok.KW_IF:
+                parts.append(self.hard_line())
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            else:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
         parts.pop()
         node.gen.doc_ir = self.group(self.concat(parts))
 
@@ -772,6 +895,7 @@ class DocIRGenPass(UniPass):
     def exit_if_else_expr(self, node: uni.IfElseExpr) -> None:
         """Generate DocIR for conditional expressions."""
         parts: list[doc.DocType] = []
+        need_parens = not isinstance(node.parent, uni.AtomUnit)
 
         for i in node.kid:
             if isinstance(i, uni.Expr):
@@ -782,21 +906,60 @@ class DocIRGenPass(UniPass):
                 parts.append(self.space())
             else:
                 parts.append(i.gen.doc_ir)
-                parts.append(self.space())
+                parts.append(self.line())
         parts.pop()
-        node.gen.doc_ir = self.group(self.concat(parts))
+
+        flat = self.group(self.concat(parts))
+        parens = self.group(
+            self.concat(
+                [
+                    self.text("("),
+                    self.indent(self.concat([self.tight_line(), flat])),
+                    self.tight_line(),
+                    self.text(")"),
+                ]
+            )
+        )
+        node.gen.doc_ir = flat
+
+        if need_parens:
+            if isinstance(node.parent, uni.Assignment):
+                node.gen.doc_ir = self.if_break(
+                    break_contents=parens, flat_contents=flat
+                )
+            else:
+                node.gen.doc_ir = parens
 
     def exit_bool_expr(self, node: uni.BoolExpr) -> None:
         """Generate DocIR for boolean expressions (and/or)."""
+        exprs: list[uni.UniNode] = []
         parts: list[doc.DocType] = []
-        for i in node.kid:
-            if isinstance(i, uni.Token):
-                parts.append(i.gen.doc_ir)
-                parts.append(self.space())
+
+        def __flatten_bool_expr(expr: uni.Expr) -> list[uni.UniNode]:
+            if isinstance(expr, uni.BoolExpr):
+                out: list[uni.UniNode] = []
+                for val in expr.values:
+                    out += __flatten_bool_expr(val)
+                    out.append(expr.op)
+                out.pop()
+                return out
             else:
-                parts.append(i.gen.doc_ir)
-                parts.append(self.line())  # Potential break
-        node.gen.doc_ir = self.group(self.concat(parts))
+                return [expr]
+
+        exprs = __flatten_bool_expr(node)
+        parts += [exprs[0].gen.doc_ir, self.line()]
+        for i in range(1, len(exprs), 2):
+            (
+                op,
+                expr,
+            ) = (
+                exprs[i + 1],
+                exprs[i],
+            )
+            parts += [expr.gen.doc_ir, self.space(), op.gen.doc_ir, self.line()]
+        parts.pop()
+        flat = self.concat(parts)
+        node.gen.doc_ir = self.group(flat)
 
     def exit_unary_expr(self, node: uni.UnaryExpr) -> None:
         """Generate DocIR for unary expressions."""
@@ -860,27 +1023,49 @@ class DocIRGenPass(UniPass):
         parts: list[doc.DocType] = []
         for i in node.kid:
             parts.append(i.gen.doc_ir)
-            parts.append(self.space())
-        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."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
+            if isinstance(i, uni.InnerCompr):
+                parts.append(self.group(self.concat([self.tight_line(), i.gen.doc_ir])))
+            else:
+                parts.append(i.gen.doc_ir)
             parts.append(self.space())
         parts.pop()
-        node.gen.doc_ir = self.group(self.concat(parts))
+        node.gen.doc_ir = self.group(
+            self.concat(
+                [
+                    parts[0],
+                    self.indent(self.concat([self.tight_line(), *parts[2:-2]])),
+                    self.tight_line(),
+                    parts[-1],
+                ]
+            )
+        )
 
     def exit_set_compr(self, node: uni.SetCompr) -> None:
         """Generate DocIR for set comprehensions."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
+            if isinstance(i, uni.InnerCompr):
+                parts.append(self.group(self.concat([self.tight_line(), i.gen.doc_ir])))
+            else:
+                parts.append(i.gen.doc_ir)
             parts.append(self.space())
         parts.pop()
-        node.gen.doc_ir = self.group(self.concat(parts))
+        node.gen.doc_ir = self.group(
+            self.concat(
+                [
+                    parts[0],
+                    self.indent(self.concat([self.tight_line(), *parts[2:-2]])),
+                    self.tight_line(),
+                    parts[-1],
+                ]
+            )
+        )
 
     def exit_dict_compr(self, node: uni.DictCompr) -> None:
         """Generate DocIR for dictionary comprehensions."""
@@ -889,10 +1074,24 @@ class DocIRGenPass(UniPass):
             if isinstance(i, uni.Token) and i.name in [Tok.STAR_POW, Tok.STAR_MUL]:
                 parts.append(i.gen.doc_ir)
             else:
-                parts.append(i.gen.doc_ir)
+                if isinstance(i, uni.InnerCompr):
+                    parts.append(
+                        self.group(self.concat([self.tight_line(), i.gen.doc_ir]))
+                    )
+                else:
+                    parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         parts.pop()
-        node.gen.doc_ir = self.group(self.concat(parts))
+        node.gen.doc_ir = self.group(
+            self.concat(
+                [
+                    parts[0],
+                    self.indent(self.concat([self.tight_line(), *parts[2:-2]])),
+                    self.tight_line(),
+                    parts[-1],
+                ]
+            )
+        )
 
     def exit_k_w_pair(self, node: uni.KWPair) -> None:
         """Generate DocIR for keyword arguments."""
@@ -955,6 +1154,8 @@ class DocIRGenPass(UniPass):
         """Generate DocIR for assert statements."""
         parts: list[doc.DocType] = []
         for i in node.kid:
+            if isinstance(i, uni.Token) and i.name == Tok.SEMI and len(parts):
+                parts.pop()
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
         parts.pop()
@@ -972,6 +1173,7 @@ class DocIRGenPass(UniPass):
     def exit_global_vars(self, node: uni.GlobalVars) -> None:
         """Generate DocIR for global variables."""
         parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         for i in node.kid:
             if i == node.doc:
                 parts.append(i.gen.doc_ir)
@@ -980,15 +1182,22 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_module_code(self, node: uni.ModuleCode) -> None:
         """Generate DocIR for module code."""
         parts: list[doc.DocType] = []
         body_parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         in_body = False
         for i in node.kid:
             if node.doc and i is node.doc:
@@ -1000,6 +1209,10 @@ class DocIRGenPass(UniPass):
             elif isinstance(i, uni.Token) and i.name == Tok.COLON:
                 parts.pop()
                 parts.append(i.gen.doc_ir)
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             elif isinstance(node.body, Sequence) and i in node.body:
                 if not in_body:
                     parts.pop()
@@ -1022,6 +1235,8 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_global_stmt(self, node: uni.GlobalStmt) -> None:
         """Generate DocIR for global statements."""
@@ -1101,7 +1316,25 @@ class DocIRGenPass(UniPass):
                 parts.append(self.space())
             prev_item = i
         parts.pop()
-        node.gen.doc_ir = self.group(self.concat(parts))
+        broken = self.group(
+            self.concat(
+                [
+                    parts[0],
+                    self.indent(self.concat([self.tight_line(), *parts[1:-1]])),
+                    self.tight_line(),
+                    parts[-1],
+                ]
+            )
+        )
+        if isinstance(node.parent, (uni.Assignment)) or (
+            isinstance(node.parent, uni.IfStmt) and isinstance(node.value, uni.BoolExpr)
+        ):
+            node.gen.doc_ir = self.if_break(
+                flat_contents=self.group(self.concat(parts[1:-1])),
+                break_contents=broken,
+            )
+        else:
+            node.gen.doc_ir = broken
 
     def exit_expr_as_item(self, node: uni.ExprAsItem) -> None:
         """Generate DocIR for expression as item nodes."""
@@ -1109,6 +1342,7 @@ class DocIRGenPass(UniPass):
         for i in node.kid:
             parts.append(i.gen.doc_ir)
             parts.append(self.space())
+        parts.pop()
         node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_filter_compr(self, node: uni.FilterCompr) -> None:
@@ -1130,6 +1364,7 @@ class DocIRGenPass(UniPass):
     def exit_py_inline_code(self, node: uni.PyInlineCode) -> None:
         """Generate DocIR for Python inline code blocks."""
         parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         for i in node.kid:
             if i == node.doc:
                 parts.append(i.gen.doc_ir)
@@ -1139,14 +1374,21 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.text("::py::"))
                 parts.append(self.hard_line())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_test(self, node: uni.Test) -> None:
         """Generate DocIR for test nodes."""
         parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         for i in node.kid:
             if i == node.doc:
                 parts.append(i.gen.doc_ir)
@@ -1155,20 +1397,16 @@ class DocIRGenPass(UniPass):
                 if not i.value.startswith("_jac_gen_"):
                     parts.append(i.gen.doc_ir)
                     parts.append(self.space())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
-
-    def exit_check_stmt(self, node: uni.CheckStmt) -> None:
-        """Generate DocIR for check statements."""
-        parts: list[doc.DocType] = []
-        for i in node.kid:
-            if isinstance(i, uni.Token) and i.name == Tok.SEMI:
-                parts.pop()
-            parts.append(i.gen.doc_ir)
-            parts.append(self.space())
-        node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_match_stmt(self, node: uni.MatchStmt) -> None:
         """Generate DocIR for match statements."""
@@ -1352,6 +1590,7 @@ class DocIRGenPass(UniPass):
         """Generate DocIR for implementation definitions."""
         parts: list[doc.DocType] = []
         body_parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         in_body = False
         for i in node.kid:
             if i == node.doc or (node.decorators and i in node.decorators):
@@ -1359,6 +1598,10 @@ class DocIRGenPass(UniPass):
                 parts.append(self.hard_line())
             elif self.is_within(i, node.target):
                 parts.append(i.gen.doc_ir)
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             elif (
                 in_body
                 or isinstance(node.body, Sequence)
@@ -1386,10 +1629,13 @@ class DocIRGenPass(UniPass):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_sem_def(self, node: uni.SemDef) -> None:
         """Generate DocIR for semantic definitions."""
         parts: list[doc.DocType] = []
+        has_comment: Optional[doc.DocType] = None
         for i in node.kid:
             if i in node.target:
                 parts.append(i.gen.doc_ir)
@@ -1397,10 +1643,16 @@ class DocIRGenPass(UniPass):
                 parts.pop()
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
+            elif (
+                i == node.kid[0] and isinstance(i, uni.Token) and i.name == Tok.COMMENT
+            ):
+                has_comment = i.gen.doc_ir
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
+        if has_comment:
+            node.gen.doc_ir = self.concat([has_comment, node.gen.doc_ir])
 
     def exit_event_signature(self, node: uni.EventSignature) -> None:
         """Generate DocIR for event signatures."""
@@ -1454,14 +1706,6 @@ class DocIRGenPass(UniPass):
             and isinstance(node.parent, uni.FString)
         )
 
-        if "\n" in node.value:
-            lines = node.value.split("\n")
-            parts: list[doc.DocType] = [self.text(lines[0])]
-            for line in lines[1:]:
-                parts.append(self.hard_line())
-                parts.append(self.text(line.lstrip()))
-            node.gen.doc_ir = self.group(self.concat(parts))
-            return
         if is_escaped_curly:
             node.gen.doc_ir = self.concat(
                 [self.text(node.value), self.text(node.value)]