Merge pull request #5412 from storytold/fix/sec-trace-raster-bounds

Trace: empty and mismatched rasters, and large swatch libraries, take bounded work

crates/trace/src/mosaic.rs

@@ -7,13 +7,14 @@ const MAX_SIDE: u32 = 4096;
 
 /// Average colour (straight RGBA) of each of `cols` × `rows` tiles covering `img`, row by row.
 /// Transparent pixels count by their alpha, so a tile over transparency fades out. Tiles over the
-/// same pixels (more tiles than pixels) are averaged once, and a raster whose bytes are not
-/// `width * height * 4` gives transparent tiles.
+/// same pixels (more tiles than pixels) are averaged once, and an empty raster or one whose bytes
+/// are not `width * height * 4` gives transparent tiles.
 pub fn mosaic(img: &Raster, cols: u32, rows: u32) -> Vec<[u8; 4]> {
     let (cols, rows) = (cols.clamp(1, MAX_SIDE), rows.clamp(1, MAX_SIDE));
     let tiles = cols as usize * rows as usize;
     let (w, h) = (img.width as usize, img.height as usize);
-    if w.checked_mul(h).and_then(|n| n.checked_mul(4)) != Some(img.rgba.len()) {
+    // An empty raster has nothing to average (its rows would still be walked, up to 2^32 of them).
+    if w == 0 || h == 0 || w.checked_mul(h).and_then(|n| n.checked_mul(4)) != Some(img.rgba.len()) {
         return vec![[0; 4]; tiles];
     }
     let mut out = Vec::with_capacity(tiles);
@@ -89,6 +90,9 @@ mod tests {
         // Bytes that are not the raster's size: transparent.
         let short = Raster { width: 4, height: 4, rgba: vec![] };
         assert_eq!(mosaic(&short, 2, 2), vec![[0; 4]; 4]);
+        // No pixels at all, however tall: transparent, with no rows walked.
+        let empty = Raster { width: 0, height: u32::MAX, rgba: vec![] };
+        assert_eq!(mosaic(&empty, 2, 2), vec![[0; 4]; 4]);
         // Transparency.
         let clear = Raster::from_fn(2, 2, |_, _| [9, 9, 9, 0]);
         assert_eq!(mosaic(&clear, 1, 1), vec![[0, 0, 0, 0]]);

crates/trace/src/quantize.rs

@@ -34,8 +34,13 @@ pub(crate) fn luma(p: [u8; 4]) -> u8 {
     ((p[0] as u32 * 299 + p[1] as u32 * 587 + p[2] as u32 * 114 + 500) / 1000) as u8
 }
 
-/// Reduce `img` to a palette according to `params.mode`.
+/// Reduce `img` to a palette according to `params.mode`. A raster whose bytes are not
+/// `width * height * 4` (its fields are public) has no pixels to reduce: no labels, no palette.
 pub fn quantize(img: &Raster, params: &TraceParams) -> Quantized {
+    let want = (img.width as usize).checked_mul(img.height as usize).and_then(|n| n.checked_mul(4));
+    if want != Some(img.rgba.len()) {
+        return Quantized { labels: vec![], palette: vec![] };
+    }
     quantize_keeping(img, params).0
 }
 

crates/trace/src/quantize/palette.rs

@@ -202,13 +202,9 @@ const FLAT_BINS: usize = 96;
 /// Document Library: the `k` swatches (at most) that the most pixels are nearest to (see
 /// [`nearest_library`]).
 fn library_colors(used: &[Used], swatches: &[[u8; 3]], k: usize) -> Vec<[f64; 3]> {
-    let mut all: Vec<[f64; 3]> = vec![];
-    for s in swatches {
-        let c = s.map(f64::from);
-        if !all.contains(&c) {
-            all.push(c);
-        }
-    }
+    // Each colour once, in their order: looked up in a set, not in the colours kept so far.
+    let mut seen = std::collections::HashSet::new();
+    let all: Vec<[f64; 3]> = swatches.iter().filter(|s| seen.insert(**s)).map(|s| s.map(f64::from)).collect();
     let mut weight = vec![0.0; all.len()];
     let all_labs = labs(&all);
     for (_, c, n) in used {

crates/trace/src/tests.rs

@@ -428,6 +428,29 @@ fn a_cut_out_edge_leaves_the_color_of_the_art_inside_it_alone() {
     assert_eq!(q.palette.get(q.labels[12] as usize), Some(&[255, 255, 255]), "{q:?}");
 }
 
+#[test]
+fn a_library_of_many_repeated_swatches_is_read_once_per_color() {
+    let img = Raster::from_fn(2, 1, |x, _| if x == 0 { [0, 0, 0, 255] } else { [255, 255, 255, 255] });
+    // 200,000 swatches of 50,000 distinct colours: compared against those kept so far, billions of steps.
+    let swatches: Vec<[u8; 3]> = (0..200_000u32)
+        .map(|i| {
+            let c = i % 50_000;
+            [(c >> 16) as u8, (c >> 8) as u8, c as u8]
+        })
+        .collect();
+    let params = TraceParams { mode: Mode::Color, palette: Palette::DocumentLibrary, swatches, colors: 2, ..TraceParams::default() };
+    let start = std::time::Instant::now();
+    let q = quantize(&img, &params);
+    assert!(start.elapsed() < std::time::Duration::from_secs(30), "{:?}", start.elapsed());
+    assert_eq!(q.labels.len(), 2);
+}
+
+#[test]
+fn a_raster_whose_bytes_are_not_its_size_quantizes_to_nothing() {
+    let short = Raster { width: 60_000, height: 60_000, rgba: vec![0; 4] };
+    assert_eq!(quantize(&short, &TraceParams::default()), Quantized { labels: vec![], palette: vec![] });
+}
+
 #[test]
 fn a_semi_transparent_stripe_beside_a_cut_out_edge_keeps_its_own_color() {
     // Columns: transparent, a half-opaque gray edge (white blended with the background), a