From 494cf01408c7f5555044916187e5deea96883905 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Fri, 24 Jul 2026 14:22:45 -0700 Subject: [PATCH 1/2] Cast lasso rays in the core instead of once per polygon edge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `select_polygon` bounded the work correctly — zone-pruned bbox first, ray casting only over those candidates — and then spent all of it in NumPy, one vectorized pass per polygon EDGE. Each pass materialized full-length crossing and edge-intersection temporaries over every candidate, so a 64-vertex lasso over 160k candidates moved roughly half a gigabyte through memory to answer a question that fits in registers, and cost grew linearly with vertex count: a 2048-vertex freehand lasso (the API ceiling) took ~370 ms. Every other selection primitive was already native; this one was not. `kernels.polygon_select` walks edges inside the point loop. Crossing parity is order-independent, so the answer is unchanged. On top of that it buckets edges into y slabs, because an edge can only flip parity for a point inside its own y-extent — the index hands back a superset of the edges that can cross a given row and the rest are provably no-ops, so restricting the candidate set is not an approximation. That is what flattens the vertex scaling. vertices numpy native speedup 16 3.70ms 2.56ms 1.4x 64 12.40ms 3.20ms 3.9x 256 48.30ms 3.03ms 15.9x 1024 188.18ms 4.47ms 42.1x 2048 377.57ms 5.39ms 70.0x The full gesture unit the benchmark measures (dispatch, bbox prune, cast, wire mask reply) goes 13.59ms -> 3.21ms, min-of-5. The arithmetic is preserved term for term: `x_j - x_i` and `y_j - y_i` are hoisted per edge (exact subtractions) but the division is NOT strength- reduced to a reciprocal multiply, which would not be. The guarded division covers exactly the lanes the vectorized form computed unconditionally and then masked its inf/NaN away, since endpoints straddling the ray cannot have equal y. Bucketing is declined for a non-finite polygon, where slab bounds stop meaning anything. Verified identical to the previous predicate on convex, concave, self- intersecting (bowtie — even-odd parity is load-bearing), axis-aligned (horizontal edges are the masked-division lanes), collinear, degenerate, repeated-vertex and 2048-vertex polygons; on points landing exactly on vertices and edges; on NaN/+-inf coordinates; and on sparse candidate subsets. Both the Rust and Python tests carry the per-edge formulation as an executable reference rather than describing it. ABI 40 -> 41. --- python/xy/_native.py | 55 ++++++- python/xy/interaction.py | 24 ++- python/xy/kernels.py | 2 + spec/design/rust-engine.md | 2 +- src/kernels.rs | 281 +++++++++++++++++++++++++++++++++++ src/lib.rs | 49 +++++- tests/test_select_polygon.py | 144 ++++++++++++++++++ 7 files changed, 541 insertions(+), 16 deletions(-) create mode 100644 tests/test_select_polygon.py diff --git a/python/xy/_native.py b/python/xy/_native.py index e0298da5..b9c07c41 100644 --- a/python/xy/_native.py +++ b/python/xy/_native.py @@ -24,7 +24,7 @@ from .config import MAX_CONTOUR_WORK, MAX_SCREEN_DIM -ABI_VERSION = 40 +ABI_VERSION = 41 # Rust reports invalid arguments (and, via the ffi_guard panic shield, any # internal panic) by returning `usize::MAX` from size-returning entry points. @@ -546,6 +546,18 @@ def _load() -> ctypes.CDLL: ctypes.c_double, ctypes.c_void_p, ] + lib.xy_polygon_select.restype = ctypes.c_size_t + lib.xy_polygon_select.argtypes = [ + ctypes.c_void_p, # x + ctypes.c_void_p, # y + ctypes.c_size_t, # len + ctypes.c_void_p, # candidate row ids + ctypes.c_size_t, # candidate count + ctypes.c_void_p, # polygon x + ctypes.c_void_p, # polygon y + ctypes.c_size_t, # polygon vertices + ctypes.c_void_p, # output row IDs + ] lib.xy_sample_mask.restype = ctypes.c_int32 lib.xy_sample_mask.argtypes = [ ctypes.c_void_p, @@ -2510,6 +2522,47 @@ def range_indices( return out if written == len(out) else out[:written].copy() +def polygon_select( + x: npt.NDArray[np.float64], + y: npt.NDArray[np.float64], + rows: npt.NDArray[np.uint32], + poly_x: npt.NDArray[np.float64], + poly_y: npt.NDArray[np.float64], +) -> npt.NDArray[np.uint32]: + """Those of `rows` whose canonical point is inside the lasso polygon. + + Even-odd ray casting, order preserved. Non-finite coordinates are never + inside, matching `range_indices`.""" + x = _as_f64(x, "x") + y = _as_f64(y, "y") + if len(x) != len(y): + raise ValueError("x and y must have equal length") + poly_x = _as_f64(poly_x, "polygon x") + poly_y = _as_f64(poly_y, "polygon y") + if len(poly_x) != len(poly_y): + raise ValueError("polygon x and y must have equal length") + rows = np.ascontiguousarray(rows, dtype=np.uint32) + if rows.ndim != 1: + raise ValueError("rows must be 1-D") + out = np.empty(len(rows), dtype=np.uint32) + if len(rows) == 0: + return out + written = _lib.xy_polygon_select( + _ptr_f64(x), + _ptr_f64(y), + len(x), + rows.ctypes.data, + len(rows), + _ptr_f64(poly_x), + _ptr_f64(poly_y), + len(poly_x), + out.ctypes.data, + ) + if written == _USIZE_MAX: + raise ValueError("invalid polygon_select arguments") + return out if written == len(out) else out[:written].copy() + + def sample_mask( ids: npt.NDArray[np.uint64], seed: int, diff --git a/python/xy/interaction.py b/python/xy/interaction.py index dc43875e..89057c51 100644 --- a/python/xy/interaction.py +++ b/python/xy/interaction.py @@ -225,6 +225,14 @@ def select_polygon( The polygon's bounding box first reuses the zone-pruned range predicate; ray casting then runs only on those candidates rather than every row. + + The cast itself is native (`kernels.polygon_select`). Vectorizing it over + NumPy meant one pass per polygon EDGE, each materializing full-length + crossing and edge-intersection temporaries: a 64-vertex lasso over 160k + candidates moved roughly half a gigabyte through memory to answer a + question that fits in registers. Crossing parity is order-independent, so + the native kernel walks edges inside the point loop and returns the same + rows. """ polygon = np.asarray(points, dtype=np.float64) if polygon.ndim != 2 or polygon.shape[1:] != (2,) or not 3 <= len(polygon) <= 2048: @@ -239,25 +247,15 @@ def select_polygon( float(polygon[:, 1].max()), trace_id, ) + poly_x = np.ascontiguousarray(polygon[:, 0]) + poly_y = np.ascontiguousarray(polygon[:, 1]) out: dict[int, np.ndarray] = {} for tid, rows in candidates.items(): if len(rows) == 0: out[tid] = rows continue trace = fig.traces[tid] - x = trace.x.values[rows] - y = trace.y.values[rows] - inside = np.zeros(len(rows), dtype=np.bool_) - j = len(polygon) - 1 - for i in range(len(polygon)): - xi, yi = polygon[i] - xj, yj = polygon[j] - crosses = (yi > y) != (yj > y) - with np.errstate(divide="ignore", invalid="ignore"): - edge_x = (xj - xi) * (y - yi) / (yj - yi) + xi - inside ^= crosses & (x < edge_x) - j = i - out[tid] = rows[inside] + out[tid] = kernels.polygon_select(trace.x.values, trace.y.values, rows, poly_x, poly_y) return out diff --git a/python/xy/kernels.py b/python/xy/kernels.py index e53cedf3..520663f2 100644 --- a/python/xy/kernels.py +++ b/python/xy/kernels.py @@ -62,6 +62,7 @@ valid_indices_f64 = _impl.valid_indices_f64 remap_u8 = _impl.remap_u8 range_indices = _impl.range_indices +polygon_select = _impl.polygon_select sample_mask = _impl.sample_mask sample_range_indices = _impl.sample_range_indices stratified_sample_range_u8 = _impl.stratified_sample_range_u8 @@ -121,6 +122,7 @@ "marching_triangles", "min_max", "normalize_f32", + "polygon_select", "polygon_triangles", "pyramid_append", "pyramid_build", diff --git a/spec/design/rust-engine.md b/spec/design/rust-engine.md index b8a2e3b5..0911821a 100644 --- a/spec/design/rust-engine.md +++ b/spec/design/rust-engine.md @@ -25,7 +25,7 @@ clear ImportError when it can't load, with no pure-Python fallback. | Concern | Today | Verdict | |---|---|---| -| zone maps, encode_f32, m4, bin_2d, bin_2d_mean_color, min_max, histogram_uniform, normalize_f32, range/validity indices, local_log_density | Rust (ABI v40) | correct — new equal-length x/y columns use a paired zone-map call with bit-identical per-column reductions; full-domain density first paint fuses binning with uniform or counted-u8 overlay sampling while retaining exact standalone outputs; mean-color binning (LOD doc §2) is an integer-only pipeline (checked-in sRGB⇄linear-u16 tables, alpha-weighted u64 sums) so grids are bitwise deterministic across thread counts and platforms; mesh/rectangle validity scans consume only columns not already proven finite by zone metadata | +| zone maps, encode_f32, m4, bin_2d, bin_2d_mean_color, min_max, histogram_uniform, normalize_f32, range/validity indices, polygon (lasso) selection, local_log_density | Rust (ABI v41) | correct — new equal-length x/y columns use a paired zone-map call with bit-identical per-column reductions; lasso ray casting (`xy_polygon_select`, §34) walks edges inside the point loop instead of one NumPy pass per edge, and buckets edges by y so a point tests only those spanning its row — the answer is the same crossing parity, without the per-edge full-length temporaries that made a 2048-vertex lasso cost ~370 ms over 160k candidates; full-domain density first paint fuses binning with uniform or counted-u8 overlay sampling while retaining exact standalone outputs; mean-color binning (LOD doc §2) is an integer-only pipeline (checked-in sRGB⇄linear-u16 tables, alpha-weighted u64 sums) so grids are bitwise deterministic across thread counts and platforms; mesh/rectangle validity scans consume only columns not already proven finite by zone metadata | | fixed-width string/bytes/bool factorization | Rust (ABI v36) | correct — compact palettes use a bounded L1-resident codebook with full-record collision checks and emit exact counts; U1 uses a direct Unicode-scalar table with endian support; ≥512k rows probe a prefix then encode disjoint chunks in parallel, merging late labels by canonical first-row order before any retry; Python sees only unique labels and retains display-label ordering policy | | static display-list raster, row-banded polyline/point/segment paint, batched fill+stroke triangle meshes, affine scatter projection plus typed color/size resolution, density/heatmap colormap and sampling | Rust (ABI v36) | correct — commands borrow f32/u8 payload or canonical spans synchronously; compact stratified sampling reuses factorization counts; batched/banded output is byte-identical | | signal processing: `xy_rfft`, `xy_welch_spectra`, `xy_spectrogram` | Rust (ABI v36) | correct — O(N) transforms over sample columns; window/segment policy stays in Python | diff --git a/src/kernels.rs b/src/kernels.rs index 339e99e8..74b4f6e1 100644 --- a/src/kernels.rs +++ b/src/kernels.rs @@ -4778,6 +4778,176 @@ pub fn range_indices( range_indices_impl(x, y, lo_x, hi_x, lo_y, hi_y, par_threads(x.len()), out) } +/// Which of `rows` land inside the lasso polygon, by even-odd ray casting +/// (§34 selection). Writes the surviving canonical row ids to `out` (capacity +/// `rows.len()`, order preserved) and returns how many. +/// +/// The predicate is edge-for-edge the one the Python selection path used: for +/// each edge from vertex `j` to vertex `i`, the horizontal ray is crossed when +/// `(y_i > y) != (y_j > y)` and `x < (x_j - x_i) * (y - y_i) / (y_j - y_i) + +/// x_i`, and membership is the parity of the crossings. `y_j - y_i` cannot be +/// zero where the endpoints straddle the ray, so the guarded division here +/// covers exactly the lanes the vectorized form computed unconditionally and +/// then masked its inf/NaN away. `x_j - x_i` and `y_j - y_i` are hoisted per +/// edge — exact subtractions, so the arithmetic is unchanged; the division is +/// NOT turned into a reciprocal multiply, which would not be. +/// +/// Parity is order-independent, so walking edges inside the point loop rather +/// than points inside the edge loop gives the same answer while holding the +/// whole polygon in registers: no per-edge full-length temporaries (a +/// 64-vertex lasso over 160k candidates moved ~0.5 GB through memory). +/// +/// A non-finite coordinate fails every comparison and so is never inside, +/// matching both the vectorized predicate and `range_indices`. +pub fn polygon_select( + x: &[f64], + y: &[f64], + rows: &[u32], + poly_x: &[f64], + poly_y: &[f64], + out: &mut [u32], +) -> usize { + assert_eq!(x.len(), y.len()); + assert_eq!(poly_x.len(), poly_y.len()); + assert!(out.len() >= rows.len()); + let n = poly_x.len(); + if n < 3 { + return 0; + } + // (x_i, y_i, y_j, x_j - x_i, y_j - y_i) per edge, hoisted so the inner + // loop is a bounds-check-free walk over a couple of KB. + let edges: Vec<(f64, f64, f64, f64, f64)> = (0..n) + .map(|i| { + let j = if i == 0 { n - 1 } else { i - 1 }; + ( + poly_x[i], + poly_y[i], + poly_y[j], + poly_x[j] - poly_x[i], + poly_y[j] - poly_y[i], + ) + }) + .collect(); + let index = PolygonSlabs::build(poly_y, &edges); + let mut write = 0; + for &r in rows { + let (xv, yv) = (x[r as usize], y[r as usize]); + let mut inside = false; + // An edge can only flip parity for a point inside its own y-extent, + // so the slab index hands back a superset of the edges that can + // cross this row and the rest are provably no-ops — same crossings, + // same parity, but a freehand lasso's hundreds of vertices stop + // being scanned in full for every candidate point. + for &e in index.candidates(yv) { + let (xi, yi, yj, dx, dy) = edges[e as usize]; + if ((yi > yv) != (yj > yv)) && xv < dx * (yv - yi) / dy + xi { + inside = !inside; + } + } + if inside { + out[write] = r; + write += 1; + } + } + write +} + +/// Edges bucketed by y, so a point tests only those spanning its row. +struct PolygonSlabs { + y0: f64, + y1: f64, + inv_h: f64, + n: usize, + /// CSR over slabs into `items`; empty when the index was declined. + starts: Vec, + items: Vec, + /// Every edge, in order — the answer when the index was declined. + all: Vec, +} + +impl PolygonSlabs { + /// Bucketing is declined for a non-finite polygon (slab bounds stop + /// meaning anything) and for a vertex count too small to pay for the + /// build. Slab count tracks the vertex count and is capped, bounding the + /// index at `vertices x 256` u32s in the pathological case where every + /// edge spans the full height. + fn build(poly_y: &[f64], edges: &[(f64, f64, f64, f64, f64)]) -> Self { + let n_edges = edges.len(); + let unindexed = || Self { + y0: 0.0, + y1: 0.0, + inv_h: 0.0, + n: 0, + starts: Vec::new(), + items: Vec::new(), + all: (0..n_edges as u32).collect(), + }; + if n_edges < 16 || !poly_y.iter().all(|v| v.is_finite()) { + return unindexed(); + } + let mut y0 = f64::INFINITY; + let mut y1 = f64::NEG_INFINITY; + for &v in poly_y { + y0 = y0.min(v); + y1 = y1.max(v); + } + let height = y1 - y0; + if height <= 0.0 { + return unindexed(); + } + let n = n_edges.clamp(16, 256); + let inv_h = n as f64 / height; + // Half-open [lo, hi) per edge, matching the crossing test's own + // convention, so an edge lands in exactly the slabs it can cross. + let slab_of = |v: f64| (((v - y0) * inv_h) as usize).min(n - 1); + let span = |e: &(f64, f64, f64, f64, f64)| { + let (lo, hi) = if e.1 <= e.2 { (e.1, e.2) } else { (e.2, e.1) }; + (slab_of(lo), slab_of(hi)) + }; + let mut counts = vec![0u32; n + 1]; + for e in edges { + let (a, b) = span(e); + for c in counts.iter_mut().take(b + 2).skip(a + 1) { + *c += 1; + } + } + for s in 0..n { + counts[s + 1] += counts[s]; + } + let mut fill = counts.clone(); + let mut items = vec![0u32; counts[n] as usize]; + for (i, e) in edges.iter().enumerate() { + let (a, b) = span(e); + for slot in fill.iter_mut().take(b + 1).skip(a) { + items[*slot as usize] = i as u32; + *slot += 1; + } + } + Self { + y0, + y1, + inv_h, + n, + starts: counts, + items, + all: Vec::new(), + } + } + + /// Edges that could cross the ray at `yv`. A y outside the polygon's own + /// extent (or NaN) crosses nothing. + fn candidates(&self, yv: f64) -> &[u32] { + if self.n == 0 { + return &self.all; + } + if !(yv >= self.y0 && yv < self.y1) { + return &[]; + } + let s = (((yv - self.y0) * self.inv_h) as usize).min(self.n - 1); + &self.items[self.starts[s] as usize..self.starts[s + 1] as usize] + } +} + /// Append the global indices (`base + i`) of in-window rows to `out`, /// returning the match count. NaN fails every comparison → skipped, matching /// the historical behavior. Dispatches to the AVX2 clone when available. @@ -5080,6 +5250,117 @@ fn min_max_impl(data: &[f64], threads: usize) -> Option<(f64, f64)> { mod tests { use super::*; + /// The vectorized formulation `polygon_select` replaced, kept as the + /// executable definition of the predicate: one pass per edge, crossing + /// parity accumulated over every candidate. + fn polygon_reference( + x: &[f64], + y: &[f64], + rows: &[u32], + poly_x: &[f64], + poly_y: &[f64], + ) -> Vec { + let n = poly_x.len(); + let mut inside = vec![false; rows.len()]; + let mut j = n - 1; + for i in 0..n { + let (xi, yi) = (poly_x[i], poly_y[i]); + let (xj, yj) = (poly_x[j], poly_y[j]); + for (k, &r) in rows.iter().enumerate() { + let (xv, yv) = (x[r as usize], y[r as usize]); + let crosses = (yi > yv) != (yj > yv); + let edge_x = (xj - xi) * (yv - yi) / (yj - yi) + xi; + inside[k] ^= crosses && xv < edge_x; + } + j = i; + } + rows.iter() + .zip(&inside) + .filter(|(_, &b)| b) + .map(|(&r, _)| r) + .collect() + } + + fn polygon_case(poly_x: &[f64], poly_y: &[f64], x: &[f64], y: &[f64]) { + let rows: Vec = (0..x.len() as u32).collect(); + let mut out = vec![0u32; rows.len()]; + let n = polygon_select(x, y, &rows, poly_x, poly_y, &mut out); + assert_eq!( + &out[..n], + &polygon_reference(x, y, &rows, poly_x, poly_y)[..], + "polygon with {} vertices", + poly_x.len() + ); + } + + /// The y-slab index must only ever restrict the candidate edge set, never + /// the answer: indexed (>= 16 vertices) and unindexed paths agree with the + /// per-edge reference on convex, concave, self-intersecting, degenerate + /// and axis-aligned polygons, and on non-finite data. + #[test] + fn polygon_select_matches_the_per_edge_reference_indexed_and_not() { + let mut x = Vec::new(); + let mut y = Vec::new(); + // Deterministic lattice plus points landing exactly on vertices/edges. + for i in 0..60 { + for j in 0..60 { + x.push(i as f64 * 2.0); + y.push(j as f64 * 2.0); + } + } + x.extend_from_slice(&[f64::NAN, 0.0, f64::INFINITY, -f64::INFINITY, 50.0]); + y.extend_from_slice(&[50.0, f64::NAN, 50.0, 50.0, f64::INFINITY]); + + // Triangle and axis-aligned box: below the index threshold, and the + // box's horizontal edges are the zero-height ones. + polygon_case(&[10.0, 90.0, 50.0], &[10.0, 20.0, 80.0], &x, &y); + polygon_case(&[20.0, 80.0, 80.0, 20.0], &[20.0, 20.0, 80.0, 80.0], &x, &y); + // Bowtie: self-intersecting, so even-odd parity is load-bearing. + polygon_case(&[10.0, 90.0, 90.0, 10.0], &[10.0, 90.0, 10.0, 90.0], &x, &y); + // Collinear and zero-height polygons decline the index. + polygon_case(&[0.0, 50.0, 100.0], &[0.0, 50.0, 100.0], &x, &y); + polygon_case(&[0.0, 50.0, 100.0], &[30.0, 30.0, 30.0], &x, &y); + + // Indexed paths: circles and a concave star past the 16-vertex gate. + for n in [16usize, 17, 64, 257] { + let (px, py): (Vec, Vec) = (0..n) + .map(|k| { + let t = std::f64::consts::TAU * k as f64 / n as f64; + (50.0 + 30.0 * t.cos(), 50.0 + 30.0 * t.sin()) + }) + .unzip(); + polygon_case(&px, &py, &x, &y); + } + let (px, py): (Vec, Vec) = (0..24) + .map(|k| { + let t = std::f64::consts::TAU * k as f64 / 24.0; + let r = if k % 2 == 0 { 40.0 } else { 12.0 }; + (50.0 + r * t.cos(), 50.0 + r * t.sin()) + }) + .unzip(); + polygon_case(&px, &py, &x, &y); + + // A non-finite vertex declines the index; the scan must still agree. + polygon_case( + &[10.0, 90.0, 90.0, 10.0, 50.0], + &[10.0, 10.0, 90.0, 90.0, f64::NAN], + &x, + &y, + ); + + // Fewer than three vertices selects nothing. + let rows: Vec = (0..x.len() as u32).collect(); + let mut out = vec![0u32; rows.len()]; + assert_eq!( + polygon_select(&x, &y, &rows, &[0.0, 9.0], &[0.0, 9.0], &mut out), + 0 + ); + assert_eq!( + polygon_select(&x, &y, &[], &[0.0, 9.0, 5.0], &[0.0, 9.0, 5.0], &mut out), + 0 + ); + } + #[test] fn factorize_fixed_preserves_first_seen_codes_and_full_record_identity() { let rows = [b"ab\0", b"xy\0", b"ab\0", b"abx", b"xy\0"]; diff --git a/src/lib.rs b/src/lib.rs index d568d998..e3482cbf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -82,7 +82,7 @@ unsafe fn borrowed_byte_spans<'a>( /// ABI version — bumped on any signature change. The Python wrapper checks this /// at load time and refuses a mismatched library loudly (§33 comm-versioning /// rule, applied to the in-process boundary). -pub const ABI_VERSION: u32 = 40; +pub const ABI_VERSION: u32 = 41; const FACTORIZE_CAPACITY_EXCEEDED: usize = usize::MAX - 1; #[no_mangle] @@ -2797,6 +2797,53 @@ pub unsafe extern "C" fn xy_range_indices( }) } +/// Canonical row ids from `rows` that fall inside the lasso polygon, by +/// even-odd ray casting. Returns the count written; `out` must hold +/// `n_rows` u32s. A polygon of fewer than 3 vertices selects nothing. +/// +/// # Safety +/// `x`/`y` must point to `len` readable f64s, `rows` to `n_rows` readable +/// u32s, `poly_x`/`poly_y` to `n_poly` readable f64s, and `out` to `n_rows` +/// writable u32s. +#[no_mangle] +#[allow(clippy::too_many_arguments)] +pub unsafe extern "C" fn xy_polygon_select( + x: *const f64, + y: *const f64, + len: usize, + rows: *const u32, + n_rows: usize, + poly_x: *const f64, + poly_y: *const f64, + n_poly: usize, + out: *mut u32, +) -> usize { + // u32 index ceiling — see xy_m4_indices. + if len > u32::MAX as usize { + return usize::MAX; + } + if n_rows == 0 { + return 0; + } + if x.is_null() || y.is_null() || rows.is_null() || out.is_null() { + return usize::MAX; + } + if n_poly != 0 && (poly_x.is_null() || poly_y.is_null()) { + return usize::MAX; + } + let x = std::slice::from_raw_parts(x, len); + let y = std::slice::from_raw_parts(y, len); + let rows = std::slice::from_raw_parts(rows, n_rows); + let poly_x = std::slice::from_raw_parts(poly_x, n_poly); + let poly_y = std::slice::from_raw_parts(poly_y, n_poly); + let out = std::slice::from_raw_parts_mut(out, n_rows); + // An out-of-range row id panics the bounds check inside; ffi_guard turns + // that into the error sentinel rather than unwinding across the ABI. + ffi_guard(usize::MAX, || { + kernels::polygon_select(x, y, rows, poly_x, poly_y, out) + }) +} + /// Per-point local log density for a subset. Returns 1 on success, 0 on invalid /// grid/window arguments. /// diff --git a/tests/test_select_polygon.py b/tests/test_select_polygon.py new file mode 100644 index 00000000..98a282f8 --- /dev/null +++ b/tests/test_select_polygon.py @@ -0,0 +1,144 @@ +"""Lasso selection: the native ray cast must answer exactly what the +vectorized per-edge predicate answered (design dossier §34).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from xy import kernels +from xy._figure import Figure + + +def _per_edge_reference( + x: np.ndarray, y: np.ndarray, rows: np.ndarray, polygon: np.ndarray +) -> np.ndarray: + """One pass per polygon edge, crossing parity accumulated across rows. + + This is the definition `kernels.polygon_select` implements; keeping it + here makes the contract executable rather than a comment. + """ + xs, ys = x[rows], y[rows] + inside = np.zeros(len(rows), dtype=np.bool_) + j = len(polygon) - 1 + for i in range(len(polygon)): + xi, yi = polygon[i] + xj, yj = polygon[j] + crosses = (yi > ys) != (yj > ys) + with np.errstate(divide="ignore", invalid="ignore"): + edge_x = (xj - xi) * (ys - yi) / (yj - yi) + xi + inside ^= crosses & (xs < edge_x) + j = i + return rows[inside] + + +def _ring(n: int, *, r: float = 30.0, cx: float = 50.0, cy: float = 50.0) -> np.ndarray: + t = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + return np.stack([cx + r * np.cos(t), cy + r * np.sin(t)], axis=1) + + +def _star(n: int = 24) -> np.ndarray: + t = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + r = np.where(np.arange(n) % 2 == 0, 40.0, 12.0) + return np.stack([50.0 + r * np.cos(t), 50.0 + r * np.sin(t)], axis=1) + + +POLYGONS = { + "triangle": np.array([[10.0, 10.0], [90.0, 20.0], [50.0, 80.0]]), + # Horizontal edges are the zero-height ones the vectorized form divided + # by and masked away. + "axis_aligned_box": np.array([[20.0, 20.0], [80.0, 20.0], [80.0, 80.0], [20.0, 80.0]]), + # Self-intersecting, so even-odd parity is load-bearing. + "bowtie": np.array([[10.0, 10.0], [90.0, 90.0], [90.0, 10.0], [10.0, 90.0]]), + "collinear": np.array([[0.0, 0.0], [50.0, 50.0], [100.0, 100.0]]), + "repeated_vertex": np.array([[10.0, 10.0], [10.0, 10.0], [90.0, 90.0], [10.0, 90.0]]), + "flat": np.array([[0.0, 30.0], [50.0, 30.0], [100.0, 30.0]]), + # Straddles the >= 16-vertex gate where edges get bucketed by y. + "ring_15": _ring(15), + "ring_16": _ring(16), + "ring_64": _ring(64), + "ring_257": _ring(257), + "star_concave": _star(), + "covers_everything": _ring(32, r=500.0), + "encloses_nothing": _ring(32, r=1e-6), +} + + +@pytest.fixture(scope="module") +def sample() -> tuple[np.ndarray, np.ndarray]: + """A lattice (points land exactly on vertices and edges) plus the + non-finite rows a canonical column is allowed to carry (§19).""" + gx, gy = np.meshgrid(np.arange(0.0, 100.0, 2.0), np.arange(0.0, 100.0, 2.0)) + x = np.concatenate([gx.ravel(), [np.nan, 0.0, np.inf, -np.inf, 50.0, 20.0]]) + y = np.concatenate([gy.ravel(), [50.0, np.nan, 50.0, 50.0, np.inf, 20.0]]) + return np.ascontiguousarray(x), np.ascontiguousarray(y) + + +@pytest.mark.parametrize("name", sorted(POLYGONS)) +def test_native_cast_matches_the_per_edge_predicate(name, sample): + x, y = sample + polygon = POLYGONS[name] + rows = np.arange(len(x), dtype=np.uint32) + got = kernels.polygon_select( + x, y, rows, np.ascontiguousarray(polygon[:, 0]), np.ascontiguousarray(polygon[:, 1]) + ) + assert np.array_equal(got, _per_edge_reference(x, y, rows, polygon)) + assert got.dtype == np.uint32 + + +@pytest.mark.parametrize( + "rows", + [ + np.empty(0, dtype=np.uint32), + np.array([0], dtype=np.uint32), + np.arange(0, 2400, 7, dtype=np.uint32), + ], + ids=["empty", "single", "strided"], +) +def test_candidate_subsets_preserve_order_and_membership(rows, sample): + """The real call shape: `rows` is the zone-pruned bbox subset, and the + reply must stay in canonical ascending order.""" + x, y = sample + polygon = POLYGONS["ring_64"] + got = kernels.polygon_select( + x, y, rows, np.ascontiguousarray(polygon[:, 0]), np.ascontiguousarray(polygon[:, 1]) + ) + assert np.array_equal(got, _per_edge_reference(x, y, rows, polygon)) + assert np.all(np.diff(got.astype(np.int64)) > 0) + assert set(got.tolist()) <= set(rows.tolist()) + + +def test_figure_lasso_selects_the_disc_and_nothing_outside_it(): + rng = np.random.default_rng(67) + x = rng.uniform(0.0, 100.0, 20_000) + y = rng.uniform(0.0, 100.0, 20_000) + fig = Figure() + fig.scatter(x, y) + polygon = _ring(64, r=20.0).tolist() + + selected = fig.select_polygon(polygon)[0] + inside_radius = np.hypot(x[selected] - 50.0, y[selected] - 50.0) + # Every selected point is within the circumscribed radius; the inscribed + # radius bounds what the 64-gon must have caught. + assert inside_radius.max() <= 20.0 + caught = np.hypot(x - 50.0, y - 50.0) <= 20.0 * np.cos(np.pi / 64) + assert set(np.flatnonzero(caught).tolist()) <= set(selected.tolist()) + + +def test_polygon_vertex_and_row_validation(): + x = np.array([1.0, 2.0, 3.0]) + y = np.array([1.0, 2.0, 3.0]) + rows = np.array([0, 1, 2], dtype=np.uint32) + square_x = np.array([0.0, 4.0, 4.0, 0.0]) + square_y = np.array([0.0, 0.0, 4.0, 4.0]) + + # Fewer than three vertices encloses nothing rather than erroring. + assert len(kernels.polygon_select(x, y, rows, square_x[:2], square_y[:2])) == 0 + + with pytest.raises(ValueError): + kernels.polygon_select(x, y[:2], rows, square_x, square_y) + with pytest.raises(ValueError): + kernels.polygon_select(x, y, rows, square_x, square_y[:3]) + # A row id past the end of the column is a caller bug, not a silent drop. + with pytest.raises(ValueError): + kernels.polygon_select(x, y, np.array([99], dtype=np.uint32), square_x, square_y) From 640131028bcf50701d9eb44f0728eb0e0da771a2 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Fri, 24 Jul 2026 14:39:40 -0700 Subject: [PATCH 2/2] Bound the lasso slab index and fix its zero-vertex ABI path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two defects review caught in the polygon cast. `xy_polygon_select` built its polygon slices before checking the vertex count, so a caller passing null for an empty polygon reached `slice::from_raw_parts(null, 0)` — which requires a non-null pointer even at length zero. Answer the "fewer than three vertices encloses nothing" case first, before any slice exists. The slab index sized itself purely from the vertex count. `select_polygon` caps callers at 2048, but the kernel and its C export are public with no such ceiling, so a direct call with a very large polygon transiently allocated hundreds of MB, and past ~16.7M vertices the u32 span counters would have overflowed into a corrupt CSR. Cap the CSR at 2^20 entries and shed slabs to stay inside it, rather than refusing to index outright: the fallback scan is O(edges) for every point, which is exactly what hurts most at the vertex counts this guard is about. Below 8 slabs the index stops earning its build and declines. Neither path changes an answer: the budget only coarsens bucketing, and bucketing only ever restricts which edges are tested to a superset of those that can cross. Equivalence to the per-edge reference still holds across every shape, the 2048-vertex ceiling is now covered on the Python side, and a new Rust test pins the CSR invariant at 16 / 2048 / 65536 / 200000 vertices with every edge spanning the full height. --- spec/design/rust-engine.md | 2 +- src/kernels.rs | 69 +++++++++++++++++++++++++++++++++--- src/lib.rs | 9 ++++- tests/test_select_polygon.py | 2 ++ 4 files changed, 75 insertions(+), 7 deletions(-) diff --git a/spec/design/rust-engine.md b/spec/design/rust-engine.md index 0911821a..8b5641ef 100644 --- a/spec/design/rust-engine.md +++ b/spec/design/rust-engine.md @@ -25,7 +25,7 @@ clear ImportError when it can't load, with no pure-Python fallback. | Concern | Today | Verdict | |---|---|---| -| zone maps, encode_f32, m4, bin_2d, bin_2d_mean_color, min_max, histogram_uniform, normalize_f32, range/validity indices, polygon (lasso) selection, local_log_density | Rust (ABI v41) | correct — new equal-length x/y columns use a paired zone-map call with bit-identical per-column reductions; lasso ray casting (`xy_polygon_select`, §34) walks edges inside the point loop instead of one NumPy pass per edge, and buckets edges by y so a point tests only those spanning its row — the answer is the same crossing parity, without the per-edge full-length temporaries that made a 2048-vertex lasso cost ~370 ms over 160k candidates; full-domain density first paint fuses binning with uniform or counted-u8 overlay sampling while retaining exact standalone outputs; mean-color binning (LOD doc §2) is an integer-only pipeline (checked-in sRGB⇄linear-u16 tables, alpha-weighted u64 sums) so grids are bitwise deterministic across thread counts and platforms; mesh/rectangle validity scans consume only columns not already proven finite by zone metadata | +| zone maps, encode_f32, m4, bin_2d, bin_2d_mean_color, min_max, histogram_uniform, normalize_f32, range/validity indices, polygon (lasso) selection, local_log_density | Rust (ABI v41) | correct — new equal-length x/y columns use a paired zone-map call with bit-identical per-column reductions; lasso ray casting (`xy_polygon_select`, §34) walks edges inside the point loop instead of one NumPy pass per edge, and buckets edges by y so a point tests only those spanning its row — the answer is the same crossing parity, without the per-edge full-length temporaries that made a 2048-vertex lasso cost ~370 ms over 160k candidates. Bucketing is declined below 16 vertices (the build outweighs the scans it saves) and for a non-finite polygon (slab bounds stop meaning anything); above 16 the slab count tracks the vertex count, capped at 256 and shed further so the CSR never exceeds 2²⁰ entries, which bounds both the allocation and its u32 cursors independently of what a caller passes; full-domain density first paint fuses binning with uniform or counted-u8 overlay sampling while retaining exact standalone outputs; mean-color binning (LOD doc §2) is an integer-only pipeline (checked-in sRGB⇄linear-u16 tables, alpha-weighted u64 sums) so grids are bitwise deterministic across thread counts and platforms; mesh/rectangle validity scans consume only columns not already proven finite by zone metadata | | fixed-width string/bytes/bool factorization | Rust (ABI v36) | correct — compact palettes use a bounded L1-resident codebook with full-record collision checks and emit exact counts; U1 uses a direct Unicode-scalar table with endian support; ≥512k rows probe a prefix then encode disjoint chunks in parallel, merging late labels by canonical first-row order before any retry; Python sees only unique labels and retains display-label ordering policy | | static display-list raster, row-banded polyline/point/segment paint, batched fill+stroke triangle meshes, affine scatter projection plus typed color/size resolution, density/heatmap colormap and sampling | Rust (ABI v36) | correct — commands borrow f32/u8 payload or canonical spans synchronously; compact stratified sampling reuses factorization counts; batched/banded output is byte-identical | | signal processing: `xy_rfft`, `xy_welch_spectra`, `xy_spectrogram` | Rust (ABI v36) | correct — O(N) transforms over sample columns; window/segment policy stays in Python | diff --git a/src/kernels.rs b/src/kernels.rs index 74b4f6e1..d5991b19 100644 --- a/src/kernels.rs +++ b/src/kernels.rs @@ -4852,6 +4852,11 @@ pub fn polygon_select( write } +/// Ceiling on `PolygonSlabs` CSR entries (4 MB of u32 at the worst case), so +/// neither the allocation nor the u32 cursors depend on caller-supplied +/// vertex counts. +const POLYGON_INDEX_RECORD_BUDGET: usize = 1 << 20; + /// Edges bucketed by y, so a point tests only those spanning its row. struct PolygonSlabs { y0: f64, @@ -4867,10 +4872,10 @@ struct PolygonSlabs { impl PolygonSlabs { /// Bucketing is declined for a non-finite polygon (slab bounds stop - /// meaning anything) and for a vertex count too small to pay for the - /// build. Slab count tracks the vertex count and is capped, bounding the - /// index at `vertices x 256` u32s in the pathological case where every - /// edge spans the full height. + /// meaning anything), for a vertex count too small to pay for the build + /// (< 16), and for one so large that `POLYGON_INDEX_RECORD_BUDGET` leaves + /// fewer than 8 slabs to sort into. Slab count otherwise tracks the vertex + /// count, capped at 256. fn build(poly_y: &[f64], edges: &[(f64, f64, f64, f64, f64)]) -> Self { let n_edges = edges.len(); let unindexed = || Self { @@ -4895,7 +4900,20 @@ impl PolygonSlabs { if height <= 0.0 { return unindexed(); } - let n = n_edges.clamp(16, 256); + // An edge is filed once per slab its y-extent spans, so the index + // holds at most `edges x slabs` entries. Capping that product bounds + // both the allocation and the u32 cursors, and shedding slabs rather + // than refusing outright keeps a huge polygon indexed — the fallback + // scan is O(edges) for every point, which is exactly what hurts most + // at high vertex counts. `select_polygon` caps callers at 2048 + // vertices, well inside the full 256 slabs; this guards the kernel's + // own contract, which has no such ceiling. + let n = n_edges + .clamp(16, 256) + .min((POLYGON_INDEX_RECORD_BUDGET / n_edges).max(1)); + if n < 8 { + return unindexed(); + } let inv_h = n as f64 / height; // Half-open [lo, hi) per edge, matching the crossing test's own // convention, so an edge lands in exactly the slabs it can cross. @@ -5361,6 +5379,47 @@ mod tests { ); } + /// `polygon_select` is a public export with no polygon-size ceiling of its + /// own — `Figure.select_polygon` caps callers at 2048 vertices, the kernel + /// does not — so the slab index must bound its own CSR rather than trust + /// the vertex count it is handed. + #[test] + fn polygon_slab_index_stays_inside_its_record_budget() { + for n in [16usize, 2048, 65_536, 200_000] { + // Worst case for the CSR: every edge spans the full height, so + // each is filed into every slab. + let poly_x: Vec = (0..n).map(|i| i as f64).collect(); + let poly_y: Vec = (0..n) + .map(|i| if i % 2 == 0 { 0.0 } else { 100.0 }) + .collect(); + let edges: Vec<(f64, f64, f64, f64, f64)> = (0..n) + .map(|i| { + let j = if i == 0 { n - 1 } else { i - 1 }; + ( + poly_x[i], + poly_y[i], + poly_y[j], + poly_x[j] - poly_x[i], + poly_y[j] - poly_y[i], + ) + }) + .collect(); + let index = PolygonSlabs::build(&poly_y, &edges); + assert!( + index.items.len() <= POLYGON_INDEX_RECORD_BUDGET, + "{n} vertices produced {} CSR entries", + index.items.len() + ); + // Either the index was declined outright or it kept enough slabs + // to be worth consulting. + assert!( + index.n == 0 || index.n >= 8, + "{n} vertices -> {} slabs", + index.n + ); + } + } + #[test] fn factorize_fixed_preserves_first_seen_codes_and_full_record_identity() { let rows = [b"ab\0", b"xy\0", b"ab\0", b"abx", b"xy\0"]; diff --git a/src/lib.rs b/src/lib.rs index e3482cbf..0d893374 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2825,10 +2825,17 @@ pub unsafe extern "C" fn xy_polygon_select( if n_rows == 0 { return 0; } + // Fewer than three vertices encloses nothing. Answer before building any + // slice: `from_raw_parts` requires a non-null, aligned pointer even at + // length zero, so a caller passing null for an empty polygon must not + // reach the constructions below. + if n_poly < 3 { + return 0; + } if x.is_null() || y.is_null() || rows.is_null() || out.is_null() { return usize::MAX; } - if n_poly != 0 && (poly_x.is_null() || poly_y.is_null()) { + if poly_x.is_null() || poly_y.is_null() { return usize::MAX; } let x = std::slice::from_raw_parts(x, len); diff --git a/tests/test_select_polygon.py b/tests/test_select_polygon.py index 98a282f8..08ea099d 100644 --- a/tests/test_select_polygon.py +++ b/tests/test_select_polygon.py @@ -58,6 +58,8 @@ def _star(n: int = 24) -> np.ndarray: "ring_16": _ring(16), "ring_64": _ring(64), "ring_257": _ring(257), + # The API ceiling `Figure.select_polygon` enforces. + "ring_2048": _ring(2048), "star_concave": _star(), "covers_everything": _ring(32, r=500.0), "encloses_nothing": _ring(32, r=1e-6),