|
| 1 | +// Copyright (c) 2024-present, fjall-rs |
| 2 | +// This source code is licensed under both the Apache 2.0 and MIT License |
| 3 | +// (found in the LICENSE-* files in the repository) |
| 4 | + |
| 5 | +/// Returns the index of the partition point according to the given predicate |
| 6 | +/// (the index of the first element of the second partition). |
| 7 | +/// |
| 8 | +/// This seems to be faster than std's partition_point: https://github.com/rust-lang/rust/issues/138796 |
| 9 | +pub fn partition_point<T, F>(slice: &[T], pred: F) -> usize |
| 10 | +where |
| 11 | + F: Fn(&T) -> bool, |
| 12 | +{ |
| 13 | + let mut left = 0; |
| 14 | + let mut right = slice.len(); |
| 15 | + |
| 16 | + if right == 0 { |
| 17 | + return 0; |
| 18 | + } |
| 19 | + |
| 20 | + while left < right { |
| 21 | + let mid = (left + right) / 2; |
| 22 | + |
| 23 | + // SAFETY: See https://github.com/rust-lang/rust/blob/ebf0cf75d368c035f4c7e7246d203bd469ee4a51/library/core/src/slice/mod.rs#L2834-L2836 |
| 24 | + #[warn(unsafe_code)] |
| 25 | + let item = unsafe { slice.get_unchecked(mid) }; |
| 26 | + |
| 27 | + if pred(item) { |
| 28 | + left = mid + 1; |
| 29 | + } else { |
| 30 | + right = mid; |
| 31 | + } |
| 32 | + } |
| 33 | + |
| 34 | + left |
| 35 | +} |
| 36 | + |
| 37 | +#[cfg(test)] |
| 38 | +mod tests { |
| 39 | + use super::partition_point; |
| 40 | + use test_log::test; |
| 41 | + |
| 42 | + #[test] |
| 43 | + fn binary_search_first() { |
| 44 | + let items = [1, 2, 3, 4, 5]; |
| 45 | + let idx = partition_point(&items, |&x| x < 1); |
| 46 | + assert_eq!(0, idx); |
| 47 | + |
| 48 | + let std_pp_idx = items.partition_point(|&x| x < 1); |
| 49 | + assert_eq!(std_pp_idx, idx); |
| 50 | + } |
| 51 | + |
| 52 | + #[test] |
| 53 | + fn binary_search_last() { |
| 54 | + let items = [1, 2, 3, 4, 5]; |
| 55 | + let idx = partition_point(&items, |&x| x < 5); |
| 56 | + assert_eq!(4, idx); |
| 57 | + |
| 58 | + let std_pp_idx = items.partition_point(|&x| x < 5); |
| 59 | + assert_eq!(std_pp_idx, idx); |
| 60 | + } |
| 61 | + |
| 62 | + #[test] |
| 63 | + fn binary_search_middle() { |
| 64 | + let items = [1, 2, 3, 4, 5]; |
| 65 | + let idx = partition_point(&items, |&x| x < 3); |
| 66 | + assert_eq!(2, idx); |
| 67 | + |
| 68 | + let std_pp_idx = items.partition_point(|&x| x < 3); |
| 69 | + assert_eq!(std_pp_idx, idx); |
| 70 | + } |
| 71 | + |
| 72 | + #[test] |
| 73 | + fn binary_search_none() { |
| 74 | + let items = [1, 2, 3, 4, 5]; |
| 75 | + let idx = partition_point(&items, |&x| x < 10); |
| 76 | + assert_eq!(5, idx); |
| 77 | + |
| 78 | + let std_pp_idx = items.partition_point(|&x| x < 10); |
| 79 | + assert_eq!(std_pp_idx, idx); |
| 80 | + } |
| 81 | + |
| 82 | + #[test] |
| 83 | + fn binary_search_empty() { |
| 84 | + let items: [i32; 0] = []; |
| 85 | + let idx = partition_point(&items, |&x| x < 10); |
| 86 | + assert_eq!(0, idx); |
| 87 | + |
| 88 | + let std_pp_idx = items.partition_point(|&x| x < 10); |
| 89 | + assert_eq!(std_pp_idx, idx); |
| 90 | + } |
| 91 | +} |
0 commit comments