Implement Math Block Align

lib/src/lib.rs

@@ -12,8 +12,11 @@ use typst_syntax::{parse, LinkedNode};
 use Option::None;
 
 mod config;
+
 pub use config::Config;
+
 mod context;
+
 use context::Ctx;
 
 mod utils;
@@ -22,6 +25,7 @@ mod binary;
 mod code_blocks;
 mod markup;
 mod params;
+mod math;
 
 #[must_use]
 pub fn format(s: &str, config: Config) -> String {
@@ -43,7 +47,7 @@ pub fn format(s: &str, config: Config) -> String {
 /// how they will be formatted.
 ///
 /// One assumed rule is that no kind should be formatting with surrounded space
-#[instrument(skip_all,name = "V", fields(kind = format!("{:?}",node.kind())))]
+#[instrument(skip_all, name = "V", fields(kind = format!("{:?}",node.kind())))]
 fn visit(node: &LinkedNode, ctx: &mut Ctx) -> String {
     let mut res: Vec<String> = vec![];
     for child in node.children() {
@@ -68,6 +72,10 @@ fn visit(node: &LinkedNode, ctx: &mut Ctx) -> String {
             ctx.lost_context();
             node.text().to_string()
         }
+        Equation => {
+            math::format_equation(node, &res, ctx)
+        }
+        Math => math::format_math(node, &res, ctx),
         _ => format_default(node, &res, ctx),
     };
     if node.children().count() == 0 {

lib/src/math.rs

@@ -0,0 +1,146 @@
+use super::*;
+use crate::context::Ctx;
+use crate::format_comment_handling_disable;
+use std::cmp::max;
+use tracing::info;
+use typst_syntax::ast::*;
+
+#[instrument(skip_all)]
+pub(crate) fn format_equation(parent: &LinkedNode, children: &[String], ctx: &mut Ctx) -> String {
+    let mut res = String::new();
+    let first_space = parent.children().nth(1);
+    let space_type = if first_space
+        .as_ref()
+        .is_some_and(|s| s.text().contains('\n'))
+    {
+        "\n"
+    } else {
+        " "
+    };
+
+    let mut first_dollar = true;
+
+    let newline = space_type == "\n";
+
+    for (s, node) in children.iter().zip(parent.children()) {
+        match node.kind() {
+            _ if ctx.off => res.push_str(node.text()),
+            LineComment | BlockComment => {
+                let buf = format_comment_handling_disable(&node, &[], ctx);
+                ctx.push_raw_in(&buf, &mut res);
+            }
+            Dollar if first_dollar => {
+                first_dollar = false;
+                ctx.push_raw_in(s, &mut res);
+            }
+            Math => {
+                if newline {
+                    ctx.push_raw_in(ctx.get_indent().as_str(), &mut res);
+                    ctx.push_raw_indent(s, &mut res);
+                } else {
+                    ctx.push_raw_in(s, &mut res);
+                }
+            }
+            Space => {
+                ctx.push_raw_in(space_type, &mut res);
+            }
+            _ => ctx.push_raw_indent(s, &mut res),
+        }
+    }
+    res
+}
+
+#[instrument(skip_all)]
+pub(crate) fn format_math(parent: &LinkedNode, children: &[String], ctx: &mut Ctx) -> String {
+    let mut res = String::new();
+
+    let mut align_points: Vec<usize> = retrieve_align_point(parent, children);
+    let mut index = 0;
+    let mut position = 0usize;
+
+    let mut first_align = true;
+    let mut should_indent = false;
+
+    for (s, node) in children.iter().zip(parent.children()) {
+        match node.kind() {
+            _ if ctx.off => res.push_str(node.text()),
+            MathAlignPoint => {
+                debug_assert!(
+                    align_points[index] >= position,
+                    "align point {} is smaller than position {}",
+                    align_points[index],
+                    position
+                );
+
+                if position == 0 && first_align {
+                    should_indent = true;
+                }
+
+                if position == 0 && should_indent {
+                    ctx.push_raw_in(ctx.get_indent().as_str(), &mut res);
+                }
+
+                ctx.push_raw_in(
+                    " ".repeat(align_points[index] - position).as_str(),
+                    &mut res,
+                );
+                ctx.push_raw_in(s, &mut res);
+                position = align_points[index] + s.len();
+                index += 1;
+
+                first_align = false;
+            }
+            Space if s.contains('\n') => {
+                position = 0;
+                index = 0;
+                ctx.push_raw_in(s, &mut res);
+            }
+            Space => {
+                position += 1;
+                ctx.push_raw_in(" ", &mut res);
+            }
+            _ => {
+                position += s.len();
+                ctx.push_raw_in(s, &mut res)
+            }
+        }
+    }
+
+    res
+}
+
+fn retrieve_align_point(parent: &LinkedNode, children: &[String]) -> Vec<usize> {
+    let mut align_points: Vec<usize> = vec![];
+    let mut index = 0;
+    let mut position = 0usize;
+
+    let mut space = false;
+
+    for (s, node) in children.iter().zip(parent.children()) {
+        match node.kind() {
+            MathAlignPoint => {
+                if align_points.len() <= index {
+                    align_points.push(position);
+
+                    position += s.len();
+                } else {
+                    align_points[index] = max(align_points[index], position);
+                }
+                index += 1
+            }
+            Space if s.contains('\n') => {
+                position = 0;
+                index = 0;
+            }
+            Space => {
+                space = true;
+                position += 1;
+            }
+            _ => {
+                space = false;
+                position += s.len();
+            }
+        }
+    }
+    align_points
+}

lib/src/tests/math.rs

@@ -0,0 +1,61 @@
+use super::*;
+
+make_test!(mathblock1,
+r#"$
+    #xx(a,b) &= 1 \
+    &= 2 \
+
+$"#);
+
+
+make_test!(mathblock2,
+r#"$x$"#);
+
+
+make_test!(mathblock3,
+r#"$
+    #xx(a,b) &= 1 \
+    &= 2 \
+    &=3 \
+    &=4 \
+
+$"#);
+
+make_test!(mathblock4,
+r#"$
+    #xx(a,b) &= 1 &=3 \
+    &= 2 &=2 \
+    &=3 \
+    &=4 \
+    &=5 \
+    $"#);
+
+make_test!(mathblock5,
+r#"$
+#xx(a,b)
+    &= 1 &=3 \
+    &= 2 &=2 \
+    &=3 \
+    &=4 \
+    &=5 \
+    $"#);
+
+make_test!(mathblock6,
+r#"$
+    #xx(a,b) &= 1 &=3 \
+                    &= 2 &=2 \
+    &=3 \
+    &=4 \
+    &=5 \
+    $"#);
+
+make_test!(mathblock7,
+r#"$
+mat(
+    1,2,3,4;
+5,2,2,3
+)
+$"#);
+
+make_test!(mathblock8,
+r#"$#xx(a,b) &= 1 \ #xx(a,b) &= 1 \ &= 2 \ $"#);

lib/src/tests/mod.rs

@@ -170,3 +170,4 @@ mod lists;
 mod markup;
 mod params;
 mod snippets;
+mod math;

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock1__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\n    #xx(a,b) &= 1 \\\\\\n    &= 2 \\\\\\n\\n$\"\n===\n$\n    #xx(a,b) &= 1 \\\n    &= 2 \\\n\n$\n===\nFORMATTED\n===\n$\n  #xx(a, b) &= 1 \\\n            &= 2 \\\n$"
+expression: formatted
+---
+"$\n  #xx(a, b) &= 1 \\\n            &= 2 \\\n$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock2__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$x$\"\n===\n$x$\n===\nFORMATTED\n===\n$x$"
+expression: formatted
+---
+"$x$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock3__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\n    #xx(a,b) &= 1 \\\\\\n    &= 2 \\\\\\n    &=3 \\\\\\n    &=4 \\\\\\n\\n$\"\n===\n$\n    #xx(a,b) &= 1 \\\n    &= 2 \\\n    &=3 \\\n    &=4 \\\n\n$\n===\nFORMATTED\n===\n$\n  #xx(a, b) &= 1 \\\n            &= 2 \\\n            &=3 \\\n            &=4 \\\n$"
+expression: formatted
+---
+"$\n  #xx(a, b) &= 1 \\\n            &= 2 \\\n            &=3 \\\n            &=4 \\\n$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock4__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\n    #xx(a,b) &= 1 &=3 \\\\\\n    &= 2 &=2 \\\\\\n    &=3 \\\\\\n    &=4 \\\\\\n    &=5 \\\\\\n    $\"\n===\n$\n    #xx(a,b) &= 1 &=3 \\\n    &= 2 &=2 \\\n    &=3 \\\n    &=4 \\\n    &=5 \\\n    $\n===\nFORMATTED\n===\n$\n  #xx(a, b) &= 1 &=3 \\\n            &= 2 &=2 \\\n            &=3 \\\n            &=4 \\\n            &=5 \\\n$"
+expression: formatted
+---
+"$\n  #xx(a, b) &= 1 &=3 \\\n            &= 2 &=2 \\\n            &=3 \\\n            &=4 \\\n            &=5 \\\n$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock5__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\n#xx(a,b)\\n    &= 1 &=3 \\\\\\n    &= 2 &=2 \\\\\\n    &=3 \\\\\\n    &=4 \\\\\\n    &=5 \\\\\\n    $\"\n===\n$\n#xx(a,b)\n    &= 1 &=3 \\\n    &= 2 &=2 \\\n    &=3 \\\n    &=4 \\\n    &=5 \\\n    $\n===\nFORMATTED\n===\n$\n  #xx(a, b)\n    &= 1 &=3 \\\n    &= 2 &=2 \\\n    &=3 \\\n    &=4 \\\n    &=5 \\\n$"
+expression: formatted
+---
+"$\n  #xx(a, b)\n    &= 1 &=3 \\\n    &= 2 &=2 \\\n    &=3 \\\n    &=4 \\\n    &=5 \\\n$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock6__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\n    #xx(a,b) &= 1 &=3 \\\\\\n                    &= 2 &=2 \\\\\\n    &=3 \\\\\\n    &=4 \\\\\\n    &=5 \\\\\\n    $\"\n===\n$\n    #xx(a,b) &= 1 &=3 \\\n                    &= 2 &=2 \\\n    &=3 \\\n    &=4 \\\n    &=5 \\\n    $\n===\nFORMATTED\n===\n$\n  #xx(a, b) &= 1 &=3 \\\n            &= 2 &=2 \\\n            &=3 \\\n            &=4 \\\n            &=5 \\\n$"
+expression: formatted
+---
+"$\n  #xx(a, b) &= 1 &=3 \\\n            &= 2 &=2 \\\n            &=3 \\\n            &=4 \\\n            &=5 \\\n$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock7__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\nmat(\\n    1,2,3,4;\\n5,2,2,3\\n)\\n$\"\n===\n$\nmat(\n    1,2,3,4;\n5,2,2,3\n)\n$\n===\nFORMATTED\n===\n$\n  mat(1, 2, 3, 4;5, 2, 2, 3)\n$"
+expression: formatted
+---
+"$\n  mat(1, 2, 3, 4;5, 2, 2, 3)\n$"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock8__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$#xx(a,b) &= 1 \\\\ #xx(a,b) &= 1 \\\\ &= 2 \\\\ $\"\n===\n$#xx(a,b) &= 1 \\ #xx(a,b) &= 1 \\ &= 2 \\ $\n===\nFORMATTED\n===\n$#xx(a, b) &= 1 \\ #xx(a, b) &= 1 \\ &= 2 \\ $"
+expression: formatted
+---
+"$#xx(a, b) &= 1 \\ #xx(a, b) &= 1 \\ &= 2 \\ $"

lib/src/tests/snapshots/typstfmt_lib__tests__math__mathblock__snapshot.snap

@@ -0,0 +1,6 @@
+---
+source: lib/src/tests/math.rs
+description: "INPUT\n===\n\"$\\n    #xx(a,b) &= 1 \\\\\\n    &= 2 \\\\\\n\\n$\"\n===\n$\n    #xx(a,b) &= 1 \\\n    &= 2 \\\n\n$\n===\nFORMATTED\n===\n$\n  #xx(a, b) &= 1 \\\n            &= 2 \\\n$"
+expression: formatted
+---
+"$\n  #xx(a, b) &= 1 \\\n            &= 2 \\\n$"