use rand::rngs::SmallRng;
use rand::{RngExt, SeedableRng};
fn legacy_search(cum: &[f64], r: f64) -> usize {
let mut min_idx: isize = -1;
let mut max_idx: isize = (cum.len() + 1) as isize;
while min_idx < max_idx - 1 {
let mid_idx = isize::midpoint(min_idx, max_idx);
if cum[mid_idx as usize] >= r {
max_idx = mid_idx;
} else {
min_idx = mid_idx;
}
}
max_idx as usize
}
fn current_search(cum: &[f64], r: f64) -> usize {
cum.partition_point(|&x| x < r)
}
#[test]
fn r_below_all_picks_first() {
let cum = vec![0.5, 1.0, 1.5];
assert_eq!(current_search(&cum, 0.0), 0);
assert_eq!(current_search(&cum, -1e9), 0);
}
#[test]
fn r_exactly_at_boundary_picks_that_index() {
let cum = vec![0.5, 1.0, 1.5];
assert_eq!(current_search(&cum, 1.0), 1);
assert_eq!(legacy_search(&cum, 1.0), 1);
}
#[test]
fn r_just_below_boundary_picks_that_index() {
let cum = vec![0.5, 1.0, 1.5];
assert_eq!(current_search(&cum, 0.999_999), 1);
}
#[test]
fn r_near_total_picks_last() {
let cum = vec![0.5, 1.0, 1.5];
assert_eq!(current_search(&cum, 1.499_999), 2);
assert_eq!(current_search(&cum, 1.5), 2);
}
#[test]
fn single_element() {
let cum = vec![1.0];
assert_eq!(current_search(&cum, 0.0), 0);
assert_eq!(current_search(&cum, 0.5), 0);
assert_eq!(current_search(&cum, 1.0), 0);
assert_eq!(legacy_search(&cum, 0.5), 0);
}
#[test]
fn flat_regions_pick_first_of_run() {
let cum = vec![0.5, 0.5, 0.5, 1.0];
assert_eq!(current_search(&cum, 0.5), 0);
assert_eq!(current_search(&cum, 0.5 + f64::EPSILON), 3);
}
#[test]
fn matches_legacy_on_handpicked_arrays() {
let cases: &[(&[f64], &[f64])] = &[
(&[0.1, 0.3, 0.6, 1.0], &[0.0, 0.05, 0.5, 0.999]),
(&[1.0], &[0.0, 0.5, 0.999_999]),
(&[0.0, 0.0, 1.0], &[0.0, 0.5, 0.999]),
(&[0.5, 0.5, 0.5, 0.5, 1.0], &[0.0, 0.5, 0.999]),
];
for (cum, rs) in cases {
for &r in *rs {
assert_eq!(
current_search(cum, r),
legacy_search(cum, r),
"diverged on cum={cum:?}, r={r}"
);
}
}
}
#[test]
fn partition_point_handles_r_above_total_gracefully() {
let cum = vec![0.5, 1.0, 1.5];
assert_eq!(current_search(&cum, 2.0), 3);
assert_eq!(current_search(&cum, f64::INFINITY), 3);
}
#[test]
fn defensive_clamp_keeps_index_in_bounds() {
fn clamped(cum: &[f64], r: f64) -> usize {
cum.partition_point(|&x| x < r).min(cum.len() - 1)
}
let cum = vec![0.5, 1.0, 1.5];
assert_eq!(clamped(&cum, 2.0), 2);
assert_eq!(clamped(&cum, f64::INFINITY), 2);
}
#[test]
fn nan_r_is_handled_without_panicking() {
let cum = [0.5_f64, 1.0, 1.5];
let r = f64::NAN;
let idx = cum.partition_point(|&x| x < r).min(cum.len() - 1);
assert_eq!(idx, 0);
}
#[test]
fn matches_legacy_on_random_monotone_arrays() {
let mut rng = SmallRng::seed_from_u64(0xC0FF_EE42);
for _ in 0..1000 {
let n = rng.random_range(1..=32_usize);
let mut cum = Vec::with_capacity(n);
let mut acc = 0.0f64;
for _ in 0..n {
acc += rng.random_range(0.0..1.0);
cum.push(acc);
}
let total = *cum.last().unwrap();
let r = total * rng.random_range(0.0..1.0);
assert_eq!(
current_search(&cum, r),
legacy_search(&cum, r),
"diverged on cum={cum:?}, r={r}"
);
}
}
#[test]
fn algorithm_call_site_invariant_holds() {
let mut rng = SmallRng::seed_from_u64(7);
for _ in 0..1000 {
let n = rng.random_range(1..=16_usize);
let mut cum = Vec::with_capacity(n);
let mut acc = 1e-12; for _ in 0..n {
acc += rng.random_range(1e-6..1.0);
cum.push(acc);
}
let total = *cum.last().unwrap();
let r = total * rng.random_range(0.0..1.0);
let idx = current_search(&cum, r);
assert!(idx < n, "idx {idx} out of range for n={n}");
}
}