Harden UnionFind bindings and mutex-watershed neighbor bounds

B2: validate node ids at the UnionFind binding boundary. UnionFind::find/
merge/merge_to index their parent vectors without a bounds check (intentional
for the hot path), but the scalar and array bindings passed user-supplied ids
straight through, so an out-of-range id was undefined behavior / a segfault
from Python. Add check_node() and route the scalar find/merge/merge_to binds
and the merge_edges/find_nodes helpers through it, throwing std::invalid_argument
before releasing the GIL.

B1: route the mutex-watershed and semantic-mutex-watershed neighbor computation
through detail::valid_offset_target instead of an unchecked flat-index shift.
This bounds-checks the neighbor per axis, so the kernel is memory-safe and
rejects row/plane wrap-around independently of the caller's valid_edges mask.
Drops the now-unused offset_strides precompute.

Add regression tests for the new UnionFind validation and for the existing
binding-layer point-coordinate validation in non_maximum_distance_suppression.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: constantinpape

include/bioimage_cpp/non_maximum_distance_suppression.hxx

@@ -81,7 +81,9 @@ inline void non_maximum_distance_suppression(
     const auto strides = bioimage_cpp::detail::c_order_strides(distance_map.shape);
     const auto n = static_cast<std::size_t>(n_points);
 
-    // Precompute flat index and distance value at each point.
+    // Precompute flat index and distance value at each point. Point coordinates
+    // are validated against distance_map bounds in the binding layer before this
+    // is called (see bindings/distance.cxx).
     std::vector<float> point_dist(n);
     for (std::size_t i = 0; i < n; ++i) {
         const auto *row =

include/bioimage_cpp/segmentation/mutex_watershed.hxx

@@ -78,13 +78,6 @@ void mutex_watershed_grid(
     ));
     const auto spatial_strides = detail::c_order_strides(spatial_shape);
 
-    std::vector<std::ptrdiff_t> offset_strides(number_of_channels, 0);
-    for (std::size_t channel = 0; channel < number_of_channels; ++channel) {
-        for (std::size_t axis = 0; axis < spatial_ndim; ++axis) {
-            offset_strides[channel] += offsets[channel][axis] * spatial_strides[axis];
-        }
-    }
-
     const auto number_of_edges = number_of_nodes * number_of_channels;
     std::vector<WeightedGridEdge<T>> edge_order;
     edge_order.reserve(static_cast<std::size_t>(number_of_edges));
@@ -108,8 +101,13 @@ void mutex_watershed_grid(
         const auto channel = static_cast<std::size_t>(edge_id / number_of_nodes);
         const auto u = edge_id % number_of_nodes;
 
-        const auto v_signed = static_cast<std::int64_t>(u) + static_cast<std::int64_t>(offset_strides[channel]);
-        const auto v = static_cast<std::uint64_t>(v_signed);
+        // Bounds-check the neighbor per axis rather than relying solely on the
+        // valid_edges mask: this keeps the kernel memory-safe and rejects edges
+        // that would wrap across a row/plane boundary.
+        std::uint64_t v = 0;
+        if (!detail::valid_offset_target(u, offsets[channel], spatial_shape, spatial_strides, v)) {
+            continue;
+        }
         std::uint64_t root_u = sets.find(u);
         std::uint64_t root_v = sets.find(v);
         if (root_u == root_v) {

include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx

@@ -105,13 +105,6 @@ void semantic_mutex_watershed_grid(
     ));
     const auto spatial_strides = detail::c_order_strides(spatial_shape);
 
-    std::vector<std::ptrdiff_t> offset_strides(number_of_offsets, 0);
-    for (std::size_t channel = 0; channel < number_of_offsets; ++channel) {
-        for (std::size_t axis = 0; axis < spatial_ndim; ++axis) {
-            offset_strides[channel] += offsets[channel][axis] * spatial_strides[axis];
-        }
-    }
-
     struct WeightedGridEdge {
         T weight;
         std::uint64_t id;
@@ -155,9 +148,12 @@ void semantic_mutex_watershed_grid(
             continue;
         }
 
-        const auto v_signed = static_cast<std::int64_t>(u) +
-            static_cast<std::int64_t>(offset_strides[channel]);
-        const auto v = static_cast<std::uint64_t>(v_signed);
+        // Bounds-check the neighbor per axis (see mutex_watershed_grid): keeps
+        // the kernel memory-safe and rejects edges that wrap across a boundary.
+        std::uint64_t v = 0;
+        if (!detail::valid_offset_target(u, offsets[channel], spatial_shape, spatial_strides, v)) {
+            continue;
+        }
         const auto root_u = sets.find(u);
         const auto root_v = sets.find(v);
         if (root_u == root_v) {

src/bindings/util.cxx

@@ -18,6 +18,18 @@ using EdgeArray = nb::ndarray<nb::numpy, const std::uint64_t, nb::c_contig>;
 using NodeArray = nb::ndarray<nb::numpy, const std::uint64_t, nb::c_contig, nb::ndim<1>>;
 using OutputArray = nb::ndarray<nb::numpy, std::uint64_t, nb::c_contig>;
 
+// UnionFind::find/merge/merge_to index their parent/rank vectors without a
+// bounds check (intentionally, for the hot path). Validate node ids at the
+// binding boundary so out-of-range ids raise a clear error instead of UB.
+void check_node(const util::UnionFind &uf, const std::uint64_t node, const char *name) {
+    if (node >= uf.size()) {
+        throw std::invalid_argument(
+            std::string(name) + " out of range: got " + name + "="
+            + std::to_string(node) + ", size=" + std::to_string(uf.size())
+        );
+    }
+}
+
 void merge_edges(
     util::UnionFind &uf,
     EdgeArray edges
@@ -39,6 +51,11 @@ void merge_edges(
     const auto n_edges = edges.shape(0);
     const auto *data = edges.data();
 
+    for (std::size_t i = 0; i < n_edges; ++i) {
+        check_node(uf, data[2 * i], "edge endpoint");
+        check_node(uf, data[2 * i + 1], "edge endpoint");
+    }
+
     {
         nb::gil_scoped_release release;
         for (std::size_t i = 0; i < n_edges; ++i) {
@@ -49,10 +66,14 @@ void merge_edges(
 
 OutputArray find_nodes(util::UnionFind &uf, NodeArray nodes) {
     const auto n = nodes.shape(0);
+    const auto *input = nodes.data();
+    for (std::size_t i = 0; i < n; ++i) {
+        check_node(uf, input[i], "node");
+    }
+
     auto *out = new std::uint64_t[n]();
     nb::capsule owner(out, [](void *p) noexcept { delete[] static_cast<std::uint64_t *>(p); });
 
-    const auto *input = nodes.data();
     {
         nb::gil_scoped_release release;
         for (std::size_t i = 0; i < n; ++i) {
@@ -64,6 +85,23 @@ OutputArray find_nodes(util::UnionFind &uf, NodeArray nodes) {
     return OutputArray(out, 1, shape, owner);
 }
 
+std::uint64_t find_node(util::UnionFind &uf, const std::uint64_t node) {
+    check_node(uf, node, "node");
+    return uf.find(node);
+}
+
+std::uint64_t merge_pair(util::UnionFind &uf, const std::uint64_t first, const std::uint64_t second) {
+    check_node(uf, first, "first");
+    check_node(uf, second, "second");
+    return uf.merge(first, second);
+}
+
+std::uint64_t merge_to_node(util::UnionFind &uf, const std::uint64_t stable, const std::uint64_t removed) {
+    check_node(uf, stable, "stable");
+    check_node(uf, removed, "removed");
+    return uf.merge_to(stable, removed);
+}
+
 OutputArray element_labeling(util::UnionFind &uf) {
     const auto n = uf.size();
     auto *out = new std::uint64_t[n]();
@@ -98,7 +136,7 @@ void bind_util(nb::module_ &m) {
         )
         .def(
             "find",
-            &util::UnionFind::find,
+            &find_node,
             nb::arg("node"),
             "Return the (path-compressed) root of `node`."
         )
@@ -110,7 +148,7 @@ void bind_util(nb::module_ &m) {
         )
         .def(
             "merge",
-            &util::UnionFind::merge,
+            &merge_pair,
             nb::arg("first"),
             nb::arg("second"),
             "Union the sets containing `first` and `second`. Returns the new root."
@@ -123,7 +161,7 @@ void bind_util(nb::module_ &m) {
         )
         .def(
             "merge_to",
-            &util::UnionFind::merge_to,
+            &merge_to_node,
             nb::arg("stable"),
             nb::arg("removed"),
             "Union the sets containing `stable` and `removed`, forcing "

tests/distance/test_non_maximum_distance_suppression.py

@@ -33,6 +33,20 @@ def test_two_close_points_keeps_higher_value():
     assert out.tolist() == [[5, 5]]
 
 
+def test_out_of_range_point_coordinate_raises():
+    dm = np.zeros((10, 10), dtype=np.float32)
+    pts = np.array([[5, 5], [5, 12]], dtype=np.int64)
+    with pytest.raises(ValueError, match="out of bounds"):
+        nms(dm, pts)
+
+
+def test_negative_point_coordinate_raises():
+    dm = np.zeros((10, 10), dtype=np.float32)
+    pts = np.array([[-1, 5]], dtype=np.int64)
+    with pytest.raises(ValueError, match="out of bounds"):
+        nms(dm, pts)
+
+
 def test_two_far_points_both_survive():
     dm = np.zeros((20, 20), dtype=np.float32)
     dm[2, 2] = 1.0

tests/test_util_union_find.py

@@ -141,3 +141,33 @@ def test_bulk_merge_rejects_three_columns():
     uf = bic.utils.UnionFind(3)
     with pytest.raises(Exception, match=r"\(N, 2\)"):
         uf.merge(np.zeros((2, 3), dtype=np.uint64))
+
+
+def test_scalar_find_rejects_out_of_range_node():
+    uf = bic.utils.UnionFind(3)
+    with pytest.raises(ValueError, match="out of range"):
+        uf.find(3)
+
+
+def test_scalar_merge_rejects_out_of_range_node():
+    uf = bic.utils.UnionFind(3)
+    with pytest.raises(ValueError, match="out of range"):
+        uf.merge(0, 5)
+
+
+def test_merge_to_rejects_out_of_range_node():
+    uf = bic.utils.UnionFind(3)
+    with pytest.raises(ValueError, match="out of range"):
+        uf.merge_to(7, 0)
+
+
+def test_bulk_find_rejects_out_of_range_node():
+    uf = bic.utils.UnionFind(3)
+    with pytest.raises(ValueError, match="out of range"):
+        uf.find(np.array([0, 1, 9], dtype=np.uint64))
+
+
+def test_bulk_merge_rejects_out_of_range_node():
+    uf = bic.utils.UnionFind(3)
+    with pytest.raises(ValueError, match="out of range"):
+        uf.merge(np.array([[0, 1], [2, 8]], dtype=np.uint64))