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

jaclang/compiler/passes/tool/doc_ir_gen_pass.py

@@ -3,7 +3,7 @@
 This is a pass for generating DocIr for Jac code.
 """
 
-from typing import List, Optional
+from typing import List, Optional, Sequence
 
 import jaclang.compiler.passes.tool.doc_ir as doc
 import jaclang.compiler.unitree as uni
@@ -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,14 +119,15 @@ class DocIRGenPass(UniPass):
         """Exit import node."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            if isinstance(i, uni.SubNodeList) and i.items:
-                parts.append(
-                    self.indent(self.concat([self.tight_line(), i.gen.doc_ir]))
-                )
-                parts.append(self.line())
+            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())
@@ -179,13 +160,45 @@ class DocIRGenPass(UniPass):
     def exit_archetype(self, node: uni.Archetype) -> None:
         """Generate DocIR for archetypes."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        prev_item = None
+        in_body = False
         for i in node.kid:
-            if i in [node.doc, node.decorators]:
+            if (node.doc and i is node.doc) or (
+                node.decorators and i in node.decorators
+            ):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.hard_line())
             elif i == node.name:
                 parts.append(i.gen.doc_ir)
                 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())
+                if (prev_item and type(prev_item) is not type(i)) or (
+                    prev_item and not self.is_one_line(prev_item)
+                ):
+                    body_parts.append(self.hard_line())
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                prev_item = i
+                in_body = True
+            elif in_body:
+                in_body = False
+                body_parts.pop()
+                parts.append(self.indent(self.concat(body_parts)))
+                parts.append(self.hard_line())
+                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)
@@ -193,43 +206,85 @@ 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_ability(self, node: uni.Ability) -> None:
         """Generate DocIR for abilities."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        in_body = False
         for i in node.kid:
-            if i in [node.doc, node.decorators]:
+            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.name_ref:
                 parts.append(i.gen.doc_ir)
                 if not isinstance(node.signature, uni.FuncSignature):
                     parts.append(self.space())
+            elif isinstance(node.body, Sequence) and i in node.body:
+                if not in_body:
+                    parts.pop()
+                    body_parts.append(self.hard_line())
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                in_body = True
+            elif in_body:
+                in_body = False
+                body_parts.pop()
+                parts.append(self.indent(self.concat(body_parts)))
+                parts.append(self.hard_line())
+                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)
+                parts.append(self.space())
             else:
                 parts.append(i.gen.doc_ir)
-                if not isinstance(i, uni.CommentToken):
-                    parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+                parts.append(self.space())
+        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:
+            if isinstance(i, uni.Token) and i.name == Tok.LPAREN and node.params:
+                in_params = True
                 parts.append(i.gen.doc_ir)
-            elif i == node.params:
+            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(), i.gen.doc_ir]))
+                    self.indent(self.concat([self.tight_line(), *indent_parts]))
                 )
                 parts.append(self.tight_line())
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            elif isinstance(i, uni.Token) and i.name == Tok.RPAREN:
+                parts.pop()
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            elif in_params:
+                if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                    indent_parts.append(i.gen.doc_ir)
+                    indent_parts.append(self.space())
+                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."""
@@ -241,51 +296,140 @@ 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."""
+        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:
+                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."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        in_body = False
         for i in node.kid:
-            if isinstance(i, uni.Token) and i.name == Tok.SEMI:
+            if isinstance(node.body, Sequence) and i in node.body:
+                if not in_body:
+                    parts.pop()
+                    body_parts.append(self.hard_line())
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                in_body = True
+            elif in_body:
+                in_body = False
+                body_parts.pop()
+                parts.append(self.indent(self.concat(body_parts)))
+                parts.append(self.hard_line())
+                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)
                 parts.append(self.space())
             else:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
-
-    def exit_if_stmt(self, node: uni.IfStmt) -> None:
-        """Generate DocIR for if statements."""
-        parts: list[doc.DocType] = []
-        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_else_if(self, node: uni.ElseIf) -> None:
         """Generate DocIR for else if statements."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        in_body = 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(node.body, Sequence) and i in node.body:
+                if not in_body:
+                    parts.pop()
+                    body_parts.append(self.hard_line())
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                in_body = True
+            elif in_body:
+                in_body = False
+                body_parts.pop()
+                parts.append(self.indent(self.concat(body_parts)))
+                parts.append(self.hard_line())
+                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)
+                parts.append(self.space())
+            else:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_else_stmt(self, node: uni.ElseStmt) -> None:
         """Generate DocIR for else statements."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        in_body = 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(node.body, Sequence) and i in node.body:
