feat: format args and hashed expressions in math

crates/typstyle-core/src/attr.rs

@@ -1,17 +1,12 @@
 use rustc_hash::FxHashMap;
 
 use typst_syntax::{
-    ast::{Args, AstNode, Math, Raw},
+    ast::{Args, AstNode},
     Span, SyntaxKind, SyntaxNode,
 };
 
 use crate::ext::StrExt;
 
-#[derive(Clone, Copy)]
-struct State {
-    is_math: bool,
-}
-
 #[derive(Debug, Clone, Default)]
 pub struct Attributes {
     /// Indicates whether formatting is explicitly disabled (`@typstyle off`) or always ignored.
@@ -118,30 +113,10 @@ impl AttrStore {
     }
 
     fn compute_no_format(&mut self, root: &SyntaxNode) {
-        self.compute_no_format_impl(root, State { is_math: false });
+        self.compute_no_format_impl(root);
     }
 
-    fn compute_no_format_impl(&mut self, node: &SyntaxNode, state: State) {
-        let state = if node.is::<Math>() {
-            State { is_math: true }
-        } else {
-            state
-        };
-
-        // no format multiline single backtick raw block
-        if node
-            .cast::<Raw>()
-            .is_some_and(|raw| !raw.block() && raw.lines().count() > 1)
-        {
-            self.set_format_disabled(node);
-            return;
-        }
-        // no format args in math blocks
-        if node.kind() == SyntaxKind::Args && state.is_math {
-            self.set_format_disabled(node);
-            return;
-        }
-
+    fn compute_no_format_impl(&mut self, node: &SyntaxNode) {
         let mut disable_next = false;
         let mut commented = false;
         for child in node.children() {
@@ -161,12 +136,7 @@ impl AttrStore {
                 disable_next = false;
                 continue;
             }
-            // no format hash related nodes in math blocks
-            if child_kind == SyntaxKind::Hash && state.is_math {
-                self.set_format_disabled(node);
-                break;
-            }
-            self.compute_no_format_impl(child, state);
+            self.compute_no_format_impl(child);
         }
         if commented {
             self.set_commented(node);

crates/typstyle-core/src/pretty/code_flow.rs

@@ -11,9 +11,6 @@ impl<'a> PrettyPrinter<'a> {
         self.convert_flow_like(named.to_untyped(), |child| {
             if child.kind() == SyntaxKind::Colon {
                 FlowItem::tight_spaced(self.arena.text(":"))
-            } else if child.kind() == SyntaxKind::Hash {
-                // name
-                FlowItem::spaced_tight(self.arena.text("#"))
             } else if let Some(expr) = child.cast() {
                 // expr
                 FlowItem::spaced_before(self.convert_expr(expr), seen_name.replace(true))

crates/typstyle-core/src/pretty/code_list.rs

@@ -77,14 +77,32 @@ impl<'a> PrettyPrinter<'a> {
             })
     }
 
+    /// In math mode, we have `$fun(1, 2; 3, 4)$ == $fun(#(1, 2), #(3, 4))$`.
     pub(super) fn convert_array(&'a self, array: Array<'a>) -> ArenaDoc<'a> {
         let _g = self.with_mode(Mode::CodeCont);
 
+        // Whether the array has parens.
+        // This is also used to determine whether we need to add a trailing comma.
+        // Note that we should not strip trailing commas in math.
+        let is_explicit = array
+            .to_untyped()
+            .children()
+            .next()
+            .is_some_and(|child| child.kind() == SyntaxKind::LeftParen);
+        let ends_with_comma = !is_explicit
+            && array
+                .to_untyped()
+                .children()
+                .last()
+                .is_some_and(|child| child.kind() == SyntaxKind::Comma);
+
         ListStylist::new(self)
             .with_fold_style(self.get_fold_style(array))
             .process_list(array.to_untyped(), |node| self.convert_array_item(node))
             .print_doc(ListStyle {
-                add_trailing_sep_single: true,
+                add_trailing_sep_single: is_explicit,
+                add_trailing_sep_always: ends_with_comma,
+                delim: if is_explicit { ("(", ")") } else { ("", "") },
                 ..Default::default()
             })
     }

crates/typstyle-core/src/pretty/flow.rs

@@ -3,7 +3,7 @@ use typst_syntax::{SyntaxKind, SyntaxNode};
 
 use crate::ext::StrExt;
 
-use super::{util::is_comment_node, ArenaDoc, PrettyPrinter};
+use super::{util::is_comment_node, ArenaDoc, Mode, PrettyPrinter};
 
 /// An item in the flow. A space is added only when the item before and the item after both allow it.
 pub struct FlowItem<'a>(Option<FlowItemRepr<'a>>);
@@ -100,34 +100,41 @@ impl<'a> PrettyPrinter<'a> {
         node: &'a SyntaxNode,
         producer: impl FnMut(&'a SyntaxNode) -> FlowItem<'a>,
     ) -> ArenaDoc<'a> {
-        self.convert_flow_like_sliced(node.children(), producer)
+        self.convert_flow_like_iter(node.children(), producer)
     }
 
-    pub(super) fn convert_flow_like_sliced(
+    pub(super) fn convert_flow_like_iter(
         &'a self,
-        children: std::slice::Iter<'a, SyntaxNode>,
+        children: impl Iterator<Item = &'a SyntaxNode>,
         mut producer: impl FnMut(&'a SyntaxNode) -> FlowItem<'a>,
     ) -> ArenaDoc<'a> {
         let mut flow = FlowStylist::new(self);
-        let mut seen_line_comment = false;
+        let mut peek_line_comment = false;
+        let mut peek_hash = false;
         for child in children {
-            let peek_line_comment = seen_line_comment;
-            seen_line_comment = false;
+            let at_line_comment = peek_line_comment;
+            peek_line_comment = false;
+            let at_hash = peek_hash;
+            peek_hash = false;
             if child.kind().is_keyword()
                 && !matches!(child.kind(), SyntaxKind::None | SyntaxKind::Auto)
             {
                 flow.push_doc(self.arena.text(child.text().as_str()), true, true);
             } else if is_comment_node(child) {
                 if child.kind() == SyntaxKind::LineComment {
-                    seen_line_comment = true; // defers the linebreak
+                    peek_line_comment = true; // defers the linebreak
                 }
                 flow.push_comment(child);
-            } else if peek_line_comment
+            } else if at_line_comment
                 && child.kind() == SyntaxKind::Space
                 && child.text().has_linebreak()
             {
                 flow.push_doc(self.arena.hardline(), false, false);
+            } else if child.kind() == SyntaxKind::Hash {
+                flow.push_doc(self.arena.text("#"), true, false);
+                peek_hash = true;
             } else {
+                let _g = self.with_mode_if(Mode::Code, at_hash);
                 let item = producer(child);
                 if let Some(repr) = item.0 {
                     flow.push_doc(repr.doc, repr.space_before, repr.space_after);

crates/typstyle-core/src/pretty/func_call.rs

@@ -1,6 +1,9 @@
 use pretty::DocAllocator;
-use typst_syntax::{ast::*, SyntaxKind};
+use typst_syntax::{ast::*, SyntaxKind, SyntaxNode};
 
+use crate::ext::StrExt;
+
+use super::flow::FlowItem;
 use super::list::{ListStyle, ListStylist};
 use super::mode::Mode;
 use super::plain::PlainStylist;
@@ -31,17 +34,16 @@ impl<'a> PrettyPrinter<'a> {
 
     fn convert_func_call_args(&'a self, func_call: FuncCall<'a>, args: Args<'a>) -> ArenaDoc<'a> {
         if self.current_mode().is_math() {
-            return self.format_disabled(args.to_untyped());
-        }
-        let _g = self.with_mode(Mode::CodeCont);
+            return self.convert_args_in_math(args);
+        };
 
         let mut doc = self.arena.nil();
         let has_parenthesized_args = has_parenthesized_args(args);
         if table::is_table(func_call) {
             if let Some(cols) = table::is_formatable_table(func_call) {
                 doc += self.convert_table(func_call, cols);
             } else if has_parenthesized_args {
-                doc += self.convert_parenthesized_args_as_is(args);
+                doc += self.convert_parenthesized_args_as_list(args);
             }
         } else if has_parenthesized_args {
             doc += self.convert_parenthesized_args(args);
@@ -105,7 +107,7 @@ impl<'a> PrettyPrinter<'a> {
             })
     }
 
-    fn convert_parenthesized_args_as_is(&'a self, args: Args<'a>) -> ArenaDoc<'a> {
+    fn convert_parenthesized_args_as_list(&'a self, args: Args<'a>) -> ArenaDoc<'a> {
         let _g = self.with_mode(Mode::CodeCont);
 
         let inner = PlainStylist::new(self)
@@ -116,6 +118,65 @@ impl<'a> PrettyPrinter<'a> {
         inner.nest(self.config.tab_spaces as isize).parens()
     }
 
+    fn convert_args_in_math(&'a self, args: Args<'a>) -> ArenaDoc<'a> {
+        // strip spaces
+        let children = {
+            let children = args.to_untyped().children().as_slice();
+            let i = children
+                .iter()
+                .position(|child| {
+                    !matches!(child.kind(), SyntaxKind::LeftParen | SyntaxKind::Space)
+                })
+                .unwrap_or(0);
+            let j = children
+                .iter()
+                .rposition(|child| {
+                    !matches!(child.kind(), SyntaxKind::RightParen | SyntaxKind::Space)
+                })
+                .unwrap_or(children.len().saturating_sub(1));
+            children[i..=j].iter()
+        };
+
+        let mut peek_hashed_arg = false;
+        let inner = self.convert_flow_like_iter(children, |child| {
+            let at_hashed_arg = peek_hashed_arg;
+            peek_hashed_arg = false;
+            match child.kind() {
+                SyntaxKind::Comma => FlowItem::tight_spaced(self.arena.text(",")),
+                SyntaxKind::Semicolon => {
+                    // We should avoid the semicolon counted the terminator of the previous hashed arg.
+                    FlowItem::new(self.arena.text(";"), at_hashed_arg, true)
+                }
+                SyntaxKind::Space => {
+                    peek_hashed_arg = at_hashed_arg;
+                    if child.text().has_linebreak() {
+                        FlowItem::tight(self.arena.hardline())
+                    } else {
+                        FlowItem::none()
+                    }
+                }
+                _ => {
+                    if let Some(arg) = child.cast::<Arg>() {
+                        if is_ends_with_hashed_expr(arg.to_untyped().children()) {
+                            peek_hashed_arg = true;
+                        }
+                        FlowItem::spaced(self.convert_arg(arg))
+                    } else {
+                        FlowItem::none()
+                    }
+                }
+            }
+        });
+        if self.attr_store.is_multiline(args.to_untyped()) {
+            ((self.arena.line_() + inner).nest(self.config.tab_spaces as isize)
+                + self.arena.line_())
+            .group()
+            .parens()
+        } else {
+            inner.parens()
+        }
+    }
+
     /// Handle additional content blocks
     fn convert_additional_args(&'a self, args: Args<'a>, has_paren: bool) -> ArenaDoc<'a> {
         let args = args
@@ -141,3 +202,10 @@ impl<'a> PrettyPrinter<'a> {
         }
     }
 }
+
+fn is_ends_with_hashed_expr(mut children: std::slice::Iter<'_, SyntaxNode>) -> bool {
+    children.next_back().is_some_and(|it| it.is::<Expr>())
+        && children
+            .next_back()
+            .is_some_and(|it| it.kind() == SyntaxKind::Hash)
+}

crates/typstyle-core/src/pretty/import.rs

@@ -27,7 +27,7 @@ impl<'a> PrettyPrinter<'a> {
             };
 
         // Convert the prefix section.
-        let prefix_doc = self.convert_flow_like_sliced(prefix_part.iter(), |child| {
+        let prefix_doc = self.convert_flow_like_iter(prefix_part.iter(), |child| {
             match child.kind() {
                 SyntaxKind::Colon => FlowItem::tight_spaced(self.arena.text(":")),
                 SyntaxKind::Star => FlowItem::spaced(self.arena.text("*")), // wildcard import