feat: better formatting for call chains (#106)

* feat: better formatting for call chains

* feat: correctly parse dot chain

* fix: workaround weird typst parser issue

* fix: break for loop iter and target into separate lines

* test: add 80line test case

* test: apply closure format

* test: split compileable test

* test: add 0 width test cases

* fix: dont touch markup mode chain

* test: add last arg no indent bad case

* test: update ugrad math

src/pretty/dot_chain.rs

@@ -0,0 +1,44 @@
+use pretty::BoxDoc;
+use typst_syntax::ast::*;
+
+use crate::PrettyPrinter;
+
+impl PrettyPrinter {
+    pub fn resolve_dot_chain<'a>(
+        &'a self,
+        field_access: FieldAccess<'a>,
+    ) -> Option<Vec<BoxDoc<'a, ()>>> {
+        let mut nodes = vec![];
+        let mut push_node = |node: BoxDoc<'a, ()>, last_is_field_access: bool| {
+            if last_is_field_access {
+                nodes.push(node);
+            } else {
+                let last = nodes.pop().unwrap();
+                nodes.push(node.append(last));
+            }
+        };
+        let mut current = field_access.to_untyped();
+        let mut last_is_field_access = true;
+        loop {
+            if let Some(field_access) = current.cast::<FieldAccess>() {
+                let rhs = self.convert_ident(field_access.field());
+                push_node(rhs, last_is_field_access);
+                last_is_field_access = true;
+                current = field_access.target().to_untyped();
+                continue;
+            }
+            if let Some(func_call) = current.cast::<FuncCall>() {
+                let args = self.convert_args(func_call.args());
+                push_node(args, last_is_field_access);
+                last_is_field_access = false;
+                current = func_call.callee().to_untyped();
+                continue;
+            }
+            let lhs = self.convert_expr(current.cast::<Expr>().unwrap());
+            push_node(lhs, last_is_field_access);
+            break;
+        }
+        nodes.reverse();
+        Some(nodes)
+    }
+}

src/pretty/func_call.rs