+                if not in_body:
+                    parts.pop()
+                    body_parts.append(self.hard_line())
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                in_body = True
+            elif in_body:
+                in_body = False
+                body_parts.pop()
+                parts.append(self.indent(self.concat(body_parts)))
+                parts.append(self.hard_line())
+                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)
+                parts.append(self.space())
+            else:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_binary_expr(self, node: uni.BinaryExpr) -> None:
         """Generate DocIR for binary expressions."""
@@ -293,23 +437,14 @@ 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."""
         parts: list[doc.DocType] = []
-        is_fstring = (
-            node.parent
-            and isinstance(node.parent, uni.SubNodeList)
-            and node.parent.parent
-            and isinstance(node.parent.parent, uni.FString)
-        )
         for i in node.kid:
-            if is_fstring:
-                parts.append(self.text("{"))
             parts.append(i.gen.doc_ir)
-            if is_fstring:
-                parts.append(self.text("}"))
         node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_concurrent_expr(self, node: uni.ConcurrentExpr) -> None:
@@ -318,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."""
@@ -330,13 +465,32 @@ 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."""
         parts: list[doc.DocType] = []
+        indent_parts: list[doc.DocType] = []
+        in_params = False
         for i in node.kid:
-            parts.append(i.gen.doc_ir)
+            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
+                parts.append(
+                    self.indent(self.concat([self.tight_line(), *indent_parts]))
+                )
+                parts.append(self.tight_line())
+                parts.append(i.gen.doc_ir)
+            elif in_params:
+                if isinstance(i, uni.Token) and i.name == Tok.COMMA:
+                    indent_parts.append(i.gen.doc_ir)
+                    indent_parts.append(self.space())
+                else:
+                    indent_parts.append(i.gen.doc_ir)
+            else:
+                parts.append(i.gen.doc_ir)
         node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_atom_trailer(self, node: uni.AtomTrailer) -> None:
@@ -350,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."""
@@ -363,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."""
@@ -371,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."""
@@ -395,16 +566,14 @@ class DocIRGenPass(UniPass):
             if i == node.doc:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.hard_line())
-            elif (
-                isinstance(i, uni.SubNodeList)
-                and i == node.vars
-                and isinstance(i.gen.doc_ir, doc.Concat)
-            ):
-                parts.append(self.align(i.gen.doc_ir))
+            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_while_stmt(self, node: uni.WhileStmt) -> None:
         """Generate DocIR for while statements."""
