Merge pull request #5470 from vangel76/perf/trace-parallel-component-fit

Image Trace: fit the components of a layer in parallel, in chunks (#5464)

crates/trace/src/layers.rs

@@ -331,7 +331,9 @@ impl Layers<'_> {
                 }
             }
         }
-        for comp in comps {
+        // One component's paths: a pure function of its loops and the shared read-only state, so components are fitted in
+        // parallel and then added in order.
+        let fit_comp = |comp: &contour::Component| -> Option<(PathData, Option<TracedShape>)> {
             // (`hard`: some of the outline has no anti-aliasing, and is known less exactly.)
             let hard = std::cell::Cell::new(false);
             // (Each loop fitted, and whether its own outline is hard.)
@@ -355,7 +357,7 @@ impl Layers<'_> {
                 hard.set(hard.get() || is_hard);
                 Some((sp, is_hard))
             };
-            let Some((outer, _)) = fit(&comp.outer) else { continue };
+            let (outer, _) = fit(&comp.outer)?;
             let mut subs = vec![outer];
             for hole in &comp.holes {
                 if hole.area2.abs() / 2 < min_hole && params.noise > 0 {
@@ -398,9 +400,16 @@ impl Layers<'_> {
                 _ => None,
             };
             let path = shape.map_or_else(|| PathData::new(subs), |s| s.path());
-            if !add(path, comp.pixels, area_stroke, shape, false) {
-                return None;
-            }
+            Some((path, shape))
+        };
+        // In chunks, so that the anchor budget still stops the work early: a hostile image is refused after fitting at most
+        // the rest of one chunk, not after every component is fitted.
+        let accepted = fit_in_chunks(&comps, FIT_CHUNK, fit_comp, |comp, fitted| {
+            let Some((path, shape)) = fitted else { return true };
+            add(path, comp.pixels, area_stroke, shape, false)
+        });
+        if !accepted {
+            return None;
         }
         for line in lines {
             if !add(line.path, line.pixels, Some(line.width), None, line.square) {
@@ -452,6 +461,24 @@ pub(crate) fn shift_probe(img: &Raster, params: &TraceParams) -> Vec<[f64; 4]> {
     out
 }
 
+/// How many components of a layer are fitted at a time (see `fit_in_chunks`).
+const FIT_CHUNK: usize = 256;
+
+/// `fit` of each of `items` over the available cores, `chunk` items at a time, each result handed to `accept` in item
+/// order. Stops after the chunk in which `accept` returns false, and says whether it never did.
+fn fit_in_chunks<T: Sync, R: Send>(items: &[T], chunk: usize, fit: impl Fn(&T) -> R + Sync, mut accept: impl FnMut(&T, R) -> bool) -> bool {
+    for part in items.chunks(chunk.max(1)) {
+        // (A few components are not worth starting threads for.)
+        let fitted = if part.len() < 16 { part.iter().map(&fit).collect() } else { parallel_map(part, &fit) };
+        for (item, result) in part.iter().zip(fitted) {
+            if !accept(item, result) {
+                return false;
+            }
+        }
+    }
+    true
+}
+
 /// `f` of each of `items`, in order, spread over the available cores.
 pub(crate) fn parallel_map<T: Sync, R: Send>(items: &[T], f: impl Fn(&T) -> R + Sync) -> Vec<R> {
     parallel_map_at_most(usize::MAX, items, f)
@@ -492,3 +519,50 @@ pub(crate) fn parallel_map_at_most<T: Sync, R: Send>(at_a_time: usize, items: &[
 pub(crate) fn parallel_map_at_most<T: Sync, R: Send>(_: usize, items: &[T], f: impl Fn(&T) -> R + Sync) -> Vec<R> {
     items.iter().map(f).collect()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn fit_in_chunks_keeps_the_order_and_stops_at_the_first_refusal() {
+        let items: Vec<usize> = (0..1000).collect();
+        let fitted = AtomicUsize::new(0);
+        let mut seen = vec![];
+        let all = fit_in_chunks(
+            &items,
+            16,
+            |&i| {
+                fitted.fetch_add(1, Ordering::Relaxed);
+                i * 2
+            },
+            |&i, r| {
+                assert_eq!(r, i * 2);
+                seen.push(i);
+                seen.len() <= 10
+            },
+        );
+        assert!(!all);
+        // Items come in order, up to and including the one refused.
+        assert_eq!(seen, (0..=10).collect::<Vec<_>>());
+        // The early stop is kept: only the chunk holding the refusal was fitted, not the other 984 items.
+        assert_eq!(fitted.load(Ordering::Relaxed), 16);
+    }
+
+    #[test]
+    fn fit_in_chunks_accepts_everything_in_order() {
+        let items: Vec<usize> = (0..1000).collect();
+        let mut out = vec![];
+        assert!(fit_in_chunks(
+            &items,
+            7,
+            |&i| i + 1,
+            |_, r| {
+                out.push(r);
+                true
+            }
+        ));
+        assert_eq!(out, (1..=1000).collect::<Vec<_>>());
+        assert!(fit_in_chunks(&[] as &[usize], 7, |&i| i, |_, _| false));
+    }
+}

crates/trace/src/tests.rs

@@ -609,6 +609,32 @@ fn a_trace_past_its_anchor_budget_is_refused() {
     assert_eq!(trace_within(&img, &p, anchors).unwrap(), full);
 }
 
+/// Components of a layer are fitted in parallel, a chunk at a time, and added in the order they were found (top to
+/// bottom, left to right), also across the chunks. The anchor budget still refuses a trace with many components.
+#[test]
+fn many_components_come_out_in_scan_order_and_still_meet_the_budget() {
+    // 20 x 20 black dots, more than one chunk of components.
+    let img = Raster::from_fn(120, 120, |x, y| if x % 6 < 3 && y % 6 < 3 { BLACK } else { WHITE });
+    let p = TraceParams { noise: 1, ..bw() };
+    let full = trace(&img, &p);
+    assert_eq!(full.paths.len(), 400);
+    let tops: Vec<(i64, i64)> = full
+        .paths
+        .iter()
+        .map(|t| {
+            let b = t.path.bounds().unwrap();
+            (b.y0.round() as i64, b.x0.round() as i64)
+        })
+        .collect();
+    assert!(tops.windows(2).all(|w| w[0] < w[1]), "not in scan order: {tops:?}");
+    assert_eq!(trace(&img, &p), full, "tracing twice differs");
+    let anchors = full.anchor_count();
+    assert!(matches!(trace_within(&img, &p, anchors / 2), Err(TraceError::TooComplex { max }) if max == anchors / 2));
+    // The last component alone is over the budget: a refusal in a later chunk.
+    assert!(matches!(trace_within(&img, &p, anchors - 1), Err(TraceError::TooComplex { .. })));
+    assert_eq!(trace_within(&img, &p, anchors).unwrap(), full);
+}
+
 /// #3098: a watched trace says how far it has got, layer by layer, and stops when told to.
 #[test]
 fn a_watched_trace_counts_its_layers_and_stops_when_told() {