Proximity Match: 1.5-3.8x faster stroke searches, smaller mask tables (#2101)

Co-authored-by: Brandon Thomas <echelon@gmail.com>

crates/algo/src/inpaint.rs

@@ -20,6 +20,9 @@
 
 use crate::poisson::membrane_fill;
 
+mod proximity_integral;
+use proximity_integral::HoleIntegral;
+
 /// Below this many grid cells a level runs single-threaded (thread hand-off costs more than it saves).
 #[cfg(not(target_arch = "wasm32"))]
 const PAR_MIN: usize = 128 * 128;
@@ -529,7 +532,9 @@ pub fn best_offset(w: usize, h: usize, ch: usize, img: &[f32], hole: &[bool], ri
         return None;
     }
     let ring = ring.max(1).min(extent as usize) as i32;
-    let hs = integral(w, h, hole);
+    // Every set cell lies in the hole bounds. Counts outside that rectangle are zero,
+    // so source rejection needs a table for the hole, not for the whole photo.
+    let hs = HoleIntegral::new(w, hole, (bx0, by0, bx1, by1))?;
     // Sample points: the hole's ring (dilated minus hole), subsampled for speed.
     let mut pts = Vec::new();
     let (x0, y0) = (bx0 as i32 - ring, by0 as i32 - ring);
@@ -539,35 +544,43 @@ pub fn best_offset(w: usize, h: usize, ch: usize, img: &[f32], hole: &[bool], ri
             if hole[y as usize * w + x as usize] {
                 continue;
             }
-            if window_count(&hs, w, h, x - ring, y - ring, x + ring + 1, y + ring + 1) > 0 {
+            if hs.count(x - ring, y - ring, x + ring + 1, y + ring + 1) > 0 {
                 pts.push((x, y));
             }
         }
     }
     let stride = (pts.len() / 2000).max(1);
     let pts: Vec<(i32, i32)> = pts.into_iter().step_by(stride).collect();
-    let eval = |dx: i32, dy: i32| -> Option<f32> {
+    let eval = |dx: i32, dy: i32, cutoff: f32| -> Option<f32> {
         let (hx0, hy0, hx1, hy1) = (x0 + dx, y0 + dy, x1 + dx, y1 + dy);
         let inside = hx0 >= 0 && hy0 >= 0 && hx1 <= w as i32 && hy1 <= h as i32;
-        if (dx, dy) == (0, 0) || !inside || window_count(&hs, w, h, hx0, hy0, hx1, hy1) != 0 {
+        if (dx, dy) == (0, 0) || !inside || hs.count(hx0, hy0, hx1, hy1) != 0 {
             return None;
         }
         let mut e = 0.0f32;
-        for &(x, y) in &pts {
+        let n = pts.len().max(1) as f32;
+        let penalty = 1.0 + 1e-3 * ((dx as f32).hypot(dy as f32) / max_radius as f32);
+        for (i, &(x, y)) in pts.iter().enumerate() {
             let (a, b) = ((y as usize * w + x as usize) * ch, ((y + dy) as usize * w + (x + dx) as usize) * ch);
             for c in 0..ch {
                 let d = img[a + c] - img[b + c];
                 e += d * d;
             }
+            // SSD terms are nonnegative. The partial score is a lower bound on the
+            // final score; use the same division and penalty to keep float rounding
+            // and strict tie ordering unchanged. Check in batches to keep this cheap.
+            if i % 32 == 31 && cutoff.is_finite() && e / n * penalty >= cutoff {
+                return None;
+            }
         }
         // Mild preference for nearer sources on ties.
-        Some(e / pts.len().max(1) as f32 * (1.0 + 1e-3 * ((dx as f32).hypot(dy as f32) / max_radius as f32)))
+        Some(e / n * penalty)
     };
     // Coarse grid over the search window, then a full-resolution refinement around the winner.
     let step = (max_radius / 24).max(1);
     let mut best: Option<((i32, i32), f32)> = None;
     let consider = |dx: i32, dy: i32, best: &mut Option<((i32, i32), f32)>| {
-        if let Some(e) = eval(dx, dy)
+        if let Some(e) = eval(dx, dy, best.map_or(f32::INFINITY, |(_, error)| error))
             && best.is_none_or(|(_, be)| e < be)
         {
             *best = Some(((dx, dy), e));

crates/algo/src/inpaint/proximity_integral.rs

@@ -0,0 +1,46 @@
+//! Summed-area table restricted to the nonzero bounds of a healing hole.
+
+pub(super) struct HoleIntegral {
+    x0: i32,
+    y0: i32,
+    w: usize,
+    h: usize,
+    data: Vec<u32>,
+}
+
+impl HoleIntegral {
+    pub(super) fn new(stride: usize, hole: &[bool], (x0, y0, x1, y1): (usize, usize, usize, usize)) -> Option<Self> {
+        let (w, h) = (x1.checked_sub(x0)?, y1.checked_sub(y0)?);
+        let row_stride = w.checked_add(1)?;
+        let len = row_stride.checked_mul(h.checked_add(1)?)?;
+        let mut data = Vec::new();
+        data.try_reserve_exact(len).ok()?;
+        data.resize(len, 0u32);
+        for (y, sy) in (y0..y1).enumerate() {
+            let base = sy.checked_mul(stride)?;
+            let row = hole.get(base.checked_add(x0)?..base.checked_add(x1)?)?;
+            let mut count = 0u32;
+            for (x, &set) in row.iter().enumerate() {
+                count = count.checked_add(u32::from(set))?;
+                let above = *data.get(y * row_stride + x + 1)?;
+                *data.get_mut((y + 1) * row_stride + x + 1)? = above.checked_add(count)?;
+            }
+        }
+        Some(Self { x0: i32::try_from(x0).ok()?, y0: i32::try_from(y0).ok()?, w, h, data })
+    }
+
+    pub(super) fn count(&self, x0: i32, y0: i32, x1: i32, y1: i32) -> u32 {
+        super::window_count(
+            &self.data,
+            self.w,
+            self.h,
+            x0.saturating_sub(self.x0),
+            y0.saturating_sub(self.y0),
+            x1.saturating_sub(self.x0),
+            y1.saturating_sub(self.y0),
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests;

crates/algo/src/inpaint/proximity_integral/tests.rs

@@ -0,0 +1,238 @@
+use super::*;
+use crate::inpaint::{best_offset, hole_bbox, integral, window_count};
+
+#[test]
+fn cropped_counts_match_full_image_for_clipped_and_disjoint_windows() {
+    let (w, h) = (37, 29);
+    for kind in 0..5 {
+        let hole: Vec<_> = (0..w * h)
+            .map(|i| {
+                let (x, y) = (i % w, i / w);
+                match kind {
+                    0 => (11..16).contains(&x) && (8..14).contains(&y),
+                    1 => x < 3 && y < 4,
+                    2 => x >= w - 3 && y >= h - 4,
+                    3 => (x == 3 && y == 8) || (x == 31 && y == 25),
+                    _ => (x * 37 + y * 71) % 17 == 0,
+                }
+            })
+            .collect();
+        let cropped = HoleIntegral::new(w, &hole, hole_bbox(w, h, &hole).unwrap()).unwrap();
+        let full = integral(w, h, &hole);
+        for x0 in [-50, -1, 0, 3, 11, 15, 36, 40] {
+            for x1 in [-2, 0, 3, 13, 16, 37, 50] {
+                for y0 in [-30, -1, 0, 8, 13, 28, 35] {
+                    for y1 in [-1, 0, 10, 14, 29, 40] {
+                        assert_eq!(cropped.count(x0, y0, x1, y1), window_count(&full, w, h, x0, y0, x1, y1), "kind {kind}: {x0},{y0}..{x1},{y1}");
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[test]
+fn table_size_depends_on_the_hole_bounds() {
+    let (w, h) = (1000, 700);
+    let hole: Vec<_> = (0..w * h).map(|i| (503..514).contains(&(i % w)) && (307..320).contains(&(i / w))).collect();
+    let table = HoleIntegral::new(w, &hole, hole_bbox(w, h, &hole).unwrap()).unwrap();
+    assert_eq!((table.w, table.h, table.data.len()), (11, 13, 12 * 14));
+    assert_eq!(table.count(0, 0, 1000, 700), 143);
+    assert_eq!(table.count(i32::MIN, i32::MIN, i32::MAX, i32::MAX), 143);
+    assert!(HoleIntegral::new(w, &hole, (10, 5, 9, 6)).is_none());
+    assert!(HoleIntegral::new(w, &[], (3, 4, 5, 6)).is_none());
+    assert!(HoleIntegral::new(w, &[], (0, 0, usize::MAX, 2)).is_none());
+}
+
+#[test]
+fn proximity_offsets_match_full_tables_at_several_depths_and_channels() {
+    let (w, h) = (61, 49);
+    for ch in [1, 3, 4, 5] {
+        for depth in [8, 16, 32] {
+            let img: Vec<_> = (0..w * h * ch)
+                .map(|i| {
+                    let value = ((i * 101 + i / (w * ch) * 37) % 251) as f32 / 250.0;
+                    match depth {
+                        8 => (value * 255.0).round() / 255.0,
+                        16 => (value * 65535.0).round() / 65535.0,
+                        _ => value * 2.0 - 0.5,
+                    }
+                })
+                .collect();
+            for kind in 0..5 {
+                let hole: Vec<_> = (0..w * h)
+                    .map(|i| {
+                        let (x, y) = (i % w, i / w);
+                        match kind {
+                            0 => (27..32).contains(&x) && (20..25).contains(&y),
+                            1 => x < 2 && y < 3,
+                            2 => (x == 11 && y == 12) || (x == 38 && y == 30),
+                            3 => x == 30 && y == 24,
+                            _ => false,
+                        }
+                    })
+                    .collect();
+                for (ring, radius) in [(1, 5), (3, 17), (5, 30), (1, 49)] {
+                    assert_eq!(
+                        best_offset(w, h, ch, &img, &hole, ring, radius),
+                        old_best_offset(w, h, ch, &img, &hole, ring, radius),
+                        "ch {ch} depth {depth} hole {kind} ring {ring} radius {radius}"
+                    );
+                }
+            }
+        }
+    }
+}
+
+#[test]
+fn ties_nonfinite_samples_and_full_holes_keep_the_previous_result() {
+    let (w, h, ch) = (40, 35, 3);
+    let hole: Vec<_> = (0..w * h).map(|i| (18..22).contains(&(i % w)) && (16..20).contains(&(i / w))).collect();
+    for value in [0.5, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
+        let img = vec![value; w * h * ch];
+        assert_eq!(best_offset(w, h, ch, &img, &hole, 3, 20), old_best_offset(w, h, ch, &img, &hole, 3, 20));
+    }
+    let full = vec![true; w * h];
+    assert_eq!(best_offset(w, h, ch, &vec![0.5; w * h * ch], &full, 3, 20), None);
+}
+
+#[test]
+#[ignore = "release comparison on 24 and 36 MP images"]
+#[allow(clippy::assertions_on_constants)]
+fn proximity_integral_release_comparison() {
+    use std::{hint::black_box, time::Instant};
+    assert!(!cfg!(debug_assertions), "run with --release");
+    for (w, h, diameters, radius) in [
+        (182usize, 182usize, vec![50usize], 66i32),
+        (332, 332, vec![100], 116),
+        (632, 632, vec![200], 216),
+        (932, 932, vec![300], 316),
+        (6000, 4000, vec![32, 256], 512),
+        (7200, 5000, vec![32, 256], 512),
+    ] {
+        for diameter in diameters {
+            let ch = 3;
+            let img: Vec<_> = (0..w * h * ch)
+                .map(|i| {
+                    let (x, y, c) = ((i / ch) % w, (i / ch) / w, i % ch);
+                    ((x % 53 * 7 + y % 47 * 11 + c * 19) % 251) as f32 / 250.0
+                })
+                .collect();
+            let hole: Vec<_> = (0..w * h)
+                .map(|i| {
+                    let dx = (i % w) as f64 - (w / 2) as f64;
+                    let dy = (i / w) as f64 - (h / 2) as f64;
+                    dx * dx + dy * dy < (diameter as f64 * 0.5).powi(2)
+                })
+                .collect();
+            let bbox = hole_bbox(w, h, &hole).unwrap();
+            let cropped = HoleIntegral::new(w, &hole, bbox).unwrap();
+            let old = old_best_offset(w, h, ch, &img, &hole, (diameter / 8).clamp(3, 16), radius);
+            assert_eq!(best_offset(w, h, ch, &img, &hole, (diameter / 8).clamp(3, 16), radius), old);
+            let mut old_ms = Vec::new();
+            let mut new_ms = Vec::new();
+            for run in 0..5 {
+                let mut measure = |fast: bool| {
+                    let start = Instant::now();
+                    let result = if fast {
+                        best_offset(w, h, ch, black_box(&img), &hole, (diameter / 8).clamp(3, 16), radius)
+                    } else {
+                        old_best_offset(w, h, ch, black_box(&img), &hole, (diameter / 8).clamp(3, 16), radius)
+                    };
+                    let ms = start.elapsed().as_secs_f64() * 1000.0;
+                    assert_eq!(result, old);
+                    if fast {
+                        new_ms.push(ms);
+                    } else {
+                        old_ms.push(ms);
+                    }
+                };
+                measure(run % 2 == 0);
+                measure(run % 2 != 0);
+            }
+            old_ms.sort_by(f64::total_cmp);
+            new_ms.sort_by(f64::total_cmp);
+            println!(
+                "PROXIMITY {w}x{h} hole_diameter={diameter} old_ms={old_ms:?} new_ms={new_ms:?} speedup={:.2} full_table_bytes={} cropped_table_bytes={} offset={old:?}",
+                old_ms[2] / new_ms[2],
+                (w + 1) * (h + 1) * 4,
+                cropped.data.len() * 4
+            );
+        }
+    }
+}
+fn old_best_offset(w: usize, h: usize, ch: usize, img: &[f32], hole: &[bool], ring: usize, max_radius: i32) -> Option<(i32, i32)> {
+    let (bx0, by0, bx1, by1) = hole_bbox(w, h, hole)?;
+    // No valid translation can move farther than the source image's extent.
+    // Bound arbitrary public API inputs before constructing i32 search ranges:
+    // usize -> i32 truncation and signed increments otherwise overflow or hang.
+    let extent = w.max(h).min((i32::MAX / 4) as usize) as i32;
+    let max_radius = max_radius.max(0).min(extent);
+    if max_radius == 0 {
+        return None;
+    }
+    let ring = ring.max(1).min(extent as usize) as i32;
+    let hs = integral(w, h, hole);
+    // Sample points: the hole's ring (dilated minus hole), subsampled for speed.
+    let mut pts = Vec::new();
+    let (x0, y0) = (bx0 as i32 - ring, by0 as i32 - ring);
+    let (x1, y1) = (bx1 as i32 + ring, by1 as i32 + ring);
+    for y in y0.max(0)..y1.min(h as i32) {
+        for x in x0.max(0)..x1.min(w as i32) {
+            if hole[y as usize * w + x as usize] {
+                continue;
+            }
+            if window_count(&hs, w, h, x - ring, y - ring, x + ring + 1, y + ring + 1) > 0 {
+                pts.push((x, y));
+            }
+        }
+    }
+    let stride = (pts.len() / 2000).max(1);
+    let pts: Vec<(i32, i32)> = pts.into_iter().step_by(stride).collect();
+    let eval = |dx: i32, dy: i32| -> Option<f32> {
+        let (hx0, hy0, hx1, hy1) = (x0 + dx, y0 + dy, x1 + dx, y1 + dy);
+        let inside = hx0 >= 0 && hy0 >= 0 && hx1 <= w as i32 && hy1 <= h as i32;
+        if (dx, dy) == (0, 0) || !inside || window_count(&hs, w, h, hx0, hy0, hx1, hy1) != 0 {
+            return None;
+        }
+        let mut e = 0.0f32;
+        for &(x, y) in &pts {
+            let (a, b) = ((y as usize * w + x as usize) * ch, ((y + dy) as usize * w + (x + dx) as usize) * ch);
+            for c in 0..ch {
+                let d = img[a + c] - img[b + c];
+                e += d * d;
+            }
+        }
+        // Mild preference for nearer sources on ties.
+        Some(e / pts.len().max(1) as f32 * (1.0 + 1e-3 * ((dx as f32).hypot(dy as f32) / max_radius as f32)))
+    };
+    // Coarse grid over the search window, then a full-resolution refinement around the winner.
+    let step = (max_radius / 24).max(1);
+    let mut best: Option<((i32, i32), f32)> = None;
+    let consider = |dx: i32, dy: i32, best: &mut Option<((i32, i32), f32)>| {
+        if let Some(e) = eval(dx, dy)
+            && best.is_none_or(|(_, be)| e < be)
+        {
+            *best = Some(((dx, dy), e));
+        }
+    };
+    let mut dy = -max_radius;
+    while dy <= max_radius {
+        let mut dx = -max_radius;
+        while dx <= max_radius {
+            consider(dx, dy, &mut best);
+            dx += step;
+        }
+        dy += step;
+    }
+    if step > 1
+        && let Some(((bx, by), _)) = best
+    {
+        for dy in (by - step)..=(by + step) {
+            for dx in (bx - step)..=(bx + step) {
+                consider(dx, dy, &mut best);
+            }
+        }
+    }
+    best.map(|b| b.0)
+}

docs/proximity-match-performance.md

@@ -0,0 +1,71 @@
+# Proximity Match search and memory
+
+Spot Healing's Proximity Match chooses a source displacement by comparing the ring around
+the healing hole. Two changes reduce that search's work:
+
+- Source-validity counts use a summed-area table over the hole's nonzero bounding box.
+  Every excluded cell outside that box is zero, so translating and clipping queries gives
+  the same counts as a full-image table. This also covers separated holes and canvas edges.
+- Each candidate's SSD is compared with the best score after every 32 ring samples. All
+  SSD terms are nonnegative, so the partial sum divided by the final sample count and
+  multiplied by the same distance penalty is a lower bound on the final score. A candidate
+  that reaches the best score cannot win the existing strict comparison.
+
+The sample order, accumulation order, distance penalty, candidate visit order, tie handling,
+search radius and source-pixel copying stay the same. Nonfinite best scores disable pruning.
+The cropped table checks allocation, arithmetic and source slices and returns None on failure.
+The other summed-area tables used by content-aware completion and texture synthesis are unchanged.
+
+## Measurements
+
+Recorded 2026-10-10 on an AWS c7i.4xlarge, Intel Xeon Platinum 8488C, 16 vCPUs, 32 GiB RAM,
+Rust 1.99.0, normal workspace release profile. This kernel is serial. The input is a synthetic
+periodic RGB float texture and a circular hole; values are generated before timing. Each path
+gets a warmup and five measured runs, alternating old/new order. Every run checks identical
+source displacement against the old implementation.
+
+The four small regions follow the existing retouch benchmark's stroke-domain sizing:
+`width = height = diameter + 2 * (diameter + 16)`, search radius `diameter + 16`.
+Ring width is `clamp(diameter / 8, 3, 16)`. Large-image cases use search radius 512.
+These are kernel calls, excluding brush rasterization, surface reads, Poisson blending,
+document history and compositing.
+
+| Input region | Hole diameter | Before p50 | After p50 | Speedup | Old table | New table |
+|---|---:|---:|---:|---:|---:|---:|
+| 182 x 182 | 50 px | 2.924 ms | 0.765 ms | 3.82x | 133,956 B | 10,000 B |
+| 332 x 332 | 100 px | 0.828 ms | 0.409 ms | 2.02x | 443,556 B | 40,000 B |
+| 632 x 632 | 200 px | 1.959 ms | 1.272 ms | 1.54x | 1,602,756 B | 160,000 B |
+| 932 x 932 | 300 px | 4.816 ms | 1.372 ms | 3.51x | 3,481,956 B | 360,000 B |
+| 6000 x 4000 | 32 px | 67.862 ms | 14.672 ms | 4.63x | 96,040,004 B | 4,096 B |
+| 6000 x 4000 | 256 px | 96.509 ms | 15.469 ms | 6.24x | 96,040,004 B | 262,144 B |
+| 7200 x 5000 | 32 px | 99.693 ms | 21.658 ms | 4.60x | 144,048,804 B | 4,096 B |
+| 7200 x 5000 | 256 px | 132.300 ms | 22.397 ms | 5.91x | 144,048,804 B | 262,144 B |
+
+Table storage alone falls by roughly 90–93 percent in these stroke regions. On a full
+24–36 MP input it falls from 96–144 MB to 4 KiB for the 32 px hole, or 256 KiB for the
+256 px hole. Those are table sizes, not total process memory: the input image, hole mask
+and sampled ring still exist. The initial hole-bounds scan remains linear in image size.
+
+The engine already crops a normal healing stroke to a local region. The large-image cases
+show the public kernel's scaling, not an additional 4.6–6.2x UI speedup. Source textures,
+hole shapes, candidate validity and the quality of the current best match affect pruning;
+these measurements do not establish a worst-case speedup or Photoshop image-quality parity.
+A hole whose bounds cover the image still needs a full-size table.
+
+## Correctness and reproduction
+
+Tests compare cropped counts with full-table counts for rectangles wholly inside, outside,
+crossing or clipped against the hole's bounds. They include separated holes, canvas edges,
+very large query coordinates and checked table allocation/bounds failures.
+
+Search tests compare exact source displacements with the previous implementation over
+1, 3, 4 and 5 channels; 8-bit, 16-bit and float working values; several ring/search radii;
+empty, one-pixel and separated holes; uniform ties and nonfinite samples. A completely
+unknown image still has no source candidate.
+
+```sh
+cargo test --release -p photocraft-algo proximity_integral_release_comparison -- --ignored --nocapture
+```
+
+The old search is retained only in the test module. The selected source displacement is
+unchanged, so downstream source-patch copying and gradient-domain blending get the same input.

docs/roadmap.md

@@ -125,6 +125,12 @@ a conservative 512 MiB working-set estimate, excluding stored Surface tiles and
 memory. The crossover is a heuristic; images dominated by near-cutoff alpha can need extra
 scalar work. These measurements do not establish Photoshop filter parity.
 
+2026-10-10: Proximity Match search medians on an AWS c7i.4xlarge improved 1.54–3.82x
+for four synthetic stroke regions (50–300 px holes), with identical source displacements.
+Bounding the mask table to the hole and pruning losing SSD candidates reduced table storage
+by about 90–93 percent there. Full 24–36 MP kernel inputs improved 4.60–6.24x, but normal
+engine strokes already use a cropped region. See [method and limits](proximity-match-performance.md).
+
 2026-10-08: the tools update above is checked against the current toolbar groups in
 [`panels.rs`](../crates/ui-egui/src/panels.rs), the Pencil and tool-cycle tests in
 [`pencil_tests.rs`](../crates/ui-egui/src/pencil_tests.rs), Patch's live-preview test in