Fix broken matrices with alignment and optimize code while we're at it (#935)

4 files changed, 31 insertions, 34 deletions

diff --git a/library/src/math/align.rs b/library/src/math/align.rs
index bbdda5fd..4e4a76e9 100644
--- a/library/src/math/align.rs
+++ b/library/src/math/align.rs
@@ -15,10 +15,16 @@ impl LayoutMath for AlignPointElem {
     }
 }
 
+pub(super) struct AlignmentResult {
+    pub points: Vec<Abs>,
+    pub width: Abs,
+}
+
 /// Determine the position of the alignment points.
-pub(super) fn alignments(rows: &[MathRow]) -> Vec<Abs> {
+pub(super) fn alignments(rows: &[MathRow]) -> AlignmentResult {
     let mut widths = Vec::<Abs>::new();
 
+    let mut pending_width = Abs::zero();
     for row in rows {
         let mut width = Abs::zero();
         let mut alignment_index = 0;
@@ -28,6 +34,7 @@ pub(super) fn alignments(rows: &[MathRow]) -> Vec<Abs> {
                     widths[alignment_index].set_max(width);
                 } else {
                     widths.push(width);
+                    pending_width = Abs::zero();
                 }
                 width = Abs::zero();
                 alignment_index += 1;
@@ -35,6 +42,7 @@ pub(super) fn alignments(rows: &[MathRow]) -> Vec<Abs> {
                 width += fragment.width();
             }
         }
+        pending_width.set_max(width);
     }
 
     let mut points = widths;
@@ -42,5 +50,8 @@ pub(super) fn alignments(rows: &[MathRow]) -> Vec<Abs> {
         let prev = points[i - 1];
         points[i] += prev;
     }
-    points
+    AlignmentResult {
+        width: points.last().copied().unwrap_or_default() + pending_width,
+        points,
+    }
 }
diff --git a/library/src/math/matrix.rs b/library/src/math/matrix.rs
index 956c7f8a..eb465a43 100644
--- a/library/src/math/matrix.rs
+++ b/library/src/math/matrix.rs
@@ -242,16 +242,13 @@ fn layout_mat_body(ctx: &mut MathContext, rows: &[Vec<Content>]) -> SourceResult
         return Ok(Frame::new(Size::zero()));
     }
 
-    let mut widths = vec![Abs::zero(); ncols];
-    let mut ascents = vec![Abs::zero(); nrows];
-    let mut descents = vec![Abs::zero(); nrows];
+    let mut heights = vec![(Abs::zero(), Abs::zero()); nrows];
 
     ctx.style(ctx.style.for_denominator());
     let mut cols = vec![vec![]; ncols];
-    for ((row, ascent), descent) in rows.iter().zip(&mut ascents).zip(&mut descents) {
-        for ((cell, rcol), col) in row.iter().zip(&mut widths).zip(&mut cols) {
+    for (row, (ascent, descent)) in rows.iter().zip(&mut heights) {
+        for (cell, col) in row.iter().zip(&mut cols) {
             let cell = ctx.layout_row(cell)?;
-            rcol.set_max(cell.width());
             ascent.set_max(cell.ascent());
             descent.set_max(cell.descent());
             col.push(cell);
@@ -259,18 +256,15 @@ fn layout_mat_body(ctx: &mut MathContext, rows: &[Vec<Content>]) -> SourceResult
     }
     ctx.unstyle();
 
-    let width = widths.iter().sum::<Abs>() + col_gap * (ncols - 1) as f64;
-    let height = ascents.iter().sum::<Abs>()
-        + descents.iter().sum::<Abs>()
-        + row_gap * (nrows - 1) as f64;
-    let size = Size::new(width, height);
-
-    let mut frame = Frame::new(size);
+    let mut frame = Frame::new(Size::new(
+        Abs::zero(),
+        heights.iter().map(|&(a, b)| a + b).sum::<Abs>() + row_gap * (nrows - 1) as f64,
+    ));
     let mut x = Abs::zero();
-    for (col, &rcol) in cols.into_iter().zip(&widths) {
-        let points = alignments(&col);
+    for col in cols {
+        let AlignmentResult { points, width: rcol } = alignments(&col);
         let mut y = Abs::zero();
-        for ((cell, &ascent), &descent) in col.into_iter().zip(&ascents).zip(&descents) {
+        for (cell, &(ascent, descent)) in col.into_iter().zip(&heights) {
             let cell = cell.into_aligned_frame(ctx, &points, Align::Center);
             let pos = Point::new(
                 if points.is_empty() { x + (rcol - cell.width()) / 2.0 } else { x },
@@ -281,6 +275,7 @@ fn layout_mat_body(ctx: &mut MathContext, rows: &[Vec<Content>]) -> SourceResult
         }
         x += rcol + col_gap;
     }
+    frame.size_mut().x = x - col_gap;
 
     Ok(frame)
 }
diff --git a/library/src/math/row.rs b/library/src/math/row.rs
index aa6f76f6..23c9d242 100644
--- a/library/src/math/row.rs
+++ b/library/src/math/row.rs
@@ -87,10 +87,6 @@ impl MathRow {
         self.iter().map(MathFragment::width).sum()
     }
 
-    pub fn height(&self) -> Abs {
-        self.ascent() + self.descent()
-    }
-
     pub fn ascent(&self) -> Abs {
         self.iter().map(MathFragment::ascent).max().unwrap_or_default()
     }
@@ -136,8 +132,7 @@ impl MathRow {
                 rows.pop();
             }
 
-            let width = rows.iter().map(|row| row.width()).max().unwrap_or_default();
-            let points = alignments(&rows);
+            let AlignmentResult { points, width } = alignments(&rows);
             let mut frame = Frame::new(Size::zero());
 
             for (i, row) in rows.into_iter().enumerate() {
diff --git a/library/src/math/underover.rs b/library/src/math/underover.rs
index 5282c6dc..89ec5cff 100644
--- a/library/src/math/underover.rs
+++ b/library/src/math/underover.rs
@@ -242,22 +242,18 @@ pub(super) fn stack(
     gap: Abs,
     baseline: usize,
 ) -> Frame {
-    let mut width = Abs::zero();
-    let mut height = rows.len().saturating_sub(1) as f64 * gap;
-
-    let points = alignments(&rows);
+    let AlignmentResult { points, width } = alignments(&rows);
     let rows: Vec<_> = rows
         .into_iter()
         .map(|row| row.into_aligned_frame(ctx, &points, align))
         .collect();
 
-    for row in &rows {
-        height += row.height();
-        width.set_max(row.width());
-    }
-
     let mut y = Abs::zero();
-    let mut frame = Frame::new(Size::new(width, height));
+    let mut frame = Frame::new(Size::new(
+        width,
+        rows.iter().map(|row| row.height()).sum::<Abs>()
+            + rows.len().saturating_sub(1) as f64 * gap,
+    ));
 
     for (i, row) in rows.into_iter().enumerate() {
         let x = align.position(width - row.width());