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,155 @@
+use std::cmp::max;
+use tracing::info;
+use super::*;
+use typst_syntax::ast::*;
+use crate::context::Ctx;
+use crate::format_comment_handling_disable;
+
+#[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;
+
+
+    info!("equation children: {:?}", children);
+
+    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;
+
+    info!("math children: {:?}", children);
+
+    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 => {
+                if align_points[index] < position {
+                    panic!("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,50 @@
+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 \
+    $"#);

lib/src/tests/mod.rs

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