Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 21ee157

Browse files
committedJul 27, 2024
Rewrite binary search implementation
This restores the original binary search implementation from rust-lang#45333 which has the nice property of having a loop count that only depends on the size of the slice. This, along with explicit conditional moves from rust-lang#128250, means that the entire binary search loop can be perfectly predicted by the branch predictor. Additionally, LLVM is able to unroll the loop when the slice length is known at compile-time. This results in a very compact code sequence of 3-4 instructions per binary search step and zero branches. Fixes rust-lang#53823
1 parent 4f92aeb commit 21ee157

File tree

2 files changed

+46
-32
lines changed

2 files changed

+46
-32
lines changed
 

‎library/core/src/slice/mod.rs

+40-26
Original file line numberDiff line numberDiff line change
@@ -2787,40 +2787,54 @@ impl<T> [T] {
27872787
where
27882788
F: FnMut(&'a T) -> Ordering,
27892789
{
2790-
// INVARIANTS:
2791-
// - 0 <= left <= left + size = right <= self.len()
2792-
// - f returns Less for everything in self[..left]
2793-
// - f returns Greater for everything in self[right..]
27942790
let mut size = self.len();
2795-
let mut left = 0;
2796-
let mut right = size;
2797-
while left < right {
2798-
let mid = left + size / 2;
2799-
2800-
// SAFETY: the while condition means `size` is strictly positive, so
2801-
// `size/2 < size`. Thus `left + size/2 < left + size`, which
2802-
// coupled with the `left + size <= self.len()` invariant means
2803-
// we have `left + size/2 < self.len()`, and this is in-bounds.
2791+
if size == 0 {
2792+
return Err(0);
2793+
}
2794+
let mut base = 0usize;
2795+
2796+
// This loop intentionally doesn't have an early exit if the comparison
2797+
// returns Equal. We want the number of loop iterations to depend *only*
2798+
// on the size of the input slice so that the CPU can reliably predict
2799+
// the loop count.
2800+
while size > 1 {
2801+
let half = size / 2;
2802+
let mid = base + half;
2803+
2804+
// SAFETY: the call is made safe by the following inconstants:
2805+
// - `mid >= 0`: by definition
2806+
// - `mid < size`: `mid = size / 2 + size / 4 + size / 8 ...`
28042807
let cmp = f(unsafe { self.get_unchecked(mid) });
28052808

28062809
// Binary search interacts poorly with branch prediction, so force
28072810
// the compiler to use conditional moves if supported by the target
28082811
// architecture.
2809-
left = select_unpredictable(cmp == Less, mid + 1, left);
2810-
right = select_unpredictable(cmp == Greater, mid, right);
2811-
if cmp == Equal {
2812-
// SAFETY: same as the `get_unchecked` above
2813-
unsafe { hint::assert_unchecked(mid < self.len()) };
2814-
return Ok(mid);
2815-
}
2816-
2817-
size = right - left;
2812+
base = select_unpredictable(cmp == Greater, base, mid);
2813+
2814+
// This is imprecise in the case where `size` is odd and the
2815+
// comparison returns Greater: the mid element still gets included
2816+
// by `size` even though it's known to be larger than the element
2817+
// being searched for.
2818+
//
2819+
// This is fine though: we gain more performance by keeping the
2820+
// loop iteration count invariant (and thus predictable) than we
2821+
// lose from considering one additional element.
2822+
size -= half;
28182823
}
28192824

2820-
// SAFETY: directly true from the overall invariant.
2821-
// Note that this is `<=`, unlike the assume in the `Ok` path.
2822-
unsafe { hint::assert_unchecked(left <= self.len()) };
2823-
Err(left)
2825+
// SAFETY: base is always in [0, size) because base <= mid.
2826+
let cmp = f(unsafe { self.get_unchecked(base) });
2827+
if cmp == Equal {
2828+
// SAFETY: same as the `get_unchecked` above.
2829+
unsafe { hint::assert_unchecked(base < self.len()) };
2830+
Ok(base)
2831+
} else {
2832+
let result = base + (cmp == Less) as usize;
2833+
// SAFETY: same as the `get_unchecked` above.
2834+
// Note that this is `<=`, unlike the assume in the `Ok` path.
2835+
unsafe { hint::assert_unchecked(result <= self.len()) };
2836+
Err(result)
2837+
}
28242838
}
28252839

28262840
/// Binary searches this slice with a key extraction function.

‎library/core/tests/slice.rs

+6-6
Original file line numberDiff line numberDiff line change
@@ -69,13 +69,13 @@ fn test_binary_search() {
6969
assert_eq!(b.binary_search(&8), Err(5));
7070

7171
let b = [(); usize::MAX];
72-
assert_eq!(b.binary_search(&()), Ok(usize::MAX / 2));
72+
assert_eq!(b.binary_search(&()), Ok(usize::MAX - 1));
7373
}
7474

7575
#[test]
7676
fn test_binary_search_by_overflow() {
7777
let b = [(); usize::MAX];
78-
assert_eq!(b.binary_search_by(|_| Ordering::Equal), Ok(usize::MAX / 2));
78+
assert_eq!(b.binary_search_by(|_| Ordering::Equal), Ok(usize::MAX - 1));
7979
assert_eq!(b.binary_search_by(|_| Ordering::Greater), Err(0));
8080
assert_eq!(b.binary_search_by(|_| Ordering::Less), Err(usize::MAX));
8181
}
@@ -87,13 +87,13 @@ fn test_binary_search_implementation_details() {
8787
let b = [1, 1, 2, 2, 3, 3, 3];
8888
assert_eq!(b.binary_search(&1), Ok(1));
8989
assert_eq!(b.binary_search(&2), Ok(3));
90-
assert_eq!(b.binary_search(&3), Ok(5));
90+
assert_eq!(b.binary_search(&3), Ok(6));
9191
let b = [1, 1, 1, 1, 1, 3, 3, 3, 3];
9292
assert_eq!(b.binary_search(&1), Ok(4));
93-
assert_eq!(b.binary_search(&3), Ok(7));
93+
assert_eq!(b.binary_search(&3), Ok(8));
9494
let b = [1, 1, 1, 1, 3, 3, 3, 3, 3];
95-
assert_eq!(b.binary_search(&1), Ok(2));
96-
assert_eq!(b.binary_search(&3), Ok(4));
95+
assert_eq!(b.binary_search(&1), Ok(3));
96+
assert_eq!(b.binary_search(&3), Ok(8));
9797
}
9898

9999
#[test]

0 commit comments

Comments
 (0)
Please sign in to comment.