fix(population): reject int64 narrowing overflow

docs/changelog.d/population-label-int64-boundary.md

@@ -0,0 +1,8 @@
+# Reject population-label int64 narrowing
+
+## Security
+
+- Reject unsigned values above the signed 64-bit boundary and floating-point
+  values that would be saturated by NumPy during population-label compaction.
+- Preserve the largest exact signed `int64` label while keeping group and
+  cluster identifiers compact before Rust-owned allocation.

python/fast_mlsirm/fit.py

@@ -33,17 +33,30 @@ def _compact_population_labels(raw, n_persons: int, name: str):
     arr = _np.asarray(raw)
     if arr.ndim != 1 or arr.shape[0] != n_persons:
         raise ValueError(f"{name} must be a 1-D array of length n_persons ({n_persons})")
-    fl = arr.astype(_np.float64)
-    if not _np.all(_np.isfinite(fl)):
+    validated = arr if arr.dtype.kind == "f" else arr.astype(_np.float64)
+    if not _np.all(_np.isfinite(validated)):
         raise ValueError(f"{name} must be finite")
-    if _np.any(fl < 0) or _np.any(fl != _np.floor(fl)):
+    if _np.any(validated < 0) or _np.any(validated != _np.floor(validated)):
         raise ValueError(f"{name} must be non-negative integers")
+    int64_max = _np.iinfo(_np.int64).max
+    if arr.dtype.kind == "u" and _np.any(arr > _np.uint64(int64_max)):
+        raise ValueError(f"{name} must fit in signed 64-bit integers")
+    if arr.dtype.kind == "f" and _np.finfo(arr.dtype).maxexp > 63:
+        signed_boundary = _np.array(2**63, dtype=arr.dtype)
+        if _np.any(arr >= signed_boundary):
+            raise ValueError(f"{name} must fit in signed 64-bit integers")
     try:
         with _np.errstate(invalid="ignore", over="ignore"):
             int_labels = arr.astype(_np.int64)
     except (OverflowError, TypeError, ValueError) as exc:
         raise ValueError(f"{name} must fit in signed 64-bit integers") from exc
-    if not _np.array_equal(int_labels.astype(_np.float64), fl):
+    if arr.dtype.kind == "f":
+        roundtrip = int_labels.astype(arr.dtype)
+        expected = arr
+    else:
+        roundtrip = int_labels.astype(_np.float64)
+        expected = validated
+    if not _np.array_equal(roundtrip, expected):
         raise ValueError(f"{name} must fit in signed 64-bit integers")
     uniq, remapped = _np.unique(int_labels, return_inverse=True)
     return remapped.astype(_np.int64), int(uniq.size)

tests/test_population_label_int64_roundtrip.py

@@ -28,6 +28,14 @@ def test_population_labels_reject_float_int64_narrowing_overflow_without_warning
             _compact_population_labels(labels, 2, "group_id")
 
 
+def test_population_labels_reject_integer_sequence_at_signed_boundary() -> None:
+    """A Python sequence containing 2**63 must fail regardless of NumPy promotion."""
+    labels = [0, 2**63]
+
+    with pytest.raises(ValueError, match="signed 64-bit"):
+        _compact_population_labels(labels, 2, "group_id")
+
+
 def test_population_labels_preserve_signed_int64_upper_boundary() -> None:
     """The largest valid signed label remains admissible and order preserving."""
     labels = np.array([0, np.iinfo(np.int64).max], dtype=np.int64)
@@ -36,3 +44,31 @@ def test_population_labels_preserve_signed_int64_upper_boundary() -> None:
 
     assert n_populations == 2
     assert ids.tolist() == [0, 1]
+
+
+def test_population_labels_preserve_extended_precision_int64_upper_boundary() -> None:
+    """A wider real dtype must preserve an exactly representable INT64_MAX label."""
+    if np.finfo(np.longdouble).nmant <= np.finfo(np.float64).nmant:
+        pytest.skip("np.longdouble has no additional precision on this platform")
+
+    labels = np.array(
+        [np.longdouble(0), np.longdouble(np.iinfo(np.int64).max)],
+        dtype=np.longdouble,
+    )
+
+    ids, n_populations = _compact_population_labels(labels, 2, "group_id")
+
+    assert n_populations == 2
+    assert ids.tolist() == [0, 1]
+
+
+def test_population_labels_preserve_float16_without_boundary_warning() -> None:
+    """Small floating dtypes must not overflow merely constructing the int64 bound."""
+    labels = np.array([0.0, 1.0], dtype=np.float16)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        ids, n_populations = _compact_population_labels(labels, 2, "group_id")
+
+    assert n_populations == 2
+    assert ids.tolist() == [0, 1]