feat: format chains with comments

src/pretty/chain.rs

@@ -1,42 +1,57 @@
 use itertools::Itertools;
 use pretty::DocAllocator;
 use std::iter;
-use typst_syntax::SyntaxNode;
+use typst_syntax::{SyntaxKind, SyntaxNode};
 
-use super::{ArenaDoc, PrettyPrinter};
+use crate::ext::StrExt;
 
+use super::{util::is_comment_node, ArenaDoc, PrettyPrinter};
+
+/// Intermediate representation in chain formatting.
 enum ChainItem<'a> {
-    Commented { body: ArenaDoc<'a> },
+    Body(ArenaDoc<'a>),
+    Op(ArenaDoc<'a>),
+    Comment(ArenaDoc<'a>),
+    Attached(ArenaDoc<'a>),
+    Linebreak,
 }
 
+/// A stylist that can format items as chains.
 pub struct ChainStylist<'a> {
     printer: &'a PrettyPrinter<'a>,
     items: Vec<ChainItem<'a>>,
-    chain_len: usize,
-    op_space: bool,
+    /// The number of chain operators in the chain.
+    chain_op_num: usize,
+    /// Whether the chain contains any line or block comment.
+    has_comment: bool,
 }
 
 #[derive(Default)]
 pub struct ChainStyle {
     /// Do not break line if the chain consists of only one operator.
     pub no_break_single: bool,
+    /// Add space before and after operators.
+    pub space_around_op: bool,
 }
 
 impl<'a> ChainStylist<'a> {
     pub fn new(printer: &'a PrettyPrinter<'a>) -> Self {
         Self {
             printer,
             items: Default::default(),
-            chain_len: 0,
-            op_space: false,
+            chain_op_num: 0,
+            has_comment: false,
         }
     }
 
-    pub fn space_around_op(mut self) -> Self {
-        self.op_space = true;
-        self
-    }
-
+    /// Processes a collection of syntax nodes directly from depth-first resolution.
+    ///
+    /// This method takes an iterator of `SyntaxNode`s, which are then processed in reverse order.
+    ///
+    /// # Parameters
+    ///
+    /// - `nodes`: An iterator over references to `SyntaxNode`s that have been resolved.
+    /// - Others: See [`Self::process`].
     pub fn process_resolved(
         self,
         nodes: impl Iterator<Item = &'a SyntaxNode>,
@@ -56,6 +71,17 @@ impl<'a> ChainStylist<'a> {
         )
     }
 
+    /// Processes a vector of syntax nodes with the provided predicates and converters
+    /// to create a structured representation.
+    ///
+    /// # Parameters
+    ///
+    /// - `nodes`: A vector of `SyntaxNode`s to be processed.
+    /// - `operand_pred`: A predicate that checks if a node is an operand.
+    /// - `op_pred`: A predicate that checks if a node is an operator.
+    /// - `rhs_converter`: A function that converts right-hand side nodes into an `Option<ArenaDoc<'a>>`.
+    /// - `fallback_converter`: A function that provides a fallback conversion for nodes that
+    ///    do not match the primary criteria. Used for sticky args and innermost expressions.
     pub fn process(
         mut self,
         nodes: Vec<&'a SyntaxNode>,
@@ -65,57 +91,129 @@ impl<'a> ChainStylist<'a> {
         fallback_converter: impl Fn(&'a SyntaxNode) -> Option<ArenaDoc<'a>>,
     ) -> Self {
         let arena = &self.printer.arena;
-
-        let mut doc = arena.nil();
+        let mut can_attach = false;
         for node in nodes {
             if operand_pred(node) {
-                self.chain_len += 1;
+                self.chain_op_num += 1;
                 let mut seen_op = false;
                 for child in node.children() {
                     if op_pred(child) {
                         seen_op = true;
-                        self.items.push(ChainItem::Commented { body: doc });
-                        doc = if self.op_space {
-                            arena.text(child.text().as_str()) + " "
+                        self.items
+                            .push(ChainItem::Op(arena.text(child.text().as_str())));
+                    } else if is_comment_node(child) {
+                        let doc = self.printer.convert_comment(child);
+                        self.items.push(if can_attach {
+                            ChainItem::Attached(doc)
                         } else {
-                            arena.text(child.text().as_str())
-                        };
+                            ChainItem::Comment(doc)
+                        });
+                        self.has_comment = true;
+                    } else if child.kind() == SyntaxKind::Space {
+                        if child.text().has_linebreak() {
+                            if self.items.last().is_some_and(|last| {
+                                matches!(*last, ChainItem::Attached(_) | ChainItem::Comment(_))
+                            }) {
+                                self.items.push(ChainItem::Linebreak);
+                            }
+                            can_attach = false;
+                        }
                     } else if seen_op {
                         if let Some(rhs) = rhs_converter(child) {
-                            doc += rhs;
+                            self.items.push(ChainItem::Body(rhs));
+                            can_attach = true;
                         }
                     }
                 }
             } else if let Some(fallback) = fallback_converter(node) {
-                doc += fallback;
+                // We must use this to handle args.
+                if let Some(ChainItem::Body(body)) = self.items.last_mut() {
+                    *body += fallback;
+                } else {
+                    self.items.push(ChainItem::Body(fallback));
+                }
             }
         }
-        self.items.push(ChainItem::Commented { body: doc });
 
         self
     }
 
+    /// Create a Doc from IR and given styles.
     pub fn print_doc(self, sty: ChainStyle) -> ArenaDoc<'a> {
         let arena = &self.printer.arena;
 
+        let op_sep = if sty.space_around_op {
+            arena.line()
+        } else {
+            arena.line_()
+        };
+
+        let use_simple_layout = self.chain_op_num == 1 && sty.no_break_single && !self.has_comment;
+
         let mut docs = vec![];
+        let mut has_break = false;
+        let mut leading = true;
+        let mut space_after = true;
         for item in self.items {
             match item {
-                ChainItem::Commented { body } => docs.push(body),
+                ChainItem::Body(body) => {
+                    if leading {
+                        docs.push(body);
+                    } else if let Some(last) = docs.last_mut() {
+                        *last += body;
+                    }
+                    leading = false;
+                    space_after = true;
+                }
+                ChainItem::Op(op) => {
+                    if !(has_break && leading || use_simple_layout) {
+                        docs.push(op_sep.clone());
+                    }
+                    has_break = false;
+                    if sty.space_around_op {
+                        docs.push(op + " ");
+                    } else {
+                        docs.push(op);
+                    }
+                    leading = false;
+                    space_after = false;
+                }
+                ChainItem::Comment(cmt) => {
+                    if leading {
+                        docs.push(cmt);
+                    } else if let Some(last) = docs.last_mut() {
+                        *last += if space_after {
+                            arena.space() + cmt
+                        } else {
+                            cmt
+                        }
+                    }
+                    leading = false;
+                    space_after = true;
+                }
+                ChainItem::Attached(cmt) => {
+                    if let Some(last) = docs.last_mut() {
+                        *last += if space_after {
+                            arena.space() + cmt
+                        } else {
+                            cmt
+                        }
+                    }
+                }
+                ChainItem::Linebreak => {
+                    has_break = true;
+                    leading = true;
+                    docs.push(arena.hardline());
+                }
             }
         }
 
-        let op_sep = if self.op_space {
-            arena.line()
-        } else {
-            arena.line_()
-        };
         let first_doc = docs.remove(0);
-        let follow_docs = arena.intersperse(docs, op_sep.clone());
-        if self.chain_len == 1 && sty.no_break_single {
+        let follow_docs = arena.concat(docs);
+        if use_simple_layout {
             (first_doc + follow_docs).group()
         } else {
-            (first_doc + (op_sep + follow_docs).nest(2)).group()
+            (first_doc + (follow_docs).nest(2)).group()
         }
     }
 }

src/pretty/code_chain.rs

@@ -5,15 +5,11 @@ use crate::PrettyPrinter;
 
 use super::{
     chain::{iterate_deep_nodes, ChainStyle, ChainStylist},
-    util::has_comment_children,
     ArenaDoc,
 };
 
 impl<'a> PrettyPrinter<'a> {
     pub(super) fn convert_field_access(&'a self, field_access: FieldAccess<'a>) -> ArenaDoc<'a> {
-        if let Some(res) = self.check_unformattable(field_access.to_untyped()) {
-            return res;
-        }
         if self.current_mode().is_code() {
             return self.convert_dot_chain(field_access.to_untyped());
         }
@@ -23,10 +19,6 @@ impl<'a> PrettyPrinter<'a> {
     }
 
     pub(super) fn convert_dot_chain(&'a self, node: &'a SyntaxNode) -> ArenaDoc<'a> {
-        if resolve_dot_chain(node).any(has_comment_children) {
-            return self.format_disabled(node);
-        }
-
         ChainStylist::new(self)
             .process_resolved(
                 resolve_dot_chain(node),
@@ -44,13 +36,13 @@ impl<'a> PrettyPrinter<'a> {
             )
             .print_doc(ChainStyle {
                 no_break_single: true,
+                ..Default::default()
             })
     }
 
     pub(super) fn convert_binary_chain(&'a self, binary: Binary<'a>) -> ArenaDoc<'a> {
         let prec = binary.op().precedence();
         ChainStylist::new(self)
-            .space_around_op()
             .process_resolved(
                 resolve_binary_chain(binary),
                 |node| {
@@ -62,6 +54,7 @@ impl<'a> PrettyPrinter<'a> {
                 |node| node.cast().map(|expr| self.convert_expr(expr)),
             )
             .print_doc(ChainStyle {
+                space_around_op: true,
                 ..Default::default()
             })
     }

src/pretty/code_flow.rs

@@ -3,10 +3,7 @@ use typst_syntax::{ast::*, SyntaxKind};
 
 use crate::ext::BoolExt;
 
-use super::{
-    code_chain::resolve_binary_chain, flow::FlowItem, util::has_comment_children, ArenaDoc,
-    PrettyPrinter,
-};
+use super::{flow::FlowItem, ArenaDoc, PrettyPrinter};
 
 impl<'a> PrettyPrinter<'a> {
     pub(super) fn convert_named(&'a self, named: Named<'a>) -> ArenaDoc<'a> {
@@ -77,8 +74,7 @@ impl<'a> PrettyPrinter<'a> {
         if matches!(
             binary.op(),
             BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::And | BinOp::Or
-        ) && !resolve_binary_chain(binary).any(has_comment_children)
-        {
+        ) {
             return self.parenthesize_if_necessary(|| self.convert_binary_chain(binary));
         }
         self.convert_flow_like(binary.to_untyped(), |child| {

tests/assets/unit/comment/comment-in-chain-binary.typ

@@ -1,8 +1,7 @@
-#{  let a = (1111111111111   // 1
- + 2222222222222/* 2 */
- +/* 3 */ 3333333333333 / 2
- /* 4 */+ 4444444444444  + 5555555555555
-/* 5 */ + 444444444
+#{  let a = (1111111111111   // -
+ + 2222222222222/* 0 */    /* 1 */
+ +/* 2 */ 3333333333333 / 2  /* 3 */ /* 4 */
+ /* 5 */+ 4444444444444  + 5555555555555
  - 6666666/* 6 */
  + /* 7 */7777777
  - 9999999 / /* 8 */ 10000000 // 9

tests/assets/unit/comment/comment-in-chain-call.typ

@@ -3,12 +3,15 @@
   .b(1)(2)
  .c   .d(3)
 
+
   a // 1
   .b(1)(2)  // 2
   // 3
 /* 4 */  ./* 5 */c  /* 6 *//* 7 */  .d(3) // 8
 
+
 /* 0 */    a/* 1 */
 /* 2 */  .b(1)(2) /* 3 */ .c   // 4
+
 /* 5 */  ./* 6 */d(3)    // 7
 }

tests/snapshots/assets__check_file@tablex.typ-0.snap

@@ -3253,8 +3253,8 @@ snapshot_kind: text
           auto,
           none,
         ) // always goes towards the right
-        or h.end >= cell.x
-          + cell.colspan // ends at or after this cell
+          or h.end >= cell.x
+            + cell.colspan // ends at or after this cell
       )
 
       (
@@ -3303,8 +3303,8 @@ snapshot_kind: text
           auto,
           none,
         ) // always goes towards the bottom
-        or v.end >= cell.y
-          + cell.rowspan // ends at or after this cell
+          or v.end >= cell.y
+            + cell.rowspan // ends at or after this cell
       )
 
       (

tests/snapshots/assets__check_file@tablex.typ-120.snap

@@ -1923,7 +1923,7 @@ snapshot_kind: text
 
       let hline_hasnt_already_ended = (
         h.end in (auto, none) // always goes towards the right
-        or h.end >= cell.x + cell.colspan // ends at or after this cell
+          or h.end >= cell.x + cell.colspan // ends at or after this cell
       )
 
       (in_top_or_bottom and hline_hasnt_already_ended)
@@ -1949,7 +1949,7 @@ snapshot_kind: text
 
       let vline_hasnt_already_ended = (
         v.end in (auto, none) // always goes towards the bottom
-        or v.end >= cell.y + cell.rowspan // ends at or after this cell
+          or v.end >= cell.y + cell.rowspan // ends at or after this cell
       )
 
       (at_left_or_right and vline_hasnt_already_ended)