@@ -412,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."""
@@ -420,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."""
@@ -428,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."""
@@ -436,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."""
@@ -444,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."""
@@ -452,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."""
@@ -472,12 +657,7 @@ class DocIRGenPass(UniPass):
         """Generate DocIR for set values."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            if isinstance(i, uni.Token):
-                parts.append(i.gen.doc_ir)
-            elif isinstance(i, uni.SubNodeList):
-                parts.append(self.if_break(self.line(), self.space()))
-                parts.append(i.gen.doc_ir)
-                parts.append(self.if_break(self.line(), self.space()))
+            parts.append(i.gen.doc_ir)
 
         node.gen.doc_ir = self.group(self.concat(parts))
 
@@ -487,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."""
@@ -495,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."""
@@ -503,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."""
@@ -526,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:
@@ -539,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:
@@ -605,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."""
@@ -613,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."""
@@ -621,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."""
@@ -632,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."""
@@ -647,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."""
@@ -655,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)."""
@@ -670,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."""
@@ -685,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."""
@@ -693,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."""
@@ -701,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."""
@@ -717,26 +909,45 @@ 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."""
         parts: list[doc.DocType] = []
+        body_parts: list[doc.DocType] = []
+        in_body = False
         for i in node.kid:
-            if i == node.doc:
+            if node.doc and i is node.doc:
                 parts.append(i.gen.doc_ir)
                 parts.append(self.hard_line())
-            elif isinstance(i, uni.Token):
+            elif i == node.name:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            elif isinstance(node.body, Sequence) and i in node.body:
+                if not in_body:
+                    parts.pop()
+                    body_parts.append(self.hard_line())
+                body_parts.append(i.gen.doc_ir)
+                body_parts.append(self.hard_line())
+                in_body = True
+            elif in_body:
+                in_body = False
+                body_parts.pop()
+                parts.append(self.indent(self.concat(body_parts)))
+                parts.append(self.hard_line())
+                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)
                 parts.append(self.space())
             else:
-                parts.append(self.concat([i.gen.doc_ir]))
-
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
         node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_global_stmt(self, node: uni.GlobalStmt) -> None:
         """Generate DocIR for global statements."""
-        # node.kid is [GLOBAL_OP_token, name_list_SubNodeList, SEMI_token]
         parts: list[doc.DocType] = []
         for i in node.kid:
             parts.append(i.gen.doc_ir)
@@ -748,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."""
@@ -761,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."""
@@ -769,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."""
@@ -777,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."""
@@ -785,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."""
@@ -793,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)."""
@@ -815,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."""
@@ -823,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."""
@@ -831,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."""
@@ -839,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."""
@@ -856,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."""
@@ -872,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."""
@@ -882,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."""
@@ -902,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."""
@@ -919,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."""
@@ -927,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."""
@@ -935,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."""
@@ -943,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."""
@@ -951,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."""
@@ -959,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."""
@@ -967,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."""
@@ -975,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)."""
@@ -983,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."""
@@ -994,21 +1217,27 @@ 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."""
         parts: list[doc.DocType] = []
         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.COMMA:
+                parts.append(i.gen.doc_ir)
+                parts.append(self.space())
+            else:
+                parts.append(i.gen.doc_ir)
+        node.gen.doc_ir = self.group(self.concat(parts))
 
     def exit_enum(self, node: uni.Enum) -> None:
         """Generate DocIR for enum declarations."""
         parts: list[doc.DocType] = []
         for i in node.kid:
-            if i in [node.doc, node.decorators]:
+            if (node.doc and i is node.doc) or (
+                node.decorators and i in node.decorators
+            ):
                 parts.append(i.gen.doc_ir)
                 parts.append(self.hard_line())
             elif isinstance(i, uni.Token) and i.name == Tok.SEMI:
@@ -1018,93 +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] = []
-        for i in node.kid:
-            parts.append(i.gen.doc_ir)
-            parts.append(self.space())
-        node.gen.doc_ir = self.finalize(parts)
+        before_colon: list[doc.DocType] = []
+        after_colon: list[doc.DocType] = []
+        seen_colon = False
 
-    def exit_sub_node_list(self, node: uni.SubNodeList) -> None:
-        """Generate DocIR for a SubNodeList."""
-        parts: list[doc.DocType] = []
-        indent_parts: list[doc.DocType] = []
-        prev_item: Optional[uni.UniNode] = None
-        in_codeblock = node.delim and node.delim.name == Tok.WS
-        in_archetype = (
-            node.parent
-            and isinstance(node.parent, (uni.Archetype, uni.Enum))
-            and node == node.parent.body
-        )
-        is_assignment = node.delim and node.delim.name == Tok.EQ
         for i in node.kid:
-            if i == node.left_enc:
-                parts.append(i.gen.doc_ir)
-                indent_parts.append(self.hard_line())
-            elif i == node.right_enc:
-                if in_codeblock:
-                    indent_parts.pop()
-                parts.append(self.indent(self.concat(indent_parts)))
-                if prev_item != node.left_enc:
-                    parts.append(self.line())
-                parts.append(i.gen.doc_ir)
-            elif (
-                isinstance(i, uni.Token)
-                and node.delim
-                and i.name == node.delim.name
-                and i.name not in [Tok.DOT, Tok.DECOR_OP]
-            ):
-                indent_parts.append(i.gen.doc_ir)
-                indent_parts.append(self.line())
+            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:
-                if (
-                    in_codeblock
-                    and prev_item
-                    and self.has_gap(prev_item, i)
-                    and isinstance(i, uni.CommentToken)
-                    and (
-                        not in_archetype
-                        or not prev_item
-                        or not (
-                            self.is_one_line(prev_item) and type(prev_item) is type(i)
-                        )
-                    )
-                ) or (
-                    in_archetype
-                    and prev_item not in [None, node.left_enc, node.right_enc]
-                    and (
-                        type(prev_item) is not type(i)
-                        or (prev_item and not self.is_one_line(prev_item))
-                    )
-                ):
-                    indent_parts.append(self.hard_line())
-                indent_parts.append(i.gen.doc_ir)
-                if in_codeblock:
-                    indent_parts.append(self.hard_line())
-                elif is_assignment:
-                    indent_parts.append(self.space())
-            prev_item = i
-        if is_assignment:
-            indent_parts.pop()
-        node.gen.doc_ir = self.concat(parts) if parts else self.concat(indent_parts)
+                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 in [node.doc, node.decorators]:
+            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."""
@@ -1112,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."""
@@ -1154,9 +1365,7 @@ class DocIRGenPass(UniPass):
         is_escaped_curly = (
             node.lit_value in ["{", "}"]
             and node.parent
-            and isinstance(node.parent, uni.SubNodeList)
-            and node.parent.parent
-            and isinstance(node.parent.parent, uni.FString)
+            and isinstance(node.parent, uni.FString)
         )
 
         if "\n" in node.value: