address review: one 2^20-cell padding budget for all the matrices of an equation

crates/docx/src/write/math.rs

@@ -11,6 +11,7 @@ use crate::xml::W;
 /// is written back unchanged when the equation came from a file. `color` (the equation's run
 /// colour) goes on every math run without its own: OMML has no equation-level colour (#1467).
 pub fn write_equation(w: &mut W, linear: &str, display: bool, math: &Math, color: Option<Rgb>) {
+    w.padded = 0;
     let tinted;
     let math = if color == omml_color(&math.omml) {
         math
@@ -179,9 +180,22 @@ fn char_prop(w: &mut W, name: &str, c: Option<char>) {
     val(w, name, &s);
 }
 
+/// Most cells the matrices of one equation are padded out to in all (rows × the widest row, of
+/// each). A ragged matrix past it (a native document's, or one parsed from a few bytes of text)
+/// writes each row as it is, rather than output that grows with the square of its size, and so
+/// do nested matrices that would each fit alone.
+const MAX_PADDED_CELLS: usize = 1 << 20;
+
 fn write_matrix(w: &mut W, n: &MNode, depth: usize) {
     let MNode::Matrix { rows, col_jc } = n else { return };
     let cols = rows.iter().map(|r| r.len()).max().unwrap_or(0).max(1);
+    let cells = rows.len().saturating_mul(cols);
+    let pad = w.padded.saturating_add(cells) <= MAX_PADDED_CELLS;
+    if pad {
+        w.padded += cells;
+    } else {
+        log::warn!("docx: a {} × {cols} matrix is too large to pad; its rows are written as they are", rows.len());
+    }
     w.open("m:m", &[]);
     w.open("m:mPr", &[]);
     w.open("m:mcs", &[]);
@@ -203,7 +217,7 @@ fn write_matrix(w: &mut W, n: &MNode, depth: usize) {
     w.close("m:mPr");
     for row in rows {
         w.open("m:mr", &[]);
-        for c in 0..cols {
+        for c in 0..if pad { cols } else { row.len().max(1) } {
             arg(w, "m:e", row.get(c).map(|v| v.as_slice()).unwrap_or(&[]), depth);
         }
         w.close("m:mr");
@@ -480,3 +494,40 @@ fn run(w: &mut W, r: &MRun) {
     w.leaf("m:t", &[("xml:space", "preserve")], &r.text);
     w.close("m:r");
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn cells(rows: Vec<Vec<Arg>>) -> usize {
+        let mut w = W::new();
+        write_matrix(&mut w, &MNode::Matrix { rows, col_jc: Vec::new() }, 0);
+        w.s.matches("<m:e/>").count() + w.s.matches("<m:e>").count()
+    }
+
+    #[test]
+    fn a_ragged_matrix_is_padded_within_a_cell_budget() {
+        // A small ragged matrix is padded to a rectangle, as before.
+        assert_eq!(cells(vec![vec![Vec::new(); 3], vec![Vec::new()]]), 6);
+        // 2,000 one-cell rows and one 2,000-cell row: 4 million cells padded, past the budget.
+        let mut rows = vec![vec![Vec::new()]; 2000];
+        rows.push(vec![Vec::new(); 2000]);
+        assert_eq!(cells(rows), 4000);
+    }
+
+    #[test]
+    fn matrices_of_one_equation_share_the_padding_budget() {
+        // Each is 700 × 700 padded (490,000 cells): two fit the budget of 2^20, a third does not.
+        let ragged = || {
+            let mut rows = vec![vec![Vec::new()]; 699];
+            rows.push(vec![Vec::new(); 700]);
+            MNode::Matrix { rows, col_jc: Vec::new() }
+        };
+        let mut w = W::new();
+        for _ in 0..3 {
+            write_matrix(&mut w, &ragged(), 0);
+        }
+        let n = w.s.matches("<m:e/>").count() + w.s.matches("<m:e>").count();
+        assert_eq!(n, 2 * 490_000 + 1399, "the third matrix is written as it is");
+    }
+}

crates/docx/src/xml.rs

@@ -513,11 +513,13 @@ pub fn esc(s: &str) -> String {
 #[derive(Default)]
 pub struct W {
     pub s: String,
+    /// Matrix cells padded so far in the equation being written (reset by `write_equation`).
+    pub padded: usize,
 }
 
 impl W {
     pub fn new() -> W {
-        W { s: String::from("<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>\n") }
+        W { s: String::from("<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>\n"), padded: 0 }
     }
     fn attrs(&mut self, attrs: &[(&str, &str)]) {
         for (k, v) in attrs {