Refactoring how MathRow is laid out into a frame while respecting alignment points (#3460)

crates/typst/src/math/align.rs

@@ -19,7 +19,7 @@ pub(super) struct AlignmentResult {
     pub width: Abs,
 }
 
-/// Determine the position of the alignment points.
+/// Determine the positions of the alignment points, according to the input rows combined.
 pub(super) fn alignments(rows: &[MathRow]) -> AlignmentResult {
     let mut widths = Vec::<Abs>::new();
 

crates/typst/src/math/matrix.rs

@@ -396,7 +396,7 @@ fn layout_vec_body(
         flat.push(ctx.layout_into_row(child, styles.chain(&denom_style))?);
     }
 
-    Ok(stack(ctx, styles, flat, align, gap, 0))
+    Ok(stack(flat, align, gap, 0))
 }
 
 /// Layout the inner contents of a matrix.
@@ -480,8 +480,7 @@ fn layout_mat_body(
         let mut y = Abs::zero();
 
         for (cell, &(ascent, descent)) in col.into_iter().zip(&heights) {
-            let cell =
-                cell.into_aligned_frame(ctx, styles, &points, FixedAlignment::Center);
+            let cell = cell.into_line_frame(&points, FixedAlignment::Center);
             let pos = Point::new(
                 if points.is_empty() { x + (rcol - cell.width()) / 2.0 } else { x },
                 y + ascent - cell.ascent(),

crates/typst/src/math/row.rs

@@ -5,8 +5,8 @@ use unicode_math_class::MathClass;
 use crate::foundations::{Resolve, StyleChain};
 use crate::layout::{Abs, AlignElem, Em, FixedAlignment, Frame, FrameKind, Point, Size};
 use crate::math::{
-    alignments, scaled_font_size, spacing, AlignmentResult, EquationElem, FrameFragment,
-    MathContext, MathFragment, MathParItem, MathSize,
+    alignments, scaled_font_size, spacing, EquationElem, FrameFragment, MathContext,
+    MathFragment, MathParItem, MathSize,
 };
 use crate::model::ParElem;
 
@@ -142,7 +142,11 @@ impl MathRow {
 
     pub fn into_frame(self, ctx: &MathContext, styles: StyleChain) -> Frame {
         let align = AlignElem::alignment_in(styles).resolve(styles).x;
-        self.into_aligned_frame(ctx, styles, &[], align)
+        if !self.is_multiline() {
+            self.into_line_frame(&[], align)
+        } else {
+            self.multiline_frame_builder(ctx, styles, align).build()
+        }
     }
 
     pub fn into_fragment(self, ctx: &MathContext, styles: StyleChain) -> MathFragment {
@@ -153,16 +157,17 @@ impl MathRow {
         }
     }
 
-    pub fn into_aligned_frame(
+    /// Returns a builder that lays out `MathFragment`s into a multi-row frame. The set
+    /// of alignment points are computed from those rows combined.
+    pub fn multiline_frame_builder(
         self,
         ctx: &MathContext,
         styles: StyleChain,
-        points: &[Abs],
         align: FixedAlignment,
-    ) -> Frame {
-        if !self.iter().any(|frag| matches!(frag, MathFragment::Linebreak)) {
-            return self.into_line_frame(points, align);
-        }
+    ) -> MathRowFrameBuilder {
+        let rows: Vec<_> = self.rows();
+        let row_count = rows.len();
+        let alignments = alignments(&rows);
 
         let leading = if EquationElem::size_in(styles) >= MathSize::Text {
             ParElem::leading_in(styles)
@@ -171,35 +176,32 @@ impl MathRow {
             TIGHT_LEADING.at(font_size)
         };
 
-        let mut rows: Vec<_> = self.rows();
-
-        if matches!(rows.last(), Some(row) if row.0.is_empty()) {
-            rows.pop();
-        }
-
-        let AlignmentResult { points, width } = alignments(&rows);
-        let mut frame = Frame::soft(Size::zero());
-
+        let mut frames: Vec<(Frame, Point)> = vec![];
+        let mut size = Size::zero();
         for (i, row) in rows.into_iter().enumerate() {
-            let sub = row.into_line_frame(&points, align);
-            let size = frame.size_mut();
+            if i == row_count - 1 && row.0.is_empty() {
+                continue;
+            }
+
+            let sub = row.into_line_frame(&alignments.points, align);
             if i > 0 {
                 size.y += leading;
             }
 
             let mut pos = Point::with_y(size.y);
-            if points.is_empty() {
-                pos.x = align.position(width - sub.width());
+            if alignments.points.is_empty() {
+                pos.x = align.position(alignments.width - sub.width());
             }
-            size.y += sub.height();
             size.x.set_max(sub.width());
-            frame.push_frame(pos, sub);
+            size.y += sub.height();
+            frames.push((sub, pos));
         }
 
-        frame
+        MathRowFrameBuilder { size, frames }
     }
 
-    fn into_line_frame(self, points: &[Abs], align: FixedAlignment) -> Frame {
+    /// Lay out `MathFragment`s into a one-row frame, with alignment points respected.
+    pub fn into_line_frame(self, points: &[Abs], align: FixedAlignment) -> Frame {
         let ascent = self.ascent();
         let mut frame = Frame::soft(Size::new(Abs::zero(), ascent + self.descent()));
         frame.set_baseline(ascent);
@@ -339,6 +341,10 @@ impl MathRow {
 
         items
     }
+
+    fn is_multiline(&self) -> bool {
+        self.iter().any(|frag| matches!(frag, MathFragment::Linebreak))
+    }
 }
 
 impl<T: Into<MathFragment>> From<T> for MathRow {
@@ -365,3 +371,23 @@ impl Iterator for LeftRightAlternator {
         r
     }
 }
+
+/// How the rows should be aligned and merged into a Frame.
+pub struct MathRowFrameBuilder {
+    /// The size of the resulting frame.
+    size: Size,
+    /// Sub frames, and the positions where they should be pushed into
+    /// the resulting frame.
+    frames: Vec<(Frame, Point)>,
+}
+
+impl MathRowFrameBuilder {
+    /// Consumes the builder and returns a `Frame`.
+    pub fn build(self) -> Frame {
+        let mut frame = Frame::soft(self.size);
+        for (sub, pos) in self.frames.into_iter() {
+            frame.push_frame(pos, sub);
+        }
+        frame
+    }
+}

crates/typst/src/math/underover.rs

@@ -290,7 +290,7 @@ fn layout_underoverspreader(
         baseline = rows.len() - 1;
     }
 
-    let frame = stack(ctx, styles, rows, FixedAlignment::Center, gap, baseline);
+    let frame = stack(rows, FixedAlignment::Center, gap, baseline);
     ctx.push(FrameFragment::new(ctx, styles, frame).with_class(body_class));
 
     Ok(())
@@ -301,8 +301,6 @@ fn layout_underoverspreader(
 /// Add a `gap` between each row and uses the baseline of the `baseline`th
 /// row for the whole frame.
 pub(super) fn stack(
-    ctx: &MathContext,
-    styles: StyleChain,
     rows: Vec<MathRow>,
     align: FixedAlignment,
     gap: Abs,
@@ -312,7 +310,7 @@ pub(super) fn stack(
     let AlignmentResult { points, width } = alignments(&rows);
     let rows: Vec<_> = rows
         .into_iter()
-        .map(|row| row.into_aligned_frame(ctx, styles, &points, align))
+        .map(|row| row.into_line_frame(&points, align))
         .collect();
 
     let mut y = Abs::zero();