@@ -87,7 +87,7 @@ impl PrettyPrinter {
         if let Some(res) = self.check_disabled(func_call.args().to_untyped()) {
             return doc.append(res);
         }
-        let has_parenthesized_args = util::has_parenthesized_args(func_call);
+        let has_parenthesized_args = util::has_parenthesized_args(func_call.args());
         if table::is_table(func_call) {
             if let Some(cols) = table::is_formatable_table(func_call) {
                 doc = doc.append(self.convert_table(func_call, cols));
@@ -100,6 +100,16 @@ impl PrettyPrinter {
         doc.append(self.convert_additional_args(func_call.args(), has_parenthesized_args))
     }
 
+    pub(super) fn convert_args<'a>(&'a self, args: Args<'a>) -> BoxDoc<'a, ()> {
+        let has_parenthesized_args = util::has_parenthesized_args(args);
+        let mut doc = BoxDoc::nil();
+        if has_parenthesized_args {
+            doc = doc.append(self.convert_parenthesized_args_as_is(args));
+        }
+        doc = doc.append(self.convert_additional_args(args, has_parenthesized_args));
+        doc
+    }
+
     pub(super) fn convert_parenthesized_args<'a>(&'a self, args: Args<'a>) -> BoxDoc<'a, ()> {
         let args = parse_args(args).collect_vec();
         let is_multiline = {

src/pretty/mod.rs

@@ -1,31 +1,40 @@
 use std::borrow::Cow;
+use std::cell::RefCell;
 use std::collections::HashMap;
 
 use itertools::Itertools;
+use mode::Mode;
 use pretty::BoxDoc;
 use typst_syntax::{ast, ast::*, SyntaxKind, SyntaxNode};
 
 use crate::attr::Attributes;
 use crate::util::{comma_seprated_items, pretty_items, FoldStyle};
 
+mod dot_chain;
 mod func_call;
+mod mode;
 mod parened_expr;
 mod table;
 mod util;
 
 #[derive(Debug, Default)]
 pub struct PrettyPrinter {
     attr_map: HashMap<SyntaxNode, Attributes>,
+    mode: RefCell<Vec<Mode>>,
 }
 
 impl PrettyPrinter {
     pub fn new(attr_map: HashMap<SyntaxNode, Attributes>) -> Self {
-        Self { attr_map }
+        Self {
+            attr_map,
+            mode: vec![].into(),
+        }
     }
 }
 
 impl PrettyPrinter {
     pub fn convert_markup<'a>(&'a self, root: Markup<'a>) -> BoxDoc<'a, ()> {
+        let _g = self.with_mode(Mode::Markup);
         let mut doc: BoxDoc<()> = BoxDoc::nil();
         #[derive(Debug, Default)]
         struct Line<'a> {
@@ -340,6 +349,7 @@ impl PrettyPrinter {
     }
 
     fn convert_equation<'a>(&'a self, equation: Equation<'a>) -> BoxDoc<'a, ()> {
+        let _g = self.with_mode(Mode::Math);
         let mut doc = BoxDoc::text("$");
         let is_multi_line = self
             .attr_map
@@ -423,6 +433,7 @@ impl PrettyPrinter {
     }
 
     fn convert_code_block<'a>(&'a self, code_block: CodeBlock<'a>) -> BoxDoc<'a, ()> {
+        let _g = self.with_mode(Mode::Code);
         let mut code_nodes = vec![];
         for node in code_block.to_untyped().children() {
             if let Some(code) = node.cast::<Code>() {
@@ -603,16 +614,32 @@ impl PrettyPrinter {
     }
 
     fn convert_field_access<'a>(&'a self, field_access: FieldAccess<'a>) -> BoxDoc<'a, ()> {
-        let left = BoxDoc::nil().append(self.convert_expr(field_access.target()));
-        let singleline_right = BoxDoc::text(".").append(self.convert_ident(field_access.field()));
-        let _multiline_right = BoxDoc::hardline()
-            .append(BoxDoc::text("."))
-            .append(self.convert_ident(field_access.field()))
-            .nest(2)
-            .group();
-        // TODO: typst doesn't support this
-        // left.append(multiline_right.flat_alt(singleline_right))
-        left.append(singleline_right)
+        let chain: Option<Vec<BoxDoc>> = self.resolve_dot_chain(field_access);
+        if chain.is_none() || matches!(self.current_mode(), Mode::Markup | Mode::Math) {
+            let left = BoxDoc::nil().append(self.convert_expr(field_access.target()));
+            let singleline_right =
+                BoxDoc::text(".").append(self.convert_ident(field_access.field()));
+            return left.append(singleline_right);
+        }
+        let mut chain = chain.unwrap();
+        if chain.len() == 2 {
+            let last = chain.pop().unwrap();
+            let first = chain.pop().unwrap();
+            return first.append(BoxDoc::text(".")).append(last);
+        }
+        let first_doc = chain.remove(0);
+        let other_doc = BoxDoc::intersperse(chain, BoxDoc::line_().append(BoxDoc::text(".")));
+        let chain = first_doc.append(
+            (BoxDoc::line_().append(BoxDoc::text(".")).append(other_doc))
+                .nest(2)
+                .group(),
+        );
+        // if matches!(self.current_mode(), Mode::Markup | Mode::Math) {
+        //     optional_paren(chain)
+        // } else {
+        //     chain
+        // }
+        chain
     }
 
     fn convert_closure<'a>(&'a self, closure: Closure<'a>) -> BoxDoc<'a, ()> {
@@ -876,13 +903,13 @@ impl PrettyPrinter {
     }
 
     fn convert_for<'a>(&'a self, for_loop: ForLoop<'a>) -> BoxDoc<'a, ()> {
-        // let mut doc = BoxDoc::nil();
         let for_pattern = BoxDoc::text("for")
             .append(BoxDoc::space())
             .append(self.convert_pattern(for_loop.pattern()))
             .append(BoxDoc::space());
         let in_iter = BoxDoc::text("in")
             .append(BoxDoc::space())
+            // .append(BoxDoc::softline()) // upstream issue: https://github.com/typst/typst/issues/4548
             .append(self.convert_expr_with_optional_paren(for_loop.iterable()))
             .append(BoxDoc::space());
         let body = self.convert_expr(for_loop.body());

src/pretty/mode.rs

@@ -0,0 +1,38 @@
+use crate::PrettyPrinter;
+
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub enum Mode {
+    #[default]
+    Markup,
+    Code,
+    Math,
+}
+
+impl PrettyPrinter {
+    pub(super) fn push_mode(&self, mode: Mode) {
+        self.mode.borrow_mut().push(mode);
+    }
+
+    pub(super) fn pop_mode(&self) {
+        self.mode.borrow_mut().pop();
+    }
+
+    pub(super) fn current_mode(&self) -> Mode {
+        *self.mode.borrow().last().unwrap_or(&Mode::Markup)
+    }
+}
+
+pub(super) struct ModeGuard<'a>(&'a PrettyPrinter);
+
+impl PrettyPrinter {
+    pub(super) fn with_mode(&self, mode: Mode) -> ModeGuard {
+        self.push_mode(mode);
+        ModeGuard(self)
+    }
+}
+
+impl<'a> Drop for ModeGuard<'a> {
+    fn drop(&mut self) {
+        self.0.pop_mode();
+    }
+}

src/pretty/parened_expr.rs

@@ -41,17 +41,22 @@ impl PrettyPrinter {
         ) {
             return self.convert_expr(expr);
         }
-        let left_paren_or_nil = BoxDoc::text("(")
-            .append(BoxDoc::line())
-            .flat_alt(BoxDoc::nil());
-        let right_paren_or_nil = BoxDoc::line()
-            .append(BoxDoc::text(")"))
-            .flat_alt(BoxDoc::nil());
         let body_expr = self.convert_expr(expr);
-        left_paren_or_nil
-            .append(body_expr)
-            .nest(2)
-            .append(right_paren_or_nil)
-            .group()
+        optional_paren(body_expr)
     }
 }
+
+/// Wrap the body with parentheses if the body is layouted on multiple lines.
+pub(super) fn optional_paren(body: BoxDoc<'_, ()>) -> BoxDoc<'_, ()> {
+    let left_paren_or_nil = BoxDoc::text("(")
+        .append(BoxDoc::line())
+        .flat_alt(BoxDoc::nil());
+    let right_paren_or_nil = BoxDoc::line()
+        .append(BoxDoc::text(")"))
+        .flat_alt(BoxDoc::nil());
+    left_paren_or_nil
+        .append(body)
+        .nest(2)
+        .append(right_paren_or_nil)
+        .group()
+}

src/pretty/util.rs

@@ -11,9 +11,8 @@ pub(super) fn func_name(node: FuncCall<'_>) -> String {
     node.callee().to_untyped().clone().into_text().to_string()
 }
 
-pub(super) fn has_parenthesized_args(node: FuncCall<'_>) -> bool {
-    node.args()
-        .to_untyped()
+pub(super) fn has_parenthesized_args(node: Args<'_>) -> bool {
+    node.to_untyped()
         .children()
         .any(|node| matches!(node.kind(), SyntaxKind::LeftParen | SyntaxKind::RightParen))
 }
@@ -32,10 +31,9 @@ pub(super) fn get_parenthesized_args(node: Args<'_>) -> impl Iterator<Item = Arg
 }
 
 #[allow(unused)]
-pub(super) fn has_additional_args(node: FuncCall<'_>) -> bool {
+pub(super) fn has_additional_args(node: Args<'_>) -> bool {
     let has_paren_args = has_parenthesized_args(node);
     let args = node
-        .args()
         .to_untyped()
         .children()
         .skip_while(|node